Skip to content

Transforms

AbstractKeypointTransform

Bases: abc.ABC

Base class for all transforms for keypoints augmentation. All transforms subclassing it should implement call method which takes image, mask and keypoints as input and returns transformed image, mask and keypoints.

Parameters:

Name Type Description Default
additional_samples_count int

Number of additional samples to generate for each image. This property is used for mixup & mosaic transforms that needs an extra samples.

 0
Source code in V3_4/src/super_gradients/training/transforms/keypoints/abstract_keypoints_transform.py 
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
class AbstractKeypointTransform(abc.ABC):
    """
    Base class for all transforms for keypoints augmentation.
    All transforms subclassing it should implement __call__ method which takes image, mask and keypoints as input and
    returns transformed image, mask and keypoints.

    :param additional_samples_count: Number of additional samples to generate for each image.
                                    This property is used for mixup & mosaic transforms that needs an extra samples.
    """

    def __init__(self, additional_samples_count: int = 0):
        """
        :param additional_samples_count: (int) number of samples that must be extra samples from dataset. Default value is 0.
        """
        self.additional_samples_count = additional_samples_count

    def __call__(
        self, image: np.ndarray, mask: np.ndarray, joints: np.ndarray, areas: Optional[np.ndarray], bboxes: Optional[np.ndarray]
    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray], Optional[np.ndarray]]:
        """
        Apply transformation to pose estimation sample passed as a tuple
        This method acts as a wrapper for apply_to_sample method to support old-style API.
        """
        sample = PoseEstimationSample(
            image=image,
            mask=mask,
            joints=joints,
            areas=areas,
            bboxes_xywh=bboxes,
            is_crowd=np.zeros(len(joints)),  # Old style API does not pass is_crowd parameter, so we set it to zeros
            additional_samples=None,
        )
        sample = self.apply_to_sample(sample)
        return sample.image, sample.mask, sample.joints, sample.areas, sample.bboxes_xywh

    @abstractmethod
    def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
        """
        Apply transformation to given pose estimation sample.
        Important note - function call may return new object, may modify it in-place.
        This is implementation dependent and if you need to keep original sample intact it
        is recommended to make a copy of it BEFORE passing it to transform.

        :param sample: Input sample to transform.
        :return:       Modified sample (It can be the same instance as input or a new object).
        """
        raise NotImplementedError

    @abstractmethod
    def get_equivalent_preprocessing(self) -> List:
        raise NotImplementedError

__call__(image, mask, joints, areas, bboxes)

Apply transformation to pose estimation sample passed as a tuple This method acts as a wrapper for apply_to_sample method to support old-style API.

Source code in V3_4/src/super_gradients/training/transforms/keypoints/abstract_keypoints_transform.py 
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
def __call__(
    self, image: np.ndarray, mask: np.ndarray, joints: np.ndarray, areas: Optional[np.ndarray], bboxes: Optional[np.ndarray]
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray], Optional[np.ndarray]]:
    """
    Apply transformation to pose estimation sample passed as a tuple
    This method acts as a wrapper for apply_to_sample method to support old-style API.
    """
    sample = PoseEstimationSample(
        image=image,
        mask=mask,
        joints=joints,
        areas=areas,
        bboxes_xywh=bboxes,
        is_crowd=np.zeros(len(joints)),  # Old style API does not pass is_crowd parameter, so we set it to zeros
        additional_samples=None,
    )
    sample = self.apply_to_sample(sample)
    return sample.image, sample.mask, sample.joints, sample.areas, sample.bboxes_xywh

__init__(additional_samples_count=0)

Parameters:

Name Type Description Default
additional_samples_count int

(int) number of samples that must be extra samples from dataset. Default value is 0.

 0
Source code in V3_4/src/super_gradients/training/transforms/keypoints/abstract_keypoints_transform.py 
20
21
22
23
24
def __init__(self, additional_samples_count: int = 0):
    """
    :param additional_samples_count: (int) number of samples that must be extra samples from dataset. Default value is 0.
    """
    self.additional_samples_count = additional_samples_count

apply_to_sample(sample) abstractmethod

Apply transformation to given pose estimation sample. Important note - function call may return new object, may modify it in-place. This is implementation dependent and if you need to keep original sample intact it is recommended to make a copy of it BEFORE passing it to transform.

Parameters:

Name Type Description Default
sample PoseEstimationSample

Input sample to transform.

 required

Returns:

Type Description
PoseEstimationSample

Modified sample (It can be the same instance as input or a new object).
Source code in V3_4/src/super_gradients/training/transforms/keypoints/abstract_keypoints_transform.py 
45
46
47
48
49
50
51
52
53
54
55
56
@abstractmethod
def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
    """
    Apply transformation to given pose estimation sample.
    Important note - function call may return new object, may modify it in-place.
    This is implementation dependent and if you need to keep original sample intact it
    is recommended to make a copy of it BEFORE passing it to transform.

    :param sample: Input sample to transform.
    :return:       Modified sample (It can be the same instance as input or a new object).
    """
    raise NotImplementedError

KeypointsBrightnessContrast

Bases: AbstractKeypointTransform

Apply brightness and contrast change to the input image using following formula: image = (image - mean_brightness) * contrast_gain + mean_brightness + brightness_gain Transformation preserves input image dtype. Saturation cast is performed at the end of the transformation.

Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_brightness_contrast.py 
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
@register_transform()
class KeypointsBrightnessContrast(AbstractKeypointTransform):
    """
    Apply brightness and contrast change to the input image using following formula:
    image = (image - mean_brightness) * contrast_gain + mean_brightness + brightness_gain
    Transformation preserves input image dtype. Saturation cast is performed at the end of the transformation.
    """

    def __init__(self, prob: float, brightness_range: Tuple[float, float], contrast_range: Tuple[float, float]):
        """

        :param prob:             Probability to apply the transform.
        :param brightness_range: Tuple of two elements, min and max brightness gain. Represents a relative range of
                                 brightness gain with respect to average image brightness. A brightness gain of 1.0
                                 indicates no change in brightness. Therefore, optimal value for this parameter is
                                 somewhere inside (0, 2) range.
        :param contrast_range:   Tuple of two elements, min and max contrast gain. Effective contrast_gain would be
                                 uniformly sampled from this range. Based on definition of contrast gain, it's optimal
                                 value is somewhere inside (0, 2) range.
        """
        if len(brightness_range) != 2:
            raise ValueError("Brightness range must be a tuple of two elements, got: " + str(brightness_range))
        if len(contrast_range) != 2:
            raise ValueError("Contrast range must be a tuple of two elements, got: " + str(contrast_range))
        super().__init__()
        self.prob = prob
        self.brightness_range = tuple(brightness_range)
        self.contrast_range = tuple(contrast_range)

    def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
        if random.random() < self.prob:
            contrast_gain = random.uniform(self.contrast_range[0], self.contrast_range[1])
            brightness_gain = random.uniform(self.brightness_range[0], self.brightness_range[1])

            input_dtype = sample.image.dtype
            image = sample.image.astype(np.float32)
            mean_brightness = np.mean(image, axis=(0, 1))

            image = (image - mean_brightness) * contrast_gain + mean_brightness * brightness_gain

            # get min & max values for the input_dtype
            min_value = np.iinfo(input_dtype).min
            max_value = np.iinfo(input_dtype).max
            sample.image = np.clip(image, a_min=min_value, a_max=max_value).astype(input_dtype)
        return sample

    def get_equivalent_preprocessing(self) -> List:
        raise RuntimeError(f"{self.__class__} does not have equivalent preprocessing because it is non-deterministic.")

__init__(prob, brightness_range, contrast_range)

Parameters:

Name Type Description Default
prob float

Probability to apply the transform.

 required
brightness_range Tuple[float, float]

Tuple of two elements, min and max brightness gain. Represents a relative range of brightness gain with respect to average image brightness. A brightness gain of 1.0 indicates no change in brightness. Therefore, optimal value for this parameter is somewhere inside (0, 2) range.

 required
contrast_range Tuple[float, float]

Tuple of two elements, min and max contrast gain. Effective contrast_gain would be uniformly sampled from this range. Based on definition of contrast gain, it's optimal value is somewhere inside (0, 2) range.

 required
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_brightness_contrast.py 
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
def __init__(self, prob: float, brightness_range: Tuple[float, float], contrast_range: Tuple[float, float]):
    """

    :param prob:             Probability to apply the transform.
    :param brightness_range: Tuple of two elements, min and max brightness gain. Represents a relative range of
                             brightness gain with respect to average image brightness. A brightness gain of 1.0
                             indicates no change in brightness. Therefore, optimal value for this parameter is
                             somewhere inside (0, 2) range.
    :param contrast_range:   Tuple of two elements, min and max contrast gain. Effective contrast_gain would be
                             uniformly sampled from this range. Based on definition of contrast gain, it's optimal
                             value is somewhere inside (0, 2) range.
    """
    if len(brightness_range) != 2:
        raise ValueError("Brightness range must be a tuple of two elements, got: " + str(brightness_range))
    if len(contrast_range) != 2:
        raise ValueError("Contrast range must be a tuple of two elements, got: " + str(contrast_range))
    super().__init__()
    self.prob = prob
    self.brightness_range = tuple(brightness_range)
    self.contrast_range = tuple(contrast_range)

KeypointsCompose

Bases: AbstractKeypointTransform

Composes several transforms together

Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_compose.py 
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
class KeypointsCompose(AbstractKeypointTransform):
    """
    Composes several transforms together
    """

    def __init__(self, transforms: List[AbstractKeypointTransform], load_sample_fn=None):
        """

        :param transforms:         List of keypoint-based transformations
        :param load_sample_fn:     A method to load additional samples if needed (for mixup & mosaic augmentations).
                                   Default value is None, which would raise an error if additional samples are needed.
        """
        for transform in transforms:
            if load_sample_fn is None and transform.additional_samples_count > 0:
                raise RuntimeError(
                    f"Detected transform {transform.__class__.__name__} that require {transform.additional_samples_count} "
                    f"additional samples, but load_sample_fn is None"
                )

        super().__init__()
        self.transforms = transforms
        self.load_sample_fn = load_sample_fn

    def __call__(
        self, image: np.ndarray, mask: np.ndarray, joints: np.ndarray, areas: Optional[np.ndarray], bboxes: Optional[np.ndarray]
    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray], Optional[np.ndarray]]:
        """
        Apply transformation to pose estimation sample passed as a tuple
        This method acts as a wrapper for apply_to_sample method to support old-style API.
        """
        for transform in self.transforms:
            if transform.additional_samples_count > 0:
                raise RuntimeError(f"{transform.__class__.__name__} require additional samples that is not supported in old-style transforms API")

        for t in self.transforms:
            image, mask, joints, areas, bboxes = t(image, mask, joints, areas, bboxes)

        return image, mask, joints, areas, bboxes

    def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
        """
        Applies the series of transformations to the input sample.
        The function may modify the input sample inplace, so input sample should not be used after the call.

        :param sample: Input sample
        :return:       Transformed sample.
        """
        sample = sample.sanitize_sample()
        sample = self._apply_transforms(sample, transforms=self.transforms, load_sample_fn=self.load_sample_fn)
        return sample

    @classmethod
    def _apply_transforms(cls, sample: PoseEstimationSample, transforms: List[AbstractKeypointTransform], load_sample_fn) -> PoseEstimationSample:
        """
        This helper method allows us to query additional samples for mixup & mosaic augmentations
        that would be also passed through augmentation pipeline. Example:

        ```
          transforms:
            - KeypointsBrightnessContrast:
                brightness_range: [ 0.8, 1.2 ]
                contrast_range: [ 0.8, 1.2 ]
                prob: 0.5
            - KeypointsHSV:
                hgain: 20
                sgain: 20
                vgain: 20
                prob: 0.5
            - KeypointsLongestMaxSize:
                max_height: ${dataset_params.image_size}
                max_width: ${dataset_params.image_size}
            - KeypointsMixup:
                prob: ${dataset_params.mixup_prob}
        ```

        In the example above all samples in mixup will be forwarded through KeypointsBrightnessContrast, KeypointsHSV,
        KeypointsLongestMaxSize and only then mixed up.

        :param sample:         Input data sample
        :param transforms:     List of transformations to apply
        :param load_sample_fn: A method to load additional samples if needed
        :return:               A data sample after applying transformations
        """
        applied_transforms_so_far = []
        for t in transforms:
            if t.additional_samples_count == 0:
                sample = t.apply_to_sample(sample)
                applied_transforms_so_far.append(t)
            else:
                additional_samples = [load_sample_fn() for _ in range(t.additional_samples_count)]
                additional_samples = [
                    cls._apply_transforms(
                        sample,
                        applied_transforms_so_far,
                        load_sample_fn=load_sample_fn,
                    )
                    for sample in additional_samples
                ]
                sample.additional_samples = additional_samples
                sample = t.apply_to_sample(sample)

        return sample

    def get_equivalent_preprocessing(self) -> List:
        preprocessing = []
        for t in self.transforms:
            preprocessing += t.get_equivalent_preprocessing()
        return preprocessing

    def __repr__(self):
        format_string = self.__class__.__name__ + "("
        for t in self.transforms:
            format_string += f"\t{repr(t)}"
        format_string += "\n)"
        return format_string

__call__(image, mask, joints, areas, bboxes)

Apply transformation to pose estimation sample passed as a tuple This method acts as a wrapper for apply_to_sample method to support old-style API.

Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_compose.py 
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
def __call__(
    self, image: np.ndarray, mask: np.ndarray, joints: np.ndarray, areas: Optional[np.ndarray], bboxes: Optional[np.ndarray]
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray], Optional[np.ndarray]]:
    """
    Apply transformation to pose estimation sample passed as a tuple
    This method acts as a wrapper for apply_to_sample method to support old-style API.
    """
    for transform in self.transforms:
        if transform.additional_samples_count > 0:
            raise RuntimeError(f"{transform.__class__.__name__} require additional samples that is not supported in old-style transforms API")

    for t in self.transforms:
        image, mask, joints, areas, bboxes = t(image, mask, joints, areas, bboxes)

    return image, mask, joints, areas, bboxes

__init__(transforms, load_sample_fn=None)

Parameters:

Name Type Description Default
transforms List[AbstractKeypointTransform]

List of keypoint-based transformations

 required
load_sample_fn

A method to load additional samples if needed (for mixup & mosaic augmentations). Default value is None, which would raise an error if additional samples are needed.

 None
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_compose.py 
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
def __init__(self, transforms: List[AbstractKeypointTransform], load_sample_fn=None):
    """

    :param transforms:         List of keypoint-based transformations
    :param load_sample_fn:     A method to load additional samples if needed (for mixup & mosaic augmentations).
                               Default value is None, which would raise an error if additional samples are needed.
    """
    for transform in transforms:
        if load_sample_fn is None and transform.additional_samples_count > 0:
            raise RuntimeError(
                f"Detected transform {transform.__class__.__name__} that require {transform.additional_samples_count} "
                f"additional samples, but load_sample_fn is None"
            )

    super().__init__()
    self.transforms = transforms
    self.load_sample_fn = load_sample_fn

apply_to_sample(sample)

Applies the series of transformations to the input sample. The function may modify the input sample inplace, so input sample should not be used after the call.

Parameters:

Name Type Description Default
sample PoseEstimationSample

Input sample

 required

Returns:

Type Description
PoseEstimationSample

Transformed sample.
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_compose.py 
48
49
50
51
52
53
54
55
56
57
58
def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
    """
    Applies the series of transformations to the input sample.
    The function may modify the input sample inplace, so input sample should not be used after the call.

    :param sample: Input sample
    :return:       Transformed sample.
    """
    sample = sample.sanitize_sample()
    sample = self._apply_transforms(sample, transforms=self.transforms, load_sample_fn=self.load_sample_fn)
    return sample

KeypointsHSV

Bases: AbstractKeypointTransform

Apply color change in HSV color space to the input image.

Parameters:

Name Type Description Default
prob float

Probability to apply the transform.

 required
hgain float

Hue gain.

 required
sgain float

Saturation gain.

 required
vgain float

Value gain.

 required
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_hsv.py 
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
@register_transform()
class KeypointsHSV(AbstractKeypointTransform):
    """
    Apply color change in HSV color space to the input image.

    :param prob:            Probability to apply the transform.
    :param hgain:           Hue gain.
    :param sgain:           Saturation gain.
    :param vgain:           Value gain.
    """

    def __init__(self, prob: float, hgain: float, sgain: float, vgain: float):
        """

        :param prob:            Probability to apply the transform.
        :param hgain:           Hue gain.
        :param sgain:           Saturation gain.
        :param vgain:           Value gain.
        """
        super().__init__()
        self.prob = prob
        self.hgain = hgain
        self.sgain = sgain
        self.vgain = vgain

    def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
        if sample.image.shape[2] != 3:
            raise ValueError("HSV transform expects image with 3 channels, got: " + str(sample.image.shape[2]))

        if random.random() < self.prob:
            image_copy = sample.image.copy()
            augment_hsv(image_copy, self.hgain, self.sgain, self.vgain, bgr_channels=(0, 1, 2))
            sample.image = image_copy
        return sample

    def get_equivalent_preprocessing(self):
        raise RuntimeError(f"{self.__class__} does not have equivalent preprocessing because it is non-deterministic.")

__init__(prob, hgain, sgain, vgain)

Parameters:

Name Type Description Default
prob float

Probability to apply the transform.

 required
hgain float

Hue gain.

 required
sgain float

Saturation gain.

 required
vgain float

Value gain.

 required
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_hsv.py 
21
22
23
24
25
26
27
28
29
30
31
32
33
def __init__(self, prob: float, hgain: float, sgain: float, vgain: float):
    """

    :param prob:            Probability to apply the transform.
    :param hgain:           Hue gain.
    :param sgain:           Saturation gain.
    :param vgain:           Value gain.
    """
    super().__init__()
    self.prob = prob
    self.hgain = hgain
    self.sgain = sgain
    self.vgain = vgain

KeypointsImageNormalize

Bases: AbstractKeypointTransform

Normalize image with mean and std using formula (image - mean) / std. Output image will allways have dtype of np.float32.

Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_image_normalize.py 
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
@register_transform(Transforms.KeypointsImageNormalize)
class KeypointsImageNormalize(AbstractKeypointTransform):
    """
    Normalize image with mean and std using formula `(image - mean) / std`.
    Output image will allways have dtype of np.float32.
    """

    def __init__(self, mean: Union[float, List[float], ListConfig], std: Union[float, List[float], ListConfig]):
        """

        :param mean: (float, List[float]) A constant bias to be subtracted from the image.
                     If it is a list, it should have the same length as the number of channels in the image.
        :param std:  (float, List[float]) A scaling factor to be applied to the image after subtracting mean.
                     If it is a list, it should have the same length as the number of channels in the image.
        """
        super().__init__()
        self.mean = np.array(list(mean)).reshape((1, 1, -1)).astype(np.float32)
        self.std = np.array(list(std)).reshape((1, 1, -1)).astype(np.float32)

    def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
        """
        Apply transformation to given pose estimation sample

        :param sample: A pose estimation sample
        :return:       Same pose estimation sample with normalized image
        """
        sample.image = np.divide(sample.image - self.mean, self.std, dtype=np.float32)
        return sample

    def __repr__(self):
        return self.__class__.__name__ + f"(mean={self.mean}, std={self.std})"

    def get_equivalent_preprocessing(self) -> List:
        return [{Processings.NormalizeImage: {"mean": self.mean, "std": self.std}}]

__init__(mean, std)

Parameters:

Name Type Description Default
mean Union[float, List[float], ListConfig]

(float, List[float]) A constant bias to be subtracted from the image. If it is a list, it should have the same length as the number of channels in the image.

 required
std Union[float, List[float], ListConfig]

(float, List[float]) A scaling factor to be applied to the image after subtracting mean. If it is a list, it should have the same length as the number of channels in the image.

 required
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_image_normalize.py 
20
21
22
23
24
25
26
27
28
29
30
def __init__(self, mean: Union[float, List[float], ListConfig], std: Union[float, List[float], ListConfig]):
    """

    :param mean: (float, List[float]) A constant bias to be subtracted from the image.
                 If it is a list, it should have the same length as the number of channels in the image.
    :param std:  (float, List[float]) A scaling factor to be applied to the image after subtracting mean.
                 If it is a list, it should have the same length as the number of channels in the image.
    """
    super().__init__()
    self.mean = np.array(list(mean)).reshape((1, 1, -1)).astype(np.float32)
    self.std = np.array(list(std)).reshape((1, 1, -1)).astype(np.float32)

apply_to_sample(sample)

Apply transformation to given pose estimation sample

Parameters:

Name Type Description Default
sample PoseEstimationSample

A pose estimation sample

 required

Returns:

Type Description
PoseEstimationSample

Same pose estimation sample with normalized image
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_image_normalize.py 
32
33
34
35
36
37
38
39
40
def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
    """
    Apply transformation to given pose estimation sample

    :param sample: A pose estimation sample
    :return:       Same pose estimation sample with normalized image
    """
    sample.image = np.divide(sample.image - self.mean, self.std, dtype=np.float32)
    return sample

KeypointsImageStandardize

Bases: AbstractKeypointTransform

Standardize image pixel values with img/max_value formula. Output image will allways have dtype of np.float32.

Parameters:

Name Type Description Default
max_value float

Current maximum value of the image pixels. (usually 255)

 255.0
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_image_standardize.py 
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
@register_transform(Transforms.KeypointsImageStandardize)
class KeypointsImageStandardize(AbstractKeypointTransform):
    """
    Standardize image pixel values with img/max_value formula.
    Output image will allways have dtype of np.float32.

    :param max_value: Current maximum value of the image pixels. (usually 255)
    """

    def __init__(self, max_value: float = 255.0):
        """

        :param max_value: A constant value to divide the image by.
        """
        super().__init__()
        self.max_value = max_value

    def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
        """
        Apply transformation to given pose estimation sample

        :param sample: A pose estimation sample
        :return:       Same pose estimation sample with standardized image
        """
        sample.image = np.divide(sample.image, self.max_value, dtype=np.float32)
        return sample

    def get_equivalent_preprocessing(self) -> List[Dict]:
        return [{Processings.StandardizeImage: {"max_value": self.max_value}}]

    def __repr__(self):
        return self.__class__.__name__ + f"(max_value={self.max_value})"

__init__(max_value=255.0)

Parameters:

Name Type Description Default
max_value float

A constant value to divide the image by.

 255.0
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_image_standardize.py 
20
21
22
23
24
25
26
def __init__(self, max_value: float = 255.0):
    """

    :param max_value: A constant value to divide the image by.
    """
    super().__init__()
    self.max_value = max_value

apply_to_sample(sample)

Apply transformation to given pose estimation sample

Parameters:

Name Type Description Default
sample PoseEstimationSample

A pose estimation sample

 required

Returns:

Type Description
PoseEstimationSample

Same pose estimation sample with standardized image
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_image_standardize.py 
28
29
30
31
32
33
34
35
36
def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
    """
    Apply transformation to given pose estimation sample

    :param sample: A pose estimation sample
    :return:       Same pose estimation sample with standardized image
    """
    sample.image = np.divide(sample.image, self.max_value, dtype=np.float32)
    return sample

KeypointsLongestMaxSize

Bases: AbstractKeypointTransform

Resize data sample to guarantee that input image dimensions is not exceeding maximum width & height

Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_longest_max_size.py 
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
@register_transform(Transforms.KeypointsLongestMaxSize)
class KeypointsLongestMaxSize(AbstractKeypointTransform):
    """
    Resize data sample to guarantee that input image dimensions is not exceeding maximum width & height
    """

    def __init__(self, max_height: int, max_width: int, interpolation: int = cv2.INTER_LINEAR, prob: float = 1.0):
        """

        :param max_height: (int) - Maximum image height
        :param max_width: (int) - Maximum image width
        :param interpolation: Used interpolation method for image
        :param prob: Probability of applying this transform
        """
        super().__init__()
        self.max_height = max_height
        self.max_width = max_width
        self.interpolation = interpolation
        self.prob = prob

    def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
        if random.random() < self.prob:
            height, width = sample.image.shape[:2]
            scale = min(self.max_height / height, self.max_width / width)
            sample.image = self.apply_to_image(sample.image, scale, cv2.INTER_LINEAR)
            sample.mask = self.apply_to_image(sample.mask, scale, cv2.INTER_NEAREST)

            if sample.image.shape[0] != self.max_height and sample.image.shape[1] != self.max_width:
                raise RuntimeError(f"Image shape is not as expected (scale={scale}, input_shape={height, width}, resized_shape={sample.image.shape[:2]})")

            if sample.image.shape[0] > self.max_height or sample.image.shape[1] > self.max_width:
                raise RuntimeError(f"Image shape is not as expected (scale={scale}, input_shape={height, width}, resized_shape={sample.image.shape[:2]}")

            sample.joints = self.apply_to_keypoints(sample.joints, scale)
            if sample.bboxes_xywh is not None:
                sample.bboxes_xywh = self.apply_to_bboxes(sample.bboxes_xywh, scale)

            if sample.areas is not None:
                sample.areas = np.multiply(sample.areas, scale**2, dtype=np.float32)

        return sample

    @classmethod
    def apply_to_image(cls, img, scale, interpolation):
        height, width = img.shape[:2]

        if scale != 1.0:
            new_height, new_width = tuple(int(dim * scale + 0.5) for dim in (height, width))
            img = cv2.resize(img, dsize=(new_width, new_height), interpolation=interpolation)
        return img

    @classmethod
    def apply_to_keypoints(cls, keypoints, scale):
        keypoints = keypoints.astype(np.float32, copy=True)
        keypoints[:, :, 0:2] *= scale
        return keypoints

    @classmethod
    def apply_to_bboxes(cls, bboxes, scale):
        return np.multiply(bboxes, scale, dtype=np.float32)

    def __repr__(self):
        return (
            self.__class__.__name__ + f"(max_height={self.max_height}, "
            f"max_width={self.max_width}, "
            f"interpolation={self.interpolation}, prob={self.prob})"
        )

    def get_equivalent_preprocessing(self) -> List:
        return [{Processings.KeypointsLongestMaxSizeRescale: {"output_shape": (self.max_height, self.max_width)}}]

__init__(max_height, max_width, interpolation=cv2.INTER_LINEAR, prob=1.0)

Parameters:

Name Type Description Default
max_height int

(int) - Maximum image height

 required
max_width int

(int) - Maximum image width

 required
interpolation int

Used interpolation method for image

 cv2.INTER_LINEAR
prob float

Probability of applying this transform

 1.0
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_longest_max_size.py 
19
20
21
22
23
24
25
26
27
28
29
30
31
def __init__(self, max_height: int, max_width: int, interpolation: int = cv2.INTER_LINEAR, prob: float = 1.0):
    """

    :param max_height: (int) - Maximum image height
    :param max_width: (int) - Maximum image width
    :param interpolation: Used interpolation method for image
    :param prob: Probability of applying this transform
    """
    super().__init__()
    self.max_height = max_height
    self.max_width = max_width
    self.interpolation = interpolation
    self.prob = prob

KeypointsMixup

Bases: AbstractKeypointTransform

Apply mixup augmentation and combine two samples into one. Images are averaged with equal weights. Targets are concatenated without any changes. This transform requires both samples have the same image size. The easiest way to achieve this is to use resize + padding before this transform:

# This will apply KeypointsLongestMaxSize and KeypointsPadIfNeeded to two samples individually
# and then apply KeypointsMixup to get a single sample.
train_dataset_params:
    transforms:
        - KeypointsLongestMaxSize:
            max_height: ${dataset_params.image_size}
            max_width: ${dataset_params.image_size}

        - KeypointsPadIfNeeded:
            min_height: ${dataset_params.image_size}
            min_width: ${dataset_params.image_size}
            image_pad_value: [127, 127, 127]
            mask_pad_value: 1
            padding_mode: center

        - KeypointsMixup:
            prob: 0.5

Parameters:

Name Type Description Default
prob float

Probability to apply the transform.

 required
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_mixup.py 
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
@register_transform()
class KeypointsMixup(AbstractKeypointTransform):
    """
    Apply mixup augmentation and combine two samples into one.
    Images are averaged with equal weights. Targets are concatenated without any changes.
    This transform requires both samples have the same image size. The easiest way to achieve this is to use resize + padding before this transform:

    ```yaml
    # This will apply KeypointsLongestMaxSize and KeypointsPadIfNeeded to two samples individually
    # and then apply KeypointsMixup to get a single sample.
    train_dataset_params:
        transforms:
            - KeypointsLongestMaxSize:
                max_height: ${dataset_params.image_size}
                max_width: ${dataset_params.image_size}

            - KeypointsPadIfNeeded:
                min_height: ${dataset_params.image_size}
                min_width: ${dataset_params.image_size}
                image_pad_value: [127, 127, 127]
                mask_pad_value: 1
                padding_mode: center

            - KeypointsMixup:
                prob: 0.5
    ```

    :param prob:            Probability to apply the transform.
    """

    def __init__(self, prob: float):
        """

        :param prob:            Probability to apply the transform.
        """
        super().__init__(additional_samples_count=1)
        self.prob = prob

    def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
        """
        Apply the transform to a single sample.

        :param sample: An input sample. It should have one additional sample in `additional_samples` field.
        :return:       A new pose estimation sample that represents the mixup sample.
        """
        if random.random() < self.prob:
            other = sample.additional_samples[0]
            if sample.image.shape != other.image.shape:
                raise RuntimeError(
                    f"KeypointsMixup requires both samples to have the same image shape. "
                    f"Got {sample.image.shape} and {other.image.shape}. "
                    f"Use KeypointsLongestMaxSize and KeypointsPadIfNeeded to resize and pad images before this transform."
                )
            sample = self._apply_mixup(sample, other)
        return sample

    def _apply_mixup(self, sample: PoseEstimationSample, other: PoseEstimationSample) -> PoseEstimationSample:
        """
        Apply mixup augmentation to a single sample.
        :param sample: First sample.
        :param other:  Second sample.
        :return:       Mixup sample.
        """
        image = (sample.image * 0.5 + other.image * 0.5).astype(sample.image.dtype)
        mask = np.logical_or(sample.mask, other.mask).astype(sample.mask.dtype)
        joints = np.concatenate([sample.joints, other.joints], axis=0)
        is_crowd = np.concatenate([sample.is_crowd, other.is_crowd], axis=0)

        bboxes = self._concatenate_arrays(sample.bboxes_xywh, other.bboxes_xywh, (0, 4))
        areas = self._concatenate_arrays(sample.areas, other.areas, (0,))
        return PoseEstimationSample(image=image, mask=mask, joints=joints, is_crowd=is_crowd, bboxes_xywh=bboxes, areas=areas, additional_samples=None)

    def _concatenate_arrays(self, arr1: Optional[np.ndarray], arr2: Optional[np.ndarray], shape_if_empty) -> Optional[np.ndarray]:
        """
        Concatenate two arrays. If one of the arrays is None, it will be replaced with array of zeros of given shape.
        This is purely utility function to simplify code of stacking arrays that may be None.
        Arrays must have same number of dims.

        :param arr1:           First array
        :param arr2:           Second array
        :param shape_if_empty: Shape of the array to create if one of the arrays is None.
        :return:               Stacked arrays along first axis. If both arrays are None, then None is returned.
        """
        if arr1 is None and arr2 is None:
            return None
        if arr1 is None:
            arr1 = np.zeros(shape_if_empty, dtype=np.float32)
        if arr2 is None:
            arr2 = np.zeros(shape_if_empty, dtype=np.float32)
        return np.concatenate([arr1, arr2], axis=0)

    def get_equivalent_preprocessing(self):
        raise RuntimeError(f"{self.__class__} does not have equivalent preprocessing because it is non-deterministic.")

__init__(prob)

Parameters:

Name Type Description Default
prob float

Probability to apply the transform.

 required
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_mixup.py 
42
43
44
45
46
47
48
def __init__(self, prob: float):
    """

    :param prob:            Probability to apply the transform.
    """
    super().__init__(additional_samples_count=1)
    self.prob = prob

apply_to_sample(sample)

Apply the transform to a single sample.

Parameters:

Name Type Description Default
sample PoseEstimationSample

An input sample. It should have one additional sample in additional_samples field.

 required

Returns:

Type Description
PoseEstimationSample

A new pose estimation sample that represents the mixup sample.
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_mixup.py 
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
    """
    Apply the transform to a single sample.

    :param sample: An input sample. It should have one additional sample in `additional_samples` field.
    :return:       A new pose estimation sample that represents the mixup sample.
    """
    if random.random() < self.prob:
        other = sample.additional_samples[0]
        if sample.image.shape != other.image.shape:
            raise RuntimeError(
                f"KeypointsMixup requires both samples to have the same image shape. "
                f"Got {sample.image.shape} and {other.image.shape}. "
                f"Use KeypointsLongestMaxSize and KeypointsPadIfNeeded to resize and pad images before this transform."
            )
        sample = self._apply_mixup(sample, other)
    return sample

KeypointsMosaic

Bases: AbstractKeypointTransform

Assemble 4 samples together to make 2x2 grid. This transform stacks input samples together to make a square with padding if necessary. This transform does not require input samples to have same size. If input samples have different sizes (H1,W1), (H2,W2), (H3,W3), (H4,W4), then resulting mosaic will have height of max(H1,H2) + max(H3,H4) and width of max(W1+W2, W2+W3), assuming the first sample is located in top left corner, second sample is in top right corner, third sample is in bottom left corner and fourth sample is in bottom right corner of mosaic.

The location of mosaic transform in the transforms list matter. It affects what transforms will be applied to all 4 samples.

In the example below, KeypointsMosaic goes after KeypointsRandomAffineTransform and KeypointsBrightnessContrast. This means that all 4 samples will be transformed with KeypointsRandomAffineTransform and KeypointsBrightnessContrast.

# This will apply KeypointsRandomAffineTransform and KeypointsBrightnessContrast to four sampls individually
# and then assemble them together in mosaic
train_dataset_params:
    transforms:
        - KeypointsRandomAffineTransform:
            min_scale: 0.75
            max_scale: 1.5

        - KeypointsBrightnessContrast:
            brightness_range: [ 0.8, 1.2 ]
            contrast_range: [ 0.8, 1.2 ]
            prob: 0.5

        - KeypointsMosaic:
            prob: 0.5

Contrary, if one puts KeypointsMosaic before KeypointsRandomAffineTransform and KeypointsBrightnessContrast, then 4 original samples will be assembled in mosaic and then transformed with KeypointsRandomAffineTransform and KeypointsBrightnessContrast:

# This will first assemble 4 samples in mosaic and then apply KeypointsRandomAffineTransform and KeypointsBrightnessContrast to the mosaic.
train_dataset_params:
    transforms:
        - KeypointsRandomAffineTransform:
            min_scale: 0.75
            max_scale: 1.5

        - KeypointsBrightnessContrast:
            brightness_range: [ 0.8, 1.2 ]
            contrast_range: [ 0.8, 1.2 ]
            prob: 0.5

        - KeypointsMosaic:
            prob: 0.5
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_mosaic.py 
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
@register_transform()
class KeypointsMosaic(AbstractKeypointTransform):
    """
    Assemble 4 samples together to make 2x2 grid.
    This transform stacks input samples together to make a square with padding if necessary.
    This transform does not require input samples to have same size.
    If input samples have different sizes (H1,W1), (H2,W2), (H3,W3), (H4,W4), then resulting mosaic will have
    height of max(H1,H2) + max(H3,H4) and width of max(W1+W2, W2+W3), assuming the first sample is located in top left corner,
    second sample is in top right corner, third sample is in bottom left corner and fourth sample is in bottom right corner of mosaic.

    The location of mosaic transform in the transforms list matter.
    It affects what transforms will be applied to all 4 samples.

    In the example below, KeypointsMosaic goes after KeypointsRandomAffineTransform and KeypointsBrightnessContrast.
    This means that all 4 samples will be transformed with KeypointsRandomAffineTransform and KeypointsBrightnessContrast.

    ```yaml
    # This will apply KeypointsRandomAffineTransform and KeypointsBrightnessContrast to four sampls individually
    # and then assemble them together in mosaic
    train_dataset_params:
        transforms:
            - KeypointsRandomAffineTransform:
                min_scale: 0.75
                max_scale: 1.5

            - KeypointsBrightnessContrast:
                brightness_range: [ 0.8, 1.2 ]
                contrast_range: [ 0.8, 1.2 ]
                prob: 0.5

            - KeypointsMosaic:
                prob: 0.5
    ```

    Contrary, if one puts KeypointsMosaic before KeypointsRandomAffineTransform and KeypointsBrightnessContrast,
    then 4 original samples will be assembled in mosaic and then transformed with KeypointsRandomAffineTransform and KeypointsBrightnessContrast:

    ```yaml
    # This will first assemble 4 samples in mosaic and then apply KeypointsRandomAffineTransform and KeypointsBrightnessContrast to the mosaic.
    train_dataset_params:
        transforms:
            - KeypointsRandomAffineTransform:
                min_scale: 0.75
                max_scale: 1.5

            - KeypointsBrightnessContrast:
                brightness_range: [ 0.8, 1.2 ]
                contrast_range: [ 0.8, 1.2 ]
                prob: 0.5

            - KeypointsMosaic:
                prob: 0.5
    ```

    """

    def __init__(self, prob: float, pad_value=(127, 127, 127)):
        """

        :param prob:     Probability to apply the transform.
        :param pad_value Value to pad the image if size of samples does not match.
        """
        super().__init__(additional_samples_count=3)
        self.prob = prob
        self.pad_value = tuple(pad_value)

    def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
        """
        Apply transformation to given estimation sample

        :param sample: A pose estimation sample. The sample must have 3 additional samples in it.
        :return:       A new pose estimation sample that represents the final mosaic.
        """
        if random.random() < self.prob:
            samples = [sample] + sample.additional_samples
            sample = self._apply_mosaic(samples)
        return sample

    def _apply_mosaic(self, samples: List[PoseEstimationSample]) -> PoseEstimationSample:
        """
        Actual method to apply mosaic to the sample.

        :param samples: List of 4 samples to make mosaic from.
        :return:        A new pose estimation sample that represents the final mosaic.
        """
        top_left, top_right, btm_left, btm_right = samples

        mosaic_sample = self._stack_samples_vertically(
            self._stack_samples_horizontally(top_left, top_right, pad_from_top=True), self._stack_samples_horizontally(btm_left, btm_right, pad_from_top=False)
        )

        return mosaic_sample

    def _pad_sample(self, sample: PoseEstimationSample, pad_top: int = 0, pad_left: int = 0, pad_right: int = 0, pad_bottom: int = 0) -> PoseEstimationSample:
        """
        Pad the sample with given padding values.

        :param sample:     Input sample. Sample is modified inplace.
        :param pad_top:    Padding in pixels from top.
        :param pad_left:   Padding in pixels from left.
        :param pad_right:  Padding in pixels from right.
        :param pad_bottom: Padding in pixels from bottom.
        :return:           Modified sample.
        """
        sample.image = cv2.copyMakeBorder(
            sample.image, top=pad_top, bottom=pad_bottom, left=pad_left, right=pad_right, borderType=cv2.BORDER_CONSTANT, value=self.pad_value
        )
        sample.mask = cv2.copyMakeBorder(sample.mask, top=pad_top, bottom=pad_bottom, left=pad_left, right=pad_right, borderType=cv2.BORDER_CONSTANT, value=1)

        sample.joints[:, :, 0] += pad_left
        sample.joints[:, :, 1] += pad_top

        sample.bboxes_xywh[:, 0] += pad_left
        sample.bboxes_xywh[:, 1] += pad_top

        return sample

    def _stack_samples_horizontally(self, left: PoseEstimationSample, right: PoseEstimationSample, pad_from_top: bool) -> PoseEstimationSample:
        """
        Stack two samples horizontally.

        :param left:         First sample (Will be located on the left side).
        :param right:        Second sample (Will be location on the right side).
        :param pad_from_top: Controls whether images should be padded from top or from bottom if they have different heights.
        :return:             A stacked sample. If first image has H1,W1 shape and second image has H2,W2 shape,
                             then resulting image will have max(H1,H2), W1+W2 shape.
        """

        max_height = max(left.image.shape[0], right.image.shape[0])
        if pad_from_top:
            left = self._pad_sample(left, pad_top=max_height - left.image.shape[0])
            right = self._pad_sample(right, pad_top=max_height - right.image.shape[0])
        else:
            left = self._pad_sample(left, pad_bottom=max_height - left.image.shape[0])
            right = self._pad_sample(right, pad_bottom=max_height - right.image.shape[0])

        image = np.concatenate([left.image, right.image], axis=1)
        mask = np.concatenate([left.mask, right.mask], axis=1)

        left_sample_width = left.image.shape[1]

        right_bboxes = right.bboxes_xywh
        if right_bboxes is None:
            right_bboxes = np.zeros((0, 4), dtype=np.float32)

        right_joints_offset = np.array([left_sample_width, 0, 0], dtype=right.joints.dtype).reshape((1, 1, 3))
        right_bboxes_offset = np.array([left_sample_width, 0, 0, 0], dtype=right_bboxes.dtype).reshape((1, 4))

        joints = np.concatenate([left.joints, right.joints + right_joints_offset], axis=0)
        bboxes = self._concatenate_arrays(left.bboxes_xywh, right_bboxes + right_bboxes_offset, shape_if_empty=(0, 4))

        is_crowd = np.concatenate([left.is_crowd, right.is_crowd], axis=0)
        areas = self._concatenate_arrays(left.areas, right.areas, shape_if_empty=(0,))
        return PoseEstimationSample(image=image, mask=mask, joints=joints, is_crowd=is_crowd, bboxes_xywh=bboxes, areas=areas, additional_samples=None)

    def _stack_samples_vertically(self, top: PoseEstimationSample, bottom: PoseEstimationSample) -> PoseEstimationSample:
        """
        Stack two samples vertically. If images have different widths, they will be padded to match the width
        of the widest image. In case padding occurs, it will be done from both sides to keep the images centered.

        :param top:    First sample (Will be located on the top).
        :param bottom: Second sample (Will be location on the bottom).
        :return:       A stacked sample. If first image has H1,W1 shape and second image has H2,W2 shape,
                       then resulting image will have H1+H2, max(W1,W2) shape.
        """
        max_width = max(top.image.shape[1], bottom.image.shape[1])

        pad_left = (max_width - top.image.shape[1]) // 2
        pad_right = max_width - top.image.shape[1] - pad_left
        top = self._pad_sample(top, pad_left=pad_left, pad_right=pad_right)

        pad_left = (max_width - bottom.image.shape[1]) // 2
        pad_right = max_width - bottom.image.shape[1] - pad_left
        bottom = self._pad_sample(bottom, pad_left=pad_left, pad_right=pad_right)

        image = np.concatenate([top.image, bottom.image], axis=0)
        mask = np.concatenate([top.mask, bottom.mask], axis=0)

        top_sample_height = top.image.shape[0]

        bottom_bboxes = bottom.bboxes_xywh
        if bottom_bboxes is None:
            bottom_bboxes = np.zeros((0, 4), dtype=np.float32)

        bottom_joints_offset = np.array([0, top_sample_height, 0], dtype=bottom.joints.dtype).reshape((1, 1, 3))
        bottom_bboxes_offset = np.array([0, top_sample_height, 0, 0], dtype=bottom_bboxes.dtype).reshape((1, 4))

        joints = np.concatenate([top.joints, bottom.joints + bottom_joints_offset], axis=0)
        bboxes = self._concatenate_arrays(top.bboxes_xywh, bottom_bboxes + bottom_bboxes_offset, shape_if_empty=(0, 4))

        is_crowd = np.concatenate([top.is_crowd, bottom.is_crowd], axis=0)
        areas = self._concatenate_arrays(top.areas, bottom.areas, shape_if_empty=(0,))
        return PoseEstimationSample(image=image, mask=mask, joints=joints, is_crowd=is_crowd, bboxes_xywh=bboxes, areas=areas, additional_samples=None)

    def _concatenate_arrays(self, arr1: Optional[np.ndarray], arr2: Optional[np.ndarray], shape_if_empty) -> Optional[np.ndarray]:
        """
        Concatenate two arrays. If one of the arrays is None, it will be replaced with array of zeros of given shape.
        This is purely utility function to simplify code of stacking arrays that may be None.
        Arrays must have same number of dims.

        :param arr1:           First array
        :param arr2:           Second array
        :param shape_if_empty: Shape of the array to create if one of the arrays is None.
        :return:               Stacked arrays along first axis. If both arrays are None, then None is returned.
        """
        if arr1 is None and arr2 is None:
            return None
        if arr1 is None:
            arr1 = np.zeros(shape_if_empty, dtype=np.float32)
        if arr2 is None:
            arr2 = np.zeros(shape_if_empty, dtype=np.float32)
        return np.concatenate([arr1, arr2], axis=0)

    def get_equivalent_preprocessing(self):
        raise RuntimeError(f"{self.__class__} does not have equivalent preprocessing because it is non-deterministic.")

__init__(prob, pad_value=(127, 127, 127))

Parameters:

Name Type Description Default
prob float

Probability to apply the transform.

 required
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_mosaic.py 
68
69
70
71
72
73
74
75
76
def __init__(self, prob: float, pad_value=(127, 127, 127)):
    """

    :param prob:     Probability to apply the transform.
    :param pad_value Value to pad the image if size of samples does not match.
    """
    super().__init__(additional_samples_count=3)
    self.prob = prob
    self.pad_value = tuple(pad_value)

apply_to_sample(sample)

Apply transformation to given estimation sample

Parameters:

Name Type Description Default
sample PoseEstimationSample

A pose estimation sample. The sample must have 3 additional samples in it.

 required

Returns:

Type Description
PoseEstimationSample

A new pose estimation sample that represents the final mosaic.
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_mosaic.py 
78
79
80
81
82
83
84
85
86
87
88
def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
    """
    Apply transformation to given estimation sample

    :param sample: A pose estimation sample. The sample must have 3 additional samples in it.
    :return:       A new pose estimation sample that represents the final mosaic.
    """
    if random.random() < self.prob:
        samples = [sample] + sample.additional_samples
        sample = self._apply_mosaic(samples)
    return sample

KeypointsPadIfNeeded

Bases: AbstractKeypointTransform

Pad image and mask to ensure that resulting image size is not less than output_size (rows, cols). Image and mask padded from right and bottom, thus joints remains unchanged.

Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_pad_if_needed.py 
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
@register_transform(Transforms.KeypointsPadIfNeeded)
class KeypointsPadIfNeeded(AbstractKeypointTransform):
    """
    Pad image and mask to ensure that resulting image size is not less than `output_size` (rows, cols).
    Image and mask padded from right and bottom, thus joints remains unchanged.
    """

    def __init__(self, min_height: int, min_width: int, image_pad_value: int, mask_pad_value: float, padding_mode: str = "bottom_right"):
        """

        :param output_size: Desired image size (rows, cols)
        :param image_pad_value: Padding value of image
        :param mask_pad_value: Padding value for mask
        """
        if padding_mode not in ("bottom_right", "center"):
            raise ValueError(f"Unknown padding mode: {padding_mode}. Supported modes: 'bottom_right', 'center'")
        super().__init__()
        self.min_height = min_height
        self.min_width = min_width
        self.image_pad_value = image_pad_value
        self.mask_pad_value = mask_pad_value
        self.padding_mode = padding_mode

    def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
        height, width = sample.image.shape[:2]
        original_dtype = sample.mask.dtype

        if self.padding_mode == "bottom_right":
            pad_left = 0
            pad_top = 0
            pad_bottom = max(0, self.min_height - height)
            pad_right = max(0, self.min_width - width)
        elif self.padding_mode == "center":
            pad_left = max(0, (self.min_width - width) // 2)
            pad_top = max(0, (self.min_height - height) // 2)
            pad_bottom = max(0, self.min_height - height - pad_top)
            pad_right = max(0, self.min_width - width - pad_left)
        else:
            raise RuntimeError(f"Unknown padding mode: {self.padding_mode}")

        image_pad_value = tuple(self.image_pad_value) if isinstance(self.image_pad_value, Iterable) else tuple([self.image_pad_value] * sample.image.shape[-1])
        sample.image = cv2.copyMakeBorder(
            sample.image, top=pad_top, bottom=pad_bottom, left=pad_left, right=pad_right, value=image_pad_value, borderType=cv2.BORDER_CONSTANT
        )

        sample.mask = cv2.copyMakeBorder(
            sample.mask.astype(np.uint8),
            top=pad_top,
            bottom=pad_bottom,
            left=pad_left,
            right=pad_right,
            value=self.mask_pad_value,
            borderType=cv2.BORDER_CONSTANT,
        ).astype(original_dtype)

        sample.joints = self.apply_to_keypoints(sample.joints, pad_left, pad_top)
        if sample.bboxes_xywh is not None:
            sample.bboxes_xywh = self.apply_to_bboxes(sample.bboxes_xywh, pad_left, pad_top)

        return sample

    def apply_to_bboxes(self, bboxes: np.ndarray, pad_left, pad_top):
        bboxes = bboxes.copy()
        bboxes[:, 0] += pad_left
        bboxes[:, 1] += pad_top
        return bboxes

    def apply_to_keypoints(self, keypoints: np.ndarray, pad_left, pad_top):
        keypoints = keypoints.copy()
        keypoints[:, :, 0] += pad_left
        keypoints[:, :, 1] += pad_top
        return keypoints

    def __repr__(self):
        return (
            self.__class__.__name__ + f"(min_height={self.min_height}, "
            f"min_width={self.min_width}, "
            f"image_pad_value={self.image_pad_value}, "
            f"mask_pad_value={self.mask_pad_value}, "
            f"padding_mode={self.padding_mode}, "
            f")"
        )

    def get_equivalent_preprocessing(self) -> List:
        if self.padding_mode == "bottom_right":
            return [{Processings.KeypointsBottomRightPadding: {"output_shape": (self.min_height, self.min_width), "pad_value": self.image_pad_value}}]
        else:
            raise RuntimeError(f"KeypointsPadIfNeeded with padding_mode={self.padding_mode} is not implemented.")

__init__(min_height, min_width, image_pad_value, mask_pad_value, padding_mode='bottom_right')

Parameters:

Name Type Description Default
output_size

Desired image size (rows, cols)

 required
image_pad_value int

Padding value of image

 required
mask_pad_value float

Padding value for mask

 required
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_pad_if_needed.py 
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
def __init__(self, min_height: int, min_width: int, image_pad_value: int, mask_pad_value: float, padding_mode: str = "bottom_right"):
    """

    :param output_size: Desired image size (rows, cols)
    :param image_pad_value: Padding value of image
    :param mask_pad_value: Padding value for mask
    """
    if padding_mode not in ("bottom_right", "center"):
        raise ValueError(f"Unknown padding mode: {padding_mode}. Supported modes: 'bottom_right', 'center'")
    super().__init__()
    self.min_height = min_height
    self.min_width = min_width
    self.image_pad_value = image_pad_value
    self.mask_pad_value = mask_pad_value
    self.padding_mode = padding_mode

KeypointsRandomAffineTransform

Bases: AbstractKeypointTransform

Apply random affine transform to image, mask and joints.

Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_random_affine.py 
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
@register_transform(Transforms.KeypointsRandomAffineTransform)
class KeypointsRandomAffineTransform(AbstractKeypointTransform):
    """
    Apply random affine transform to image, mask and joints.
    """

    def __init__(
        self,
        max_rotation: float,
        min_scale: float,
        max_scale: float,
        max_translate: float,
        image_pad_value: Union[int, float, List[int]],
        mask_pad_value: float,
        interpolation_mode: Union[int, List[int]] = cv2.INTER_LINEAR,
        prob: float = 0.5,
    ):
        """

        :param max_rotation:       Max rotation angle in degrees
        :param min_scale:          Lower bound for the scale change. For +- 20% size jitter this should be 0.8
        :param max_scale:          Lower bound for the scale change. For +- 20% size jitter this should be 1.2
        :param max_translate:      Max translation offset in percents of image size
        :param image_pad_value:    Value to pad the image during affine transform. Can be single scalar or list.
                                   If a list is provided, it should have the same length as the number of channels in the image.
        :param mask_pad_value:     Value to pad the mask during affine transform.
        :param interpolation_mode: A constant integer or list of integers, specifying the interpolation mode to use.
                                   Possible values for interpolation_mode:
                                     cv2.INTER_NEAREST = 0,
                                     cv2.INTER_LINEAR = 1,
                                     cv2.INTER_CUBIC = 2,
                                     cv2.INTER_AREA = 3,
                                     cv2.INTER_LANCZOS4 = 4
                                   To use random interpolation modes on each call, set interpolation_mode = (0,1,2,3,4)
        :param prob:               Probability to apply the transform.
        """
        super().__init__()

        self.max_rotation = max_rotation
        self.min_scale = min_scale
        self.max_scale = max_scale
        self.max_translate = max_translate
        self.image_pad_value = image_pad_value
        self.mask_pad_value = mask_pad_value
        self.prob = prob
        self.interpolation_mode = tuple(interpolation_mode) if isinstance(interpolation_mode, Iterable) else (interpolation_mode,)

    def __repr__(self):
        return (
            self.__class__.__name__ + f"(max_rotation={self.max_rotation}, "
            f"min_scale={self.min_scale}, "
            f"max_scale={self.max_scale}, "
            f"max_translate={self.max_translate}, "
            f"image_pad_value={self.image_pad_value}, "
            f"mask_pad_value={self.mask_pad_value}, "
            f"prob={self.prob})"
        )

    def _get_affine_matrix(self, img: np.ndarray, angle: float, scale: float, dx: float, dy: float) -> np.ndarray:
        """
        Compute the affine matrix that combines rotation of image around center, scaling and translation
        according to given parameters. Order of operations is: scale, rotate, translate.

        :param angle: Rotation angle in degrees
        :param scale: Scaling factor
        :param dx:    Translation in x direction
        :param dy:    Translation in y direction
        :return:      Affine matrix [2,3]
        """
        height, width = img.shape[:2]
        center = (width / 2 + dx * width, height / 2 + dy * height)
        matrix = cv2.getRotationMatrix2D(center, angle, scale)

        return matrix

    def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
        """
        Apply transformation to given pose estimation sample.
        Since this transformation apply affine transform some keypoints/bboxes may be moved outside the image.
        After applying the transform, visibility status of joints is updated to reflect the new position of joints.
        Bounding boxes are clipped to image borders.
        If sample contains areas, they are scaled according to the applied affine transform.

        :param sample: A pose estimation sample
        :return:       A transformed pose estimation sample
        """

        if random.random() < self.prob:
            angle = random.uniform(-self.max_rotation, self.max_rotation)
            scale = random.uniform(self.min_scale, self.max_scale)
            dx = random.uniform(-self.max_translate, self.max_translate)
            dy = random.uniform(-self.max_translate, self.max_translate)
            interpolation = random.choice(self.interpolation_mode)

            mat_output = self._get_affine_matrix(sample.image, angle, scale, dx, dy)
            mat_output = mat_output[:2]

            image_pad_value = (
                tuple(self.image_pad_value) if isinstance(self.image_pad_value, Iterable) else tuple([self.image_pad_value] * sample.image.shape[-1])
            )

            sample.image = self.apply_to_image(
                sample.image, mat_output, interpolation=interpolation, padding_value=image_pad_value, padding_mode=cv2.BORDER_CONSTANT
            )
            sample.mask = self.apply_to_image(
                sample.mask, mat_output, interpolation=cv2.INTER_NEAREST, padding_value=self.mask_pad_value, padding_mode=cv2.BORDER_CONSTANT
            )

            sample.joints = self.apply_to_keypoints(sample.joints, mat_output, sample.image.shape[:2])

            if sample.bboxes_xywh is not None:
                sample.bboxes_xywh = self.apply_to_bboxes(sample.bboxes_xywh, mat_output)

            if sample.areas is not None:
                sample.areas = self.apply_to_areas(sample.areas, mat_output)

            sample = sample.sanitize_sample()

        return sample

    @classmethod
    def apply_to_areas(cls, areas: np.ndarray, mat: np.ndarray) -> np.ndarray:
        """
        Apply affine transform to areas.

        :param areas: [N] Single-dimension array of areas
        :param mat:   [2,3] Affine transformation matrix
        :return:      [N] Single-dimension array of areas
        """
        det = np.linalg.det(mat[:2, :2])
        return (areas * abs(det)).astype(areas.dtype)

    @classmethod
    def apply_to_bboxes(cls, bboxes_xywh: np.ndarray, mat: np.ndarray) -> np.ndarray:
        """

        :param bboxes: (N,4) array of bboxes in XYWH format
        :param mat:    [2,3] Affine transformation matrix
        :return:       (N,4) array of bboxes in XYWH format
        """

        def bbox_shift_scale_rotate(bbox, m):
            x_min, y_min, x_max, y_max = bbox[:4]

            x = np.array([x_min, x_max, x_max, x_min])
            y = np.array([y_min, y_min, y_max, y_max])
            ones = np.ones(shape=(len(x)))
            points_ones = np.vstack([x, y, ones]).transpose()

            tr_points = m.dot(points_ones.T).T

            x_min, x_max = min(tr_points[:, 0]), max(tr_points[:, 0])
            y_min, y_max = min(tr_points[:, 1]), max(tr_points[:, 1])

            return np.array([x_min, y_min, x_max, y_max])

        if len(bboxes_xywh) == 0:
            return bboxes_xywh
        bboxes_xyxy = xywh_to_xyxy(bboxes_xywh, image_shape=None)
        bboxes_xyxy = np.array([bbox_shift_scale_rotate(box, mat) for box in bboxes_xyxy])
        return xyxy_to_xywh(bboxes_xyxy, image_shape=None).astype(bboxes_xywh.dtype)

    @classmethod
    def apply_to_keypoints(cls, keypoints: np.ndarray, mat: np.ndarray, image_shape: Tuple[int, int]) -> np.ndarray:
        """
        Apply affine transform to keypoints.

        :param keypoints:   [N,K,3] array of keypoints in (x,y,visibility) format
        :param mat:         [2,3] Affine transformation matrix
        :param image_shape: Image shape after applying affine transform (height, width).
                            Used to update visibility status of keypoints.
        :return:            [N,K,3] array of keypoints in (x,y,visibility) format
        """
        keypoints_with_visibility = keypoints.copy()
        keypoints = keypoints_with_visibility[:, :, 0:2]

        shape = keypoints.shape
        dtype = keypoints.dtype
        keypoints = keypoints.reshape(-1, 2)
        keypoints = np.dot(np.concatenate((keypoints, keypoints[:, 0:1] * 0 + 1), axis=1), mat.T).reshape(shape)

        # Update visibility status of joints that were moved outside visible area
        image_height, image_width = image_shape[:2]
        outside_left = keypoints[:, :, 0] < 0
        outside_top = keypoints[:, :, 1] < 0
        outside_right = keypoints[:, :, 0] >= image_width
        outside_bottom = keypoints[:, :, 1] >= image_height

        joints_outside_image = outside_left | outside_top | outside_right | outside_bottom

        keypoints_with_visibility[:, :, 0:2] = keypoints
        keypoints_with_visibility[joints_outside_image, 2] = 0
        return keypoints_with_visibility.astype(dtype, copy=False)

    @classmethod
    def apply_to_image(cls, image: np.ndarray, mat: np.ndarray, interpolation: int, padding_value: Union[int, float, Tuple], padding_mode: int) -> np.ndarray:
        """
        Apply affine transform to image.

        :param image:          Input image
        :param mat:            [2,3] Affine transformation matrix
        :param interpolation:  Interpolation mode. See cv2.warpAffine for details.
        :param padding_value:  Value to pad the image during affine transform. See cv2.warpAffine for details.
        :param padding_mode:   Padding mode. See cv2.warpAffine for details.
        :return:               Transformed image of the same shape as input image.
        """
        return cv2.warpAffine(
            image,
            mat,
            dsize=(image.shape[1], image.shape[0]),
            flags=interpolation,
            borderValue=padding_value,
            borderMode=padding_mode,
        )

    def get_equivalent_preprocessing(self):
        raise RuntimeError(f"{self.__class__} does not have equivalent preprocessing because it is non-deterministic.")

__init__(max_rotation, min_scale, max_scale, max_translate, image_pad_value, mask_pad_value, interpolation_mode=cv2.INTER_LINEAR, prob=0.5)

Parameters:

Name Type Description Default
max_rotation float

Max rotation angle in degrees

 required
min_scale float

Lower bound for the scale change. For +- 20% size jitter this should be 0.8

 required
max_scale float

Lower bound for the scale change. For +- 20% size jitter this should be 1.2

 required
max_translate float

Max translation offset in percents of image size

 required
image_pad_value Union[int, float, List[int]]

Value to pad the image during affine transform. Can be single scalar or list. If a list is provided, it should have the same length as the number of channels in the image.

 required
mask_pad_value float

Value to pad the mask during affine transform.

 required
interpolation_mode Union[int, List[int]]

A constant integer or list of integers, specifying the interpolation mode to use. Possible values for interpolation_mode: cv2.INTER_NEAREST = 0, cv2.INTER_LINEAR = 1, cv2.INTER_CUBIC = 2, cv2.INTER_AREA = 3, cv2.INTER_LANCZOS4 = 4 To use random interpolation modes on each call, set interpolation_mode = (0,1,2,3,4)

 cv2.INTER_LINEAR
prob float

Probability to apply the transform.

 0.5
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_random_affine.py 
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
def __init__(
    self,
    max_rotation: float,
    min_scale: float,
    max_scale: float,
    max_translate: float,
    image_pad_value: Union[int, float, List[int]],
    mask_pad_value: float,
    interpolation_mode: Union[int, List[int]] = cv2.INTER_LINEAR,
    prob: float = 0.5,
):
    """

    :param max_rotation:       Max rotation angle in degrees
    :param min_scale:          Lower bound for the scale change. For +- 20% size jitter this should be 0.8
    :param max_scale:          Lower bound for the scale change. For +- 20% size jitter this should be 1.2
    :param max_translate:      Max translation offset in percents of image size
    :param image_pad_value:    Value to pad the image during affine transform. Can be single scalar or list.
                               If a list is provided, it should have the same length as the number of channels in the image.
    :param mask_pad_value:     Value to pad the mask during affine transform.
    :param interpolation_mode: A constant integer or list of integers, specifying the interpolation mode to use.
                               Possible values for interpolation_mode:
                                 cv2.INTER_NEAREST = 0,
                                 cv2.INTER_LINEAR = 1,
                                 cv2.INTER_CUBIC = 2,
                                 cv2.INTER_AREA = 3,
                                 cv2.INTER_LANCZOS4 = 4
                               To use random interpolation modes on each call, set interpolation_mode = (0,1,2,3,4)
    :param prob:               Probability to apply the transform.
    """
    super().__init__()

    self.max_rotation = max_rotation
    self.min_scale = min_scale
    self.max_scale = max_scale
    self.max_translate = max_translate
    self.image_pad_value = image_pad_value
    self.mask_pad_value = mask_pad_value
    self.prob = prob
    self.interpolation_mode = tuple(interpolation_mode) if isinstance(interpolation_mode, Iterable) else (interpolation_mode,)

apply_to_areas(areas, mat) classmethod

Apply affine transform to areas.

Parameters:

Name Type Description Default
areas np.ndarray

[N] Single-dimension array of areas

 required
mat np.ndarray

[2,3] Affine transformation matrix

 required

Returns:

Type Description
np.ndarray

[N] Single-dimension array of areas
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_random_affine.py 
134
135
136
137
138
139
140
141
142
143
144
@classmethod
def apply_to_areas(cls, areas: np.ndarray, mat: np.ndarray) -> np.ndarray:
    """
    Apply affine transform to areas.

    :param areas: [N] Single-dimension array of areas
    :param mat:   [2,3] Affine transformation matrix
    :return:      [N] Single-dimension array of areas
    """
    det = np.linalg.det(mat[:2, :2])
    return (areas * abs(det)).astype(areas.dtype)

apply_to_bboxes(bboxes_xywh, mat) classmethod

Parameters:

Name Type Description Default
bboxes

(N,4) array of bboxes in XYWH format

 required
mat np.ndarray

[2,3] Affine transformation matrix

 required

Returns:

Type Description
np.ndarray

(N,4) array of bboxes in XYWH format
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_random_affine.py 
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
@classmethod
def apply_to_bboxes(cls, bboxes_xywh: np.ndarray, mat: np.ndarray) -> np.ndarray:
    """

    :param bboxes: (N,4) array of bboxes in XYWH format
    :param mat:    [2,3] Affine transformation matrix
    :return:       (N,4) array of bboxes in XYWH format
    """

    def bbox_shift_scale_rotate(bbox, m):
        x_min, y_min, x_max, y_max = bbox[:4]

        x = np.array([x_min, x_max, x_max, x_min])
        y = np.array([y_min, y_min, y_max, y_max])
        ones = np.ones(shape=(len(x)))
        points_ones = np.vstack([x, y, ones]).transpose()

        tr_points = m.dot(points_ones.T).T

        x_min, x_max = min(tr_points[:, 0]), max(tr_points[:, 0])
        y_min, y_max = min(tr_points[:, 1]), max(tr_points[:, 1])

        return np.array([x_min, y_min, x_max, y_max])

    if len(bboxes_xywh) == 0:
        return bboxes_xywh
    bboxes_xyxy = xywh_to_xyxy(bboxes_xywh, image_shape=None)
    bboxes_xyxy = np.array([bbox_shift_scale_rotate(box, mat) for box in bboxes_xyxy])
    return xyxy_to_xywh(bboxes_xyxy, image_shape=None).astype(bboxes_xywh.dtype)

apply_to_image(image, mat, interpolation, padding_value, padding_mode) classmethod

Apply affine transform to image.

Parameters:

Name Type Description Default
image np.ndarray

Input image

 required
mat np.ndarray

[2,3] Affine transformation matrix

 required
interpolation int

Interpolation mode. See cv2.warpAffine for details.

 required
padding_value Union[int, float, Tuple]

Value to pad the image during affine transform. See cv2.warpAffine for details.

 required
padding_mode int

Padding mode. See cv2.warpAffine for details.

 required

Returns:

Type Description
np.ndarray

Transformed image of the same shape as input image.
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_random_affine.py 
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
@classmethod
def apply_to_image(cls, image: np.ndarray, mat: np.ndarray, interpolation: int, padding_value: Union[int, float, Tuple], padding_mode: int) -> np.ndarray:
    """
    Apply affine transform to image.

    :param image:          Input image
    :param mat:            [2,3] Affine transformation matrix
    :param interpolation:  Interpolation mode. See cv2.warpAffine for details.
    :param padding_value:  Value to pad the image during affine transform. See cv2.warpAffine for details.
    :param padding_mode:   Padding mode. See cv2.warpAffine for details.
    :return:               Transformed image of the same shape as input image.
    """
    return cv2.warpAffine(
        image,
        mat,
        dsize=(image.shape[1], image.shape[0]),
        flags=interpolation,
        borderValue=padding_value,
        borderMode=padding_mode,
    )

apply_to_keypoints(keypoints, mat, image_shape) classmethod

Apply affine transform to keypoints.

Parameters:

Name Type Description Default
keypoints np.ndarray

[N,K,3] array of keypoints in (x,y,visibility) format

 required
mat np.ndarray

[2,3] Affine transformation matrix

 required
image_shape Tuple[int, int]

Image shape after applying affine transform (height, width). Used to update visibility status of keypoints.

 required

Returns:

Type Description
np.ndarray

[N,K,3] array of keypoints in (x,y,visibility) format
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_random_affine.py 
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
@classmethod
def apply_to_keypoints(cls, keypoints: np.ndarray, mat: np.ndarray, image_shape: Tuple[int, int]) -> np.ndarray:
    """
    Apply affine transform to keypoints.

    :param keypoints:   [N,K,3] array of keypoints in (x,y,visibility) format
    :param mat:         [2,3] Affine transformation matrix
    :param image_shape: Image shape after applying affine transform (height, width).
                        Used to update visibility status of keypoints.
    :return:            [N,K,3] array of keypoints in (x,y,visibility) format
    """
    keypoints_with_visibility = keypoints.copy()
    keypoints = keypoints_with_visibility[:, :, 0:2]

    shape = keypoints.shape
    dtype = keypoints.dtype
    keypoints = keypoints.reshape(-1, 2)
    keypoints = np.dot(np.concatenate((keypoints, keypoints[:, 0:1] * 0 + 1), axis=1), mat.T).reshape(shape)

    # Update visibility status of joints that were moved outside visible area
    image_height, image_width = image_shape[:2]
    outside_left = keypoints[:, :, 0] < 0
    outside_top = keypoints[:, :, 1] < 0
    outside_right = keypoints[:, :, 0] >= image_width
    outside_bottom = keypoints[:, :, 1] >= image_height

    joints_outside_image = outside_left | outside_top | outside_right | outside_bottom

    keypoints_with_visibility[:, :, 0:2] = keypoints
    keypoints_with_visibility[joints_outside_image, 2] = 0
    return keypoints_with_visibility.astype(dtype, copy=False)

apply_to_sample(sample)

Apply transformation to given pose estimation sample. Since this transformation apply affine transform some keypoints/bboxes may be moved outside the image. After applying the transform, visibility status of joints is updated to reflect the new position of joints. Bounding boxes are clipped to image borders. If sample contains areas, they are scaled according to the applied affine transform.

Parameters:

Name Type Description Default
sample PoseEstimationSample

A pose estimation sample

 required

Returns:

Type Description
PoseEstimationSample

A transformed pose estimation sample
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_random_affine.py 
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
    """
    Apply transformation to given pose estimation sample.
    Since this transformation apply affine transform some keypoints/bboxes may be moved outside the image.
    After applying the transform, visibility status of joints is updated to reflect the new position of joints.
    Bounding boxes are clipped to image borders.
    If sample contains areas, they are scaled according to the applied affine transform.

    :param sample: A pose estimation sample
    :return:       A transformed pose estimation sample
    """

    if random.random() < self.prob:
        angle = random.uniform(-self.max_rotation, self.max_rotation)
        scale = random.uniform(self.min_scale, self.max_scale)
        dx = random.uniform(-self.max_translate, self.max_translate)
        dy = random.uniform(-self.max_translate, self.max_translate)
        interpolation = random.choice(self.interpolation_mode)

        mat_output = self._get_affine_matrix(sample.image, angle, scale, dx, dy)
        mat_output = mat_output[:2]

        image_pad_value = (
            tuple(self.image_pad_value) if isinstance(self.image_pad_value, Iterable) else tuple([self.image_pad_value] * sample.image.shape[-1])
        )

        sample.image = self.apply_to_image(
            sample.image, mat_output, interpolation=interpolation, padding_value=image_pad_value, padding_mode=cv2.BORDER_CONSTANT
        )
        sample.mask = self.apply_to_image(
            sample.mask, mat_output, interpolation=cv2.INTER_NEAREST, padding_value=self.mask_pad_value, padding_mode=cv2.BORDER_CONSTANT
        )

        sample.joints = self.apply_to_keypoints(sample.joints, mat_output, sample.image.shape[:2])

        if sample.bboxes_xywh is not None:
            sample.bboxes_xywh = self.apply_to_bboxes(sample.bboxes_xywh, mat_output)

        if sample.areas is not None:
            sample.areas = self.apply_to_areas(sample.areas, mat_output)

        sample = sample.sanitize_sample()

    return sample

KeypointsRandomHorizontalFlip

Bases: AbstractKeypointTransform

Flip image, mask and joints horizontally with a given probability.

Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_random_horisontal_flip.py 
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
@register_transform(Transforms.KeypointsRandomHorizontalFlip)
class KeypointsRandomHorizontalFlip(AbstractKeypointTransform):
    """
    Flip image, mask and joints horizontally with a given probability.
    """

    def __init__(self, flip_index: List[int], prob: float = 0.5):
        """

        :param flip_index: Indexes of keypoints on the flipped image. When doing left-right flip, left hand becomes right hand.
                           So this array contains order of keypoints on the flipped image. This is dataset specific and depends on
                           how keypoints are defined in dataset.
        :param prob: Probability of flipping
        """
        super().__init__()
        self.flip_index = flip_index
        self.prob = prob

    def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
        """
        Apply transformation to given pose estimation sample

        :param sample: Input pose estimation sample.
        :return:       A new pose estimation sample.
        """
        if sample.image.shape[:2] != sample.mask.shape[:2]:
            raise RuntimeError(f"Image shape ({sample.image.shape[:2]}) does not match mask shape ({sample.mask.shape[:2]}).")

        if random.random() < self.prob:
            sample.image = self.apply_to_image(sample.image)
            sample.mask = self.apply_to_image(sample.mask)
            rows, cols = sample.image.shape[:2]
            sample.joints = self.apply_to_keypoints(sample.joints, cols)

            if sample.bboxes_xywh is not None:
                sample.bboxes_xywh = self.apply_to_bboxes(sample.bboxes_xywh, cols)

        return sample

    def apply_to_image(self, image: np.ndarray) -> np.ndarray:
        """
        Flip image horizontally

        :param image: Input image
        :return:      Horizontally flipped image
        """
        return np.ascontiguousarray(np.fliplr(image))

    def apply_to_keypoints(self, keypoints: np.ndarray, cols: int) -> np.ndarray:
        """
        Flip keypoints horizontally

        :param keypoints: Input keypoints of [N,K,3] shape
        :param cols:      Image width
        :return:          Flipped keypoints  of [N,K,3] shape
        """
        keypoints = keypoints.copy()
        keypoints = keypoints[:, self.flip_index]
        keypoints[:, :, 0] = cols - keypoints[:, :, 0] - 1
        return keypoints

    def apply_to_bboxes(self, bboxes: np.ndarray, cols: int) -> np.ndarray:
        """
        Flip boxes horizontally

        :param bboxes: Input boxes of [N,4] shape in XYWH format
        :param cols:   Image width
        :return:       Flipped boxes of [N,4] shape in XYWH format
        """

        bboxes = bboxes.copy()
        bboxes[:, 0] = cols - (bboxes[:, 0] + bboxes[:, 2])
        return bboxes

    def __repr__(self):
        return self.__class__.__name__ + f"(flip_index={self.flip_index}, prob={self.prob})"

    def get_equivalent_preprocessing(self):
        raise RuntimeError(f"{self.__class__} does not have equivalent preprocessing because it is non-deterministic.")

__init__(flip_index, prob=0.5)

Parameters:

Name Type Description Default
flip_index List[int]

Indexes of keypoints on the flipped image. When doing left-right flip, left hand becomes right hand. So this array contains order of keypoints on the flipped image. This is dataset specific and depends on how keypoints are defined in dataset.

 required
prob float

Probability of flipping

 0.5
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_random_horisontal_flip.py 
17
18
19
20
21
22
23
24
25
26
27
def __init__(self, flip_index: List[int], prob: float = 0.5):
    """

    :param flip_index: Indexes of keypoints on the flipped image. When doing left-right flip, left hand becomes right hand.
                       So this array contains order of keypoints on the flipped image. This is dataset specific and depends on
                       how keypoints are defined in dataset.
    :param prob: Probability of flipping
    """
    super().__init__()
    self.flip_index = flip_index
    self.prob = prob

apply_to_bboxes(bboxes, cols)

Flip boxes horizontally

Parameters:

Name Type Description Default
bboxes np.ndarray

Input boxes of [N,4] shape in XYWH format

 required
cols int

Image width

 required

Returns:

Type Description
np.ndarray

Flipped boxes of [N,4] shape in XYWH format
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_random_horisontal_flip.py 
72
73
74
75
76
77
78
79
80
81
82
83
def apply_to_bboxes(self, bboxes: np.ndarray, cols: int) -> np.ndarray:
    """
    Flip boxes horizontally

    :param bboxes: Input boxes of [N,4] shape in XYWH format
    :param cols:   Image width
    :return:       Flipped boxes of [N,4] shape in XYWH format
    """

    bboxes = bboxes.copy()
    bboxes[:, 0] = cols - (bboxes[:, 0] + bboxes[:, 2])
    return bboxes

apply_to_image(image)

Flip image horizontally

Parameters:

Name Type Description Default
image np.ndarray

Input image

 required

Returns:

Type Description
np.ndarray

Horizontally flipped image
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_random_horisontal_flip.py 
50
51
52
53
54
55
56
57
def apply_to_image(self, image: np.ndarray) -> np.ndarray:
    """
    Flip image horizontally

    :param image: Input image
    :return:      Horizontally flipped image
    """
    return np.ascontiguousarray(np.fliplr(image))

apply_to_keypoints(keypoints, cols)

Flip keypoints horizontally

Parameters:

Name Type Description Default
keypoints np.ndarray

Input keypoints of [N,K,3] shape

 required
cols int

Image width

 required

Returns:

Type Description
np.ndarray

Flipped keypoints of [N,K,3] shape
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_random_horisontal_flip.py 
59
60
61
62
63
64
65
66
67
68
69
70
def apply_to_keypoints(self, keypoints: np.ndarray, cols: int) -> np.ndarray:
    """
    Flip keypoints horizontally

    :param keypoints: Input keypoints of [N,K,3] shape
    :param cols:      Image width
    :return:          Flipped keypoints  of [N,K,3] shape
    """
    keypoints = keypoints.copy()
    keypoints = keypoints[:, self.flip_index]
    keypoints[:, :, 0] = cols - keypoints[:, :, 0] - 1
    return keypoints

apply_to_sample(sample)

Apply transformation to given pose estimation sample

Parameters:

Name Type Description Default
sample PoseEstimationSample

Input pose estimation sample.

 required

Returns:

Type Description
PoseEstimationSample

A new pose estimation sample.
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_random_horisontal_flip.py 
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
    """
    Apply transformation to given pose estimation sample

    :param sample: Input pose estimation sample.
    :return:       A new pose estimation sample.
    """
    if sample.image.shape[:2] != sample.mask.shape[:2]:
        raise RuntimeError(f"Image shape ({sample.image.shape[:2]}) does not match mask shape ({sample.mask.shape[:2]}).")

    if random.random() < self.prob:
        sample.image = self.apply_to_image(sample.image)
        sample.mask = self.apply_to_image(sample.mask)
        rows, cols = sample.image.shape[:2]
        sample.joints = self.apply_to_keypoints(sample.joints, cols)

        if sample.bboxes_xywh is not None:
            sample.bboxes_xywh = self.apply_to_bboxes(sample.bboxes_xywh, cols)

    return sample

KeypointsRandomRotate90

Bases: AbstractKeypointTransform

Apply 90 degree rotations to the sample.

Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_random_rotate90.py 
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
@register_transform(Transforms.KeypointsRandomRotate90)
class KeypointsRandomRotate90(AbstractKeypointTransform):
    """
    Apply 90 degree rotations to the sample.
    """

    def __init__(
        self,
        prob: float = 0.5,
    ):
        """
        Initialize transform

        :param prob (float): Probability of applying the transform
        """
        super().__init__()
        self.prob = prob

    def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
        """
        :param   sample: Input PoseEstimationSample
        :return:         Result of applying the transform
        """

        if random.random() < self.prob:
            factor = random.randint(0, 3)

            image_rows, image_cols = sample.image.shape[:2]

            sample.image = self.apply_to_image(sample.image, factor)
            sample.mask = self.apply_to_image(sample.mask, factor)
            sample.joints = self.apply_to_keypoints(sample.joints, factor, image_rows, image_cols)

            if sample.bboxes_xywh is not None:
                sample.bboxes_xywh = self.apply_to_bboxes(sample.bboxes_xywh, factor, image_rows, image_cols)

        return sample

    @classmethod
    def apply_to_image(cls, image: np.ndarray, factor: int) -> np.ndarray:
        """
        Rotate image by 90 degrees

        :param image:  Input image
        :param factor: Number of 90 degree rotations to apply. Order or rotation matches np.rot90
        :return:       Rotated image
        """
        return np.rot90(image, factor)

    @classmethod
    def apply_to_bboxes(cls, bboxes_xywh: np.ndarray, factor, rows: int, cols: int) -> np.ndarray:
        """

        :param bboxes: (N, 4) array of bboxes in XYWH format
        :param factor: Number of 90 degree rotations to apply. Order or rotation matches np.rot90
        :param rows:   Number of rows (image height) of the original (input) image
        :param cols:   Number of cols (image width) of the original (input) image
        :return:       Transformed bboxes in XYWH format
        """
        from super_gradients.training.transforms.transforms import DetectionRandomRotate90

        bboxes_xyxy = xywh_to_xyxy(bboxes_xywh, image_shape=None)
        bboxes_xyxy = DetectionRandomRotate90.xyxy_bbox_rot90(bboxes_xyxy, factor, rows, cols)
        return xyxy_to_xywh(bboxes_xyxy, image_shape=None)

    @classmethod
    def apply_to_keypoints(cls, keypoints: np.ndarray, factor, rows: int, cols: int) -> np.ndarray:
        """

        :param keypoints: Input keypoints array of [Num Instances, Num Joints, 3] shape.
                          Keypoints has format (x, y, visibility)
        :param factor:    Number of 90 degree rotations to apply. Order or rotation matches np.rot90
        :param rows:      Number of rows (image height) of the original (input) image
        :param cols:      Number of cols (image width) of the original (input) image
        :return:          Transformed keypoints array of [Num Instances, Num Joints, 3] shape.
        """
        x, y, v = keypoints[:, :, 0], keypoints[:, :, 1], keypoints[:, :, 2]

        if factor == 0:
            keypoints = x, y, v
        elif factor == 1:
            keypoints = y, cols - x - 1, v
        elif factor == 2:
            keypoints = cols - x - 1, rows - y - 1, v
        elif factor == 3:
            keypoints = rows - y - 1, x, v
        else:
            raise ValueError("Parameter n must be in set {0, 1, 2, 3}")
        return np.stack(keypoints, axis=-1)

    def __repr__(self):
        return self.__class__.__name__ + f"(prob={self.prob})"

    def get_equivalent_preprocessing(self):
        raise RuntimeError(f"{self.__class__} does not have equivalent preprocessing because it is non-deterministic.")

__init__(prob=0.5)

Initialize transform

Parameters:

Name Type Description Default
(float) prob

Probability of applying the transform

 required
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_random_rotate90.py 
17
18
19
20
21
22
23
24
25
26
27
def __init__(
    self,
    prob: float = 0.5,
):
    """
    Initialize transform

    :param prob (float): Probability of applying the transform
    """
    super().__init__()
    self.prob = prob

apply_to_bboxes(bboxes_xywh, factor, rows, cols) classmethod

Parameters:

Name Type Description Default
bboxes

(N, 4) array of bboxes in XYWH format

 required
factor

Number of 90 degree rotations to apply. Order or rotation matches np.rot90

 required
rows int

Number of rows (image height) of the original (input) image

 required
cols int

Number of cols (image width) of the original (input) image

 required

Returns:

Type Description
np.ndarray

Transformed bboxes in XYWH format
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_random_rotate90.py 
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
@classmethod
def apply_to_bboxes(cls, bboxes_xywh: np.ndarray, factor, rows: int, cols: int) -> np.ndarray:
    """

    :param bboxes: (N, 4) array of bboxes in XYWH format
    :param factor: Number of 90 degree rotations to apply. Order or rotation matches np.rot90
    :param rows:   Number of rows (image height) of the original (input) image
    :param cols:   Number of cols (image width) of the original (input) image
    :return:       Transformed bboxes in XYWH format
    """
    from super_gradients.training.transforms.transforms import DetectionRandomRotate90

    bboxes_xyxy = xywh_to_xyxy(bboxes_xywh, image_shape=None)
    bboxes_xyxy = DetectionRandomRotate90.xyxy_bbox_rot90(bboxes_xyxy, factor, rows, cols)
    return xyxy_to_xywh(bboxes_xyxy, image_shape=None)

apply_to_image(image, factor) classmethod

Rotate image by 90 degrees

Parameters:

Name Type Description Default
image np.ndarray

Input image

 required
factor int

Number of 90 degree rotations to apply. Order or rotation matches np.rot90

 required

Returns:

Type Description
np.ndarray

Rotated image
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_random_rotate90.py 
49
50
51
52
53
54
55
56
57
58
@classmethod
def apply_to_image(cls, image: np.ndarray, factor: int) -> np.ndarray:
    """
    Rotate image by 90 degrees

    :param image:  Input image
    :param factor: Number of 90 degree rotations to apply. Order or rotation matches np.rot90
    :return:       Rotated image
    """
    return np.rot90(image, factor)

apply_to_keypoints(keypoints, factor, rows, cols) classmethod

Parameters:

Name Type Description Default
keypoints np.ndarray

Input keypoints array of [Num Instances, Num Joints, 3] shape. Keypoints has format (x, y, visibility)

 required
factor

Number of 90 degree rotations to apply. Order or rotation matches np.rot90

 required
rows int

Number of rows (image height) of the original (input) image

 required
cols int

Number of cols (image width) of the original (input) image

 required

Returns:

Type Description
np.ndarray

Transformed keypoints array of [Num Instances, Num Joints, 3] shape.
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_random_rotate90.py 
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
@classmethod
def apply_to_keypoints(cls, keypoints: np.ndarray, factor, rows: int, cols: int) -> np.ndarray:
    """

    :param keypoints: Input keypoints array of [Num Instances, Num Joints, 3] shape.
                      Keypoints has format (x, y, visibility)
    :param factor:    Number of 90 degree rotations to apply. Order or rotation matches np.rot90
    :param rows:      Number of rows (image height) of the original (input) image
    :param cols:      Number of cols (image width) of the original (input) image
    :return:          Transformed keypoints array of [Num Instances, Num Joints, 3] shape.
    """
    x, y, v = keypoints[:, :, 0], keypoints[:, :, 1], keypoints[:, :, 2]

    if factor == 0:
        keypoints = x, y, v
    elif factor == 1:
        keypoints = y, cols - x - 1, v
    elif factor == 2:
        keypoints = cols - x - 1, rows - y - 1, v
    elif factor == 3:
        keypoints = rows - y - 1, x, v
    else:
        raise ValueError("Parameter n must be in set {0, 1, 2, 3}")
    return np.stack(keypoints, axis=-1)

apply_to_sample(sample)

Returns:

Type Description
PoseEstimationSample

Result of applying the transform
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_random_rotate90.py 
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
    """
    :param   sample: Input PoseEstimationSample
    :return:         Result of applying the transform
    """

    if random.random() < self.prob:
        factor = random.randint(0, 3)

        image_rows, image_cols = sample.image.shape[:2]

        sample.image = self.apply_to_image(sample.image, factor)
        sample.mask = self.apply_to_image(sample.mask, factor)
        sample.joints = self.apply_to_keypoints(sample.joints, factor, image_rows, image_cols)

        if sample.bboxes_xywh is not None:
            sample.bboxes_xywh = self.apply_to_bboxes(sample.bboxes_xywh, factor, image_rows, image_cols)

    return sample

KeypointsRandomVerticalFlip

Bases: AbstractKeypointTransform

Flip image, mask and joints vertically with a given probability.

Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_random_vertical_flip.py 
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
@register_transform(Transforms.KeypointsRandomVerticalFlip)
class KeypointsRandomVerticalFlip(AbstractKeypointTransform):
    """
    Flip image, mask and joints vertically with a given probability.
    """

    def __init__(self, prob: float = 0.5):
        """

        :param prob: Probability of flipping
        """
        super().__init__()
        self.prob = prob

    def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
        """
        Apply transformation to given pose estimation sample

        :param sample: Input pose estimation sample.
        :return:       A new pose estimation sample.
        """
        if sample.image.shape[:2] != sample.mask.shape[:2]:
            raise RuntimeError(f"Image shape ({sample.image.shape[:2]}) does not match mask shape ({sample.mask.shape[:2]}).")

        if random.random() < self.prob:
            sample.image = self.apply_to_image(sample.image)
            sample.mask = self.apply_to_image(sample.mask)
            rows, cols = sample.image.shape[:2]
            sample.joints = self.apply_to_keypoints(sample.joints, rows)

            if sample.bboxes_xywh is not None:
                sample.bboxes_xywh = self.apply_to_bboxes(sample.bboxes_xywh, rows)

        return sample

    def apply_to_image(self, image: np.ndarray) -> np.ndarray:
        """
        Flip image vertically

        :param image: Input image
        :return:      Vertically flipped image
        """
        return np.ascontiguousarray(np.flipud(image))

    def apply_to_keypoints(self, keypoints: np.ndarray, rows: int) -> np.ndarray:
        """
        Flip keypoints vertically

        :param keypoints: Input keypoints of [N,K,3] shape
        :param rows:      Image height
        :return:          Flipped keypoints  of [N,K,3] shape
        """
        keypoints = keypoints.copy()
        keypoints[:, :, 1] = rows - keypoints[:, :, 1] - 1
        return keypoints

    def apply_to_bboxes(self, bboxes: np.ndarray, rows: int) -> np.ndarray:
        """
        Flip boxes vertically

        :param bboxes: Input boxes of [N,4] shape in XYWH format
        :param rows:   Image height
        :return:       Flipped boxes of [N,4] shape in XYWH format
        """

        bboxes = bboxes.copy()
        bboxes[:, 1] = rows - (bboxes[:, 1] + bboxes[:, 3]) - 1
        return bboxes

    def __repr__(self):
        return self.__class__.__name__ + f"(prob={self.prob})"

    def get_equivalent_preprocessing(self):
        raise RuntimeError(f"{self.__class__} does not have equivalent preprocessing because it is non-deterministic.")

__init__(prob=0.5)

Parameters:

Name Type Description Default
prob float

Probability of flipping

 0.5
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_random_vertical_flip.py 
16
17
18
19
20
21
22
def __init__(self, prob: float = 0.5):
    """

    :param prob: Probability of flipping
    """
    super().__init__()
    self.prob = prob

apply_to_bboxes(bboxes, rows)

Flip boxes vertically

Parameters:

Name Type Description Default
bboxes np.ndarray

Input boxes of [N,4] shape in XYWH format

 required
rows int

Image height

 required

Returns:

Type Description
np.ndarray

Flipped boxes of [N,4] shape in XYWH format
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_random_vertical_flip.py 
66
67
68
69
70
71
72
73
74
75
76
77
def apply_to_bboxes(self, bboxes: np.ndarray, rows: int) -> np.ndarray:
    """
    Flip boxes vertically

    :param bboxes: Input boxes of [N,4] shape in XYWH format
    :param rows:   Image height
    :return:       Flipped boxes of [N,4] shape in XYWH format
    """

    bboxes = bboxes.copy()
    bboxes[:, 1] = rows - (bboxes[:, 1] + bboxes[:, 3]) - 1
    return bboxes

apply_to_image(image)

Flip image vertically

Parameters:

Name Type Description Default
image np.ndarray

Input image

 required

Returns:

Type Description
np.ndarray

Vertically flipped image
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_random_vertical_flip.py 
45
46
47
48
49
50
51
52
def apply_to_image(self, image: np.ndarray) -> np.ndarray:
    """
    Flip image vertically

    :param image: Input image
    :return:      Vertically flipped image
    """
    return np.ascontiguousarray(np.flipud(image))

apply_to_keypoints(keypoints, rows)

Flip keypoints vertically

Parameters:

Name Type Description Default
keypoints np.ndarray

Input keypoints of [N,K,3] shape

 required
rows int

Image height

 required

Returns:

Type Description
np.ndarray

Flipped keypoints of [N,K,3] shape
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_random_vertical_flip.py 
54
55
56
57
58
59
60
61
62
63
64
def apply_to_keypoints(self, keypoints: np.ndarray, rows: int) -> np.ndarray:
    """
    Flip keypoints vertically

    :param keypoints: Input keypoints of [N,K,3] shape
    :param rows:      Image height
    :return:          Flipped keypoints  of [N,K,3] shape
    """
    keypoints = keypoints.copy()
    keypoints[:, :, 1] = rows - keypoints[:, :, 1] - 1
    return keypoints

apply_to_sample(sample)

Apply transformation to given pose estimation sample

Parameters:

Name Type Description Default
sample PoseEstimationSample

Input pose estimation sample.

 required

Returns:

Type Description
PoseEstimationSample

A new pose estimation sample.
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_random_vertical_flip.py 
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
    """
    Apply transformation to given pose estimation sample

    :param sample: Input pose estimation sample.
    :return:       A new pose estimation sample.
    """
    if sample.image.shape[:2] != sample.mask.shape[:2]:
        raise RuntimeError(f"Image shape ({sample.image.shape[:2]}) does not match mask shape ({sample.mask.shape[:2]}).")

    if random.random() < self.prob:
        sample.image = self.apply_to_image(sample.image)
        sample.mask = self.apply_to_image(sample.mask)
        rows, cols = sample.image.shape[:2]
        sample.joints = self.apply_to_keypoints(sample.joints, rows)

        if sample.bboxes_xywh is not None:
            sample.bboxes_xywh = self.apply_to_bboxes(sample.bboxes_xywh, rows)

    return sample

KeypointsRemoveSmallObjects

Bases: AbstractKeypointTransform

Remove pose instances from data sample that are too small or have too few visible keypoints.

Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_remove_small_objects.py 
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
@register_transform(Transforms.KeypointsRemoveSmallObjects)
class KeypointsRemoveSmallObjects(AbstractKeypointTransform):
    """
    Remove pose instances from data sample that are too small or have too few visible keypoints.
    """

    def __init__(self, min_visible_keypoints: int = 0, min_instance_area: int = 0, min_bbox_area: int = 0):
        """

        :param min_visible_keypoints: Minimum number of visible keypoints to keep the sample.
                                      Default value is 0 which means that all samples will be kept.
        :param min_instance_area:     Minimum area of instance area to keep the sample
                                      Default value is 0 which means that all samples will be kept.
        :param min_bbox_area:         Minimum area of bounding box to keep the sample
                                      Default value is 0 which means that all samples will be kept.
        """
        super().__init__()
        self.min_visible_keypoints = min_visible_keypoints
        self.min_instance_area = min_instance_area
        self.min_bbox_area = min_bbox_area

    def __call__(
        self, image: np.ndarray, mask: np.ndarray, joints: np.ndarray, areas: Optional[np.ndarray], bboxes: Optional[np.ndarray]
    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray], Optional[np.ndarray]]:
        raise RuntimeError("This transform is not supported for old-style API")

    def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
        """
        Apply transformation to given pose estimation sample.

        :param sample: Input sample to transform.
        :return:       Filtered sample.
        """
        if self.min_visible_keypoints:
            sample = sample.filter_by_visible_joints(self.min_visible_keypoints)
        if self.min_instance_area:
            sample = sample.filter_by_pose_area(self.min_instance_area)
        if self.min_bbox_area:
            sample = sample.filter_by_bbox_area(self.min_bbox_area)
        return sample

    def __repr__(self):
        return self.__class__.__name__ + (
            f"(min_visible_keypoints={self.min_visible_keypoints}, " f"min_instance_area={self.min_instance_area}, " f"min_bbox_area={self.min_bbox_area})"
        )

    def get_equivalent_preprocessing(self) -> List:
        return []

__init__(min_visible_keypoints=0, min_instance_area=0, min_bbox_area=0)

Parameters:

Name Type Description Default
min_visible_keypoints int

Minimum number of visible keypoints to keep the sample. Default value is 0 which means that all samples will be kept.

 0
min_instance_area int

Minimum area of instance area to keep the sample Default value is 0 which means that all samples will be kept.

 0
min_bbox_area int

Minimum area of bounding box to keep the sample Default value is 0 which means that all samples will be kept.

 0
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_remove_small_objects.py 
17
18
19
20
21
22
23
24
25
26
27
28
29
30
def __init__(self, min_visible_keypoints: int = 0, min_instance_area: int = 0, min_bbox_area: int = 0):
    """

    :param min_visible_keypoints: Minimum number of visible keypoints to keep the sample.
                                  Default value is 0 which means that all samples will be kept.
    :param min_instance_area:     Minimum area of instance area to keep the sample
                                  Default value is 0 which means that all samples will be kept.
    :param min_bbox_area:         Minimum area of bounding box to keep the sample
                                  Default value is 0 which means that all samples will be kept.
    """
    super().__init__()
    self.min_visible_keypoints = min_visible_keypoints
    self.min_instance_area = min_instance_area
    self.min_bbox_area = min_bbox_area

apply_to_sample(sample)

Apply transformation to given pose estimation sample.

Parameters:

Name Type Description Default
sample PoseEstimationSample

Input sample to transform.

 required

Returns:

Type Description
PoseEstimationSample

Filtered sample.
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_remove_small_objects.py 
37
38
39
40
41
42
43
44
45
46
47
48
49
50
def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
    """
    Apply transformation to given pose estimation sample.

    :param sample: Input sample to transform.
    :return:       Filtered sample.
    """
    if self.min_visible_keypoints:
        sample = sample.filter_by_visible_joints(self.min_visible_keypoints)
    if self.min_instance_area:
        sample = sample.filter_by_pose_area(self.min_instance_area)
    if self.min_bbox_area:
        sample = sample.filter_by_bbox_area(self.min_bbox_area)
    return sample

KeypointsRescale

Bases: AbstractKeypointTransform

Resize image, mask and joints to target size without preserving aspect ratio.

Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_rescale.py 
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
@register_transform(Transforms.KeypointsRescale)
class KeypointsRescale(AbstractKeypointTransform):
    """
    Resize image, mask and joints to target size without preserving aspect ratio.
    """

    def __init__(self, height: int, width: int, interpolation: int = cv2.INTER_LINEAR, prob: float = 1.0):
        """
        :param height: Target image height
        :param width: Target image width
        :param interpolation: Used interpolation method for image. See cv2.resize for details.
        :param prob: Probability of applying this transform. Default value is 1, meaning that transform is always applied.
        """
        super().__init__()
        self.height = height
        self.width = width
        self.interpolation = interpolation
        self.prob = prob

    def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
        """
        Apply transform to sample.
        :param sample: Input sample
        :return:       Output sample
        """
        if random.random() < self.prob:
            height, width = sample.image.shape[:2]
            sy = self.height / height
            sx = self.width / width

            sample.image = self.apply_to_image(sample.image, dsize=(self.width, self.height), interpolation=self.interpolation)
            sample.mask = self.apply_to_image(sample.mask, dsize=(self.width, self.height), interpolation=cv2.INTER_NEAREST)

            sample.joints = self.apply_to_keypoints(sample.joints, sx, sy)
            if sample.bboxes_xywh is not None:
                sample.bboxes_xywh = self.apply_to_bboxes(sample.bboxes_xywh, sx, sy)

            if sample.areas is not None:
                sample.areas = np.multiply(sample.areas, sx * sy, dtype=np.float32)

        return sample

    @classmethod
    def apply_to_image(cls, img, dsize: Tuple[int, int], interpolation: int) -> np.ndarray:
        """
        Resize image to target size.
        :param img:           Input image
        :param dsize:         Target size (width, height)
        :param interpolation: OpenCV interpolation method
        :return:              Resize image
        """
        img = cv2.resize(img, dsize=dsize, interpolation=interpolation)
        return img

    @classmethod
    def apply_to_keypoints(cls, keypoints: np.ndarray, sx: float, sy: float) -> np.ndarray:
        """
        Resize keypoints to target size.
        :param keypoints: [Num Instances, Num Joints, 3] Input keypoints
        :param sx:        Scale factor along the horizontal axis
        :param sy:        Scale factor along the vertical axis
        :return:          [Num Instances, Num Joints, 3] Resized keypoints
        """
        keypoints = keypoints.astype(np.float32, copy=True)
        keypoints[:, :, 0] *= sx
        keypoints[:, :, 1] *= sy
        return keypoints

    @classmethod
    def apply_to_bboxes(cls, bboxes: np.ndarray, sx: float, sy: float) -> np.ndarray:
        """
        Resize bounding boxes to target size.

        :param bboxes: Input bounding boxes in XYWH format
        :param sx:     Scale factor along the horizontal axis
        :param sy:     Scale factor along the vertical axis
        :return:       Resized bounding boxes in XYWH format
        """
        bboxes = bboxes.astype(np.float32, copy=True)
        bboxes[:, 0::2] *= sx
        bboxes[:, 1::2] *= sy
        return bboxes

    @classmethod
    def apply_to_areas(cls, areas: np.ndarray, sx: float, sy: float) -> np.ndarray:
        """
        Resize areas to target size.
        :param areas: [N] Array of instance areas
        :param sx:    Scale factor along the horizontal axis
        :param sy:    Scale factor along the vertical axis
        :return:      [N] Array of resized instance areas
        """
        return np.multiply(areas, sx * sy, dtype=np.float32)

    def __repr__(self):
        return self.__class__.__name__ + f"(height={self.height}, " f"width={self.width}, " f"interpolation={self.interpolation}, prob={self.prob})"

    def get_equivalent_preprocessing(self) -> List:
        return [{Processings.KeypointsRescale: {"output_shape": (self.height, self.width)}}]

__init__(height, width, interpolation=cv2.INTER_LINEAR, prob=1.0)

Parameters:

Name Type Description Default
height int

Target image height

 required
width int

Target image width

 required
interpolation int

Used interpolation method for image. See cv2.resize for details.

 cv2.INTER_LINEAR
prob float

Probability of applying this transform. Default value is 1, meaning that transform is always applied.

 1.0
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_rescale.py 
18
19
20
21
22
23
24
25
26
27
28
29
def __init__(self, height: int, width: int, interpolation: int = cv2.INTER_LINEAR, prob: float = 1.0):
    """
    :param height: Target image height
    :param width: Target image width
    :param interpolation: Used interpolation method for image. See cv2.resize for details.
    :param prob: Probability of applying this transform. Default value is 1, meaning that transform is always applied.
    """
    super().__init__()
    self.height = height
    self.width = width
    self.interpolation = interpolation
    self.prob = prob

apply_to_areas(areas, sx, sy) classmethod

Resize areas to target size.

Parameters:

Name Type Description Default
areas np.ndarray

[N] Array of instance areas

 required
sx float

Scale factor along the horizontal axis

 required
sy float

Scale factor along the vertical axis

 required

Returns:

Type Description
np.ndarray

[N] Array of resized instance areas
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_rescale.py 
 95
 96
 97
 98
 99
100
101
102
103
104
@classmethod
def apply_to_areas(cls, areas: np.ndarray, sx: float, sy: float) -> np.ndarray:
    """
    Resize areas to target size.
    :param areas: [N] Array of instance areas
    :param sx:    Scale factor along the horizontal axis
    :param sy:    Scale factor along the vertical axis
    :return:      [N] Array of resized instance areas
    """
    return np.multiply(areas, sx * sy, dtype=np.float32)

apply_to_bboxes(bboxes, sx, sy) classmethod

Resize bounding boxes to target size.

Parameters:

Name Type Description Default
bboxes np.ndarray

Input bounding boxes in XYWH format

 required
sx float

Scale factor along the horizontal axis

 required
sy float

Scale factor along the vertical axis

 required

Returns:

Type Description
np.ndarray

Resized bounding boxes in XYWH format
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_rescale.py 
80
81
82
83
84
85
86
87
88
89
90
91
92
93
@classmethod
def apply_to_bboxes(cls, bboxes: np.ndarray, sx: float, sy: float) -> np.ndarray:
    """
    Resize bounding boxes to target size.

    :param bboxes: Input bounding boxes in XYWH format
    :param sx:     Scale factor along the horizontal axis
    :param sy:     Scale factor along the vertical axis
    :return:       Resized bounding boxes in XYWH format
    """
    bboxes = bboxes.astype(np.float32, copy=True)
    bboxes[:, 0::2] *= sx
    bboxes[:, 1::2] *= sy
    return bboxes

apply_to_image(img, dsize, interpolation) classmethod

Resize image to target size.

Parameters:

Name Type Description Default
img

Input image

 required
dsize Tuple[int, int]

Target size (width, height)

 required
interpolation int

OpenCV interpolation method

 required

Returns:

Type Description
np.ndarray

Resize image
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_rescale.py 
54
55
56
57
58
59
60
61
62
63
64
@classmethod
def apply_to_image(cls, img, dsize: Tuple[int, int], interpolation: int) -> np.ndarray:
    """
    Resize image to target size.
    :param img:           Input image
    :param dsize:         Target size (width, height)
    :param interpolation: OpenCV interpolation method
    :return:              Resize image
    """
    img = cv2.resize(img, dsize=dsize, interpolation=interpolation)
    return img

apply_to_keypoints(keypoints, sx, sy) classmethod

Resize keypoints to target size.

Parameters:

Name Type Description Default
keypoints np.ndarray

[Num Instances, Num Joints, 3] Input keypoints

 required
sx float

Scale factor along the horizontal axis

 required
sy float

Scale factor along the vertical axis

 required

Returns:

Type Description
np.ndarray

[Num Instances, Num Joints, 3] Resized keypoints
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_rescale.py 
66
67
68
69
70
71
72
73
74
75
76
77
78
@classmethod
def apply_to_keypoints(cls, keypoints: np.ndarray, sx: float, sy: float) -> np.ndarray:
    """
    Resize keypoints to target size.
    :param keypoints: [Num Instances, Num Joints, 3] Input keypoints
    :param sx:        Scale factor along the horizontal axis
    :param sy:        Scale factor along the vertical axis
    :return:          [Num Instances, Num Joints, 3] Resized keypoints
    """
    keypoints = keypoints.astype(np.float32, copy=True)
    keypoints[:, :, 0] *= sx
    keypoints[:, :, 1] *= sy
    return keypoints

apply_to_sample(sample)

Apply transform to sample.

Parameters:

Name Type Description Default
sample PoseEstimationSample

Input sample

 required

Returns:

Type Description
PoseEstimationSample

Output sample
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_rescale.py 
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
    """
    Apply transform to sample.
    :param sample: Input sample
    :return:       Output sample
    """
    if random.random() < self.prob:
        height, width = sample.image.shape[:2]
        sy = self.height / height
        sx = self.width / width

        sample.image = self.apply_to_image(sample.image, dsize=(self.width, self.height), interpolation=self.interpolation)
        sample.mask = self.apply_to_image(sample.mask, dsize=(self.width, self.height), interpolation=cv2.INTER_NEAREST)

        sample.joints = self.apply_to_keypoints(sample.joints, sx, sy)
        if sample.bboxes_xywh is not None:
            sample.bboxes_xywh = self.apply_to_bboxes(sample.bboxes_xywh, sx, sy)

        if sample.areas is not None:
            sample.areas = np.multiply(sample.areas, sx * sy, dtype=np.float32)

    return sample

AbstractKeypointTransform

Bases: abc.ABC

Base class for all transforms for keypoints augmentation. All transforms subclassing it should implement call method which takes image, mask and keypoints as input and returns transformed image, mask and keypoints.

Parameters:

Name Type Description Default
additional_samples_count int

Number of additional samples to generate for each image. This property is used for mixup & mosaic transforms that needs an extra samples.

 0
Source code in V3_4/src/super_gradients/training/transforms/keypoints/abstract_keypoints_transform.py 
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
class AbstractKeypointTransform(abc.ABC):
    """
    Base class for all transforms for keypoints augmentation.
    All transforms subclassing it should implement __call__ method which takes image, mask and keypoints as input and
    returns transformed image, mask and keypoints.

    :param additional_samples_count: Number of additional samples to generate for each image.
                                    This property is used for mixup & mosaic transforms that needs an extra samples.
    """

    def __init__(self, additional_samples_count: int = 0):
        """
        :param additional_samples_count: (int) number of samples that must be extra samples from dataset. Default value is 0.
        """
        self.additional_samples_count = additional_samples_count

    def __call__(
        self, image: np.ndarray, mask: np.ndarray, joints: np.ndarray, areas: Optional[np.ndarray], bboxes: Optional[np.ndarray]
    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray], Optional[np.ndarray]]:
        """
        Apply transformation to pose estimation sample passed as a tuple
        This method acts as a wrapper for apply_to_sample method to support old-style API.
        """
        sample = PoseEstimationSample(
            image=image,
            mask=mask,
            joints=joints,
            areas=areas,
            bboxes_xywh=bboxes,
            is_crowd=np.zeros(len(joints)),  # Old style API does not pass is_crowd parameter, so we set it to zeros
            additional_samples=None,
        )
        sample = self.apply_to_sample(sample)
        return sample.image, sample.mask, sample.joints, sample.areas, sample.bboxes_xywh

    @abstractmethod
    def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
        """
        Apply transformation to given pose estimation sample.
        Important note - function call may return new object, may modify it in-place.
        This is implementation dependent and if you need to keep original sample intact it
        is recommended to make a copy of it BEFORE passing it to transform.

        :param sample: Input sample to transform.
        :return:       Modified sample (It can be the same instance as input or a new object).
        """
        raise NotImplementedError

    @abstractmethod
    def get_equivalent_preprocessing(self) -> List:
        raise NotImplementedError

__call__(image, mask, joints, areas, bboxes)

Apply transformation to pose estimation sample passed as a tuple This method acts as a wrapper for apply_to_sample method to support old-style API.

Source code in V3_4/src/super_gradients/training/transforms/keypoints/abstract_keypoints_transform.py 
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
def __call__(
    self, image: np.ndarray, mask: np.ndarray, joints: np.ndarray, areas: Optional[np.ndarray], bboxes: Optional[np.ndarray]
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray], Optional[np.ndarray]]:
    """
    Apply transformation to pose estimation sample passed as a tuple
    This method acts as a wrapper for apply_to_sample method to support old-style API.
    """
    sample = PoseEstimationSample(
        image=image,
        mask=mask,
        joints=joints,
        areas=areas,
        bboxes_xywh=bboxes,
        is_crowd=np.zeros(len(joints)),  # Old style API does not pass is_crowd parameter, so we set it to zeros
        additional_samples=None,
    )
    sample = self.apply_to_sample(sample)
    return sample.image, sample.mask, sample.joints, sample.areas, sample.bboxes_xywh

__init__(additional_samples_count=0)

Parameters:

Name Type Description Default
additional_samples_count int

(int) number of samples that must be extra samples from dataset. Default value is 0.

 0
Source code in V3_4/src/super_gradients/training/transforms/keypoints/abstract_keypoints_transform.py 
20
21
22
23
24
def __init__(self, additional_samples_count: int = 0):
    """
    :param additional_samples_count: (int) number of samples that must be extra samples from dataset. Default value is 0.
    """
    self.additional_samples_count = additional_samples_count

apply_to_sample(sample) abstractmethod

Apply transformation to given pose estimation sample. Important note - function call may return new object, may modify it in-place. This is implementation dependent and if you need to keep original sample intact it is recommended to make a copy of it BEFORE passing it to transform.

Parameters:

Name Type Description Default
sample PoseEstimationSample

Input sample to transform.

 required

Returns:

Type Description
PoseEstimationSample

Modified sample (It can be the same instance as input or a new object).
Source code in V3_4/src/super_gradients/training/transforms/keypoints/abstract_keypoints_transform.py 
45
46
47
48
49
50
51
52
53
54
55
56
@abstractmethod
def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
    """
    Apply transformation to given pose estimation sample.
    Important note - function call may return new object, may modify it in-place.
    This is implementation dependent and if you need to keep original sample intact it
    is recommended to make a copy of it BEFORE passing it to transform.

    :param sample: Input sample to transform.
    :return:       Modified sample (It can be the same instance as input or a new object).
    """
    raise NotImplementedError

KeypointsBrightnessContrast

Bases: AbstractKeypointTransform

Apply brightness and contrast change to the input image using following formula: image = (image - mean_brightness) * contrast_gain + mean_brightness + brightness_gain Transformation preserves input image dtype. Saturation cast is performed at the end of the transformation.

Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_brightness_contrast.py 
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
@register_transform()
class KeypointsBrightnessContrast(AbstractKeypointTransform):
    """
    Apply brightness and contrast change to the input image using following formula:
    image = (image - mean_brightness) * contrast_gain + mean_brightness + brightness_gain
    Transformation preserves input image dtype. Saturation cast is performed at the end of the transformation.
    """

    def __init__(self, prob: float, brightness_range: Tuple[float, float], contrast_range: Tuple[float, float]):
        """

        :param prob:             Probability to apply the transform.
        :param brightness_range: Tuple of two elements, min and max brightness gain. Represents a relative range of
                                 brightness gain with respect to average image brightness. A brightness gain of 1.0
                                 indicates no change in brightness. Therefore, optimal value for this parameter is
                                 somewhere inside (0, 2) range.
        :param contrast_range:   Tuple of two elements, min and max contrast gain. Effective contrast_gain would be
                                 uniformly sampled from this range. Based on definition of contrast gain, it's optimal
                                 value is somewhere inside (0, 2) range.
        """
        if len(brightness_range) != 2:
            raise ValueError("Brightness range must be a tuple of two elements, got: " + str(brightness_range))
        if len(contrast_range) != 2:
            raise ValueError("Contrast range must be a tuple of two elements, got: " + str(contrast_range))
        super().__init__()
        self.prob = prob
        self.brightness_range = tuple(brightness_range)
        self.contrast_range = tuple(contrast_range)

    def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
        if random.random() < self.prob:
            contrast_gain = random.uniform(self.contrast_range[0], self.contrast_range[1])
            brightness_gain = random.uniform(self.brightness_range[0], self.brightness_range[1])

            input_dtype = sample.image.dtype
            image = sample.image.astype(np.float32)
            mean_brightness = np.mean(image, axis=(0, 1))

            image = (image - mean_brightness) * contrast_gain + mean_brightness * brightness_gain

            # get min & max values for the input_dtype
            min_value = np.iinfo(input_dtype).min
            max_value = np.iinfo(input_dtype).max
            sample.image = np.clip(image, a_min=min_value, a_max=max_value).astype(input_dtype)
        return sample

    def get_equivalent_preprocessing(self) -> List:
        raise RuntimeError(f"{self.__class__} does not have equivalent preprocessing because it is non-deterministic.")

__init__(prob, brightness_range, contrast_range)

Parameters:

Name Type Description Default
prob float

Probability to apply the transform.

 required
brightness_range Tuple[float, float]

Tuple of two elements, min and max brightness gain. Represents a relative range of brightness gain with respect to average image brightness. A brightness gain of 1.0 indicates no change in brightness. Therefore, optimal value for this parameter is somewhere inside (0, 2) range.

 required
contrast_range Tuple[float, float]

Tuple of two elements, min and max contrast gain. Effective contrast_gain would be uniformly sampled from this range. Based on definition of contrast gain, it's optimal value is somewhere inside (0, 2) range.

 required
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_brightness_contrast.py 
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
def __init__(self, prob: float, brightness_range: Tuple[float, float], contrast_range: Tuple[float, float]):
    """

    :param prob:             Probability to apply the transform.
    :param brightness_range: Tuple of two elements, min and max brightness gain. Represents a relative range of
                             brightness gain with respect to average image brightness. A brightness gain of 1.0
                             indicates no change in brightness. Therefore, optimal value for this parameter is
                             somewhere inside (0, 2) range.
    :param contrast_range:   Tuple of two elements, min and max contrast gain. Effective contrast_gain would be
                             uniformly sampled from this range. Based on definition of contrast gain, it's optimal
                             value is somewhere inside (0, 2) range.
    """
    if len(brightness_range) != 2:
        raise ValueError("Brightness range must be a tuple of two elements, got: " + str(brightness_range))
    if len(contrast_range) != 2:
        raise ValueError("Contrast range must be a tuple of two elements, got: " + str(contrast_range))
    super().__init__()
    self.prob = prob
    self.brightness_range = tuple(brightness_range)
    self.contrast_range = tuple(contrast_range)

KeypointsCompose

Bases: AbstractKeypointTransform

Composes several transforms together

Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_compose.py 
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
class KeypointsCompose(AbstractKeypointTransform):
    """
    Composes several transforms together
    """

    def __init__(self, transforms: List[AbstractKeypointTransform], load_sample_fn=None):
        """

        :param transforms:         List of keypoint-based transformations
        :param load_sample_fn:     A method to load additional samples if needed (for mixup & mosaic augmentations).
                                   Default value is None, which would raise an error if additional samples are needed.
        """
        for transform in transforms:
            if load_sample_fn is None and transform.additional_samples_count > 0:
                raise RuntimeError(
                    f"Detected transform {transform.__class__.__name__} that require {transform.additional_samples_count} "
                    f"additional samples, but load_sample_fn is None"
                )

        super().__init__()
        self.transforms = transforms
        self.load_sample_fn = load_sample_fn

    def __call__(
        self, image: np.ndarray, mask: np.ndarray, joints: np.ndarray, areas: Optional[np.ndarray], bboxes: Optional[np.ndarray]
    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray], Optional[np.ndarray]]:
        """
        Apply transformation to pose estimation sample passed as a tuple
        This method acts as a wrapper for apply_to_sample method to support old-style API.
        """
        for transform in self.transforms:
            if transform.additional_samples_count > 0:
                raise RuntimeError(f"{transform.__class__.__name__} require additional samples that is not supported in old-style transforms API")

        for t in self.transforms:
            image, mask, joints, areas, bboxes = t(image, mask, joints, areas, bboxes)

        return image, mask, joints, areas, bboxes

    def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
        """
        Applies the series of transformations to the input sample.
        The function may modify the input sample inplace, so input sample should not be used after the call.

        :param sample: Input sample
        :return:       Transformed sample.
        """
        sample = sample.sanitize_sample()
        sample = self._apply_transforms(sample, transforms=self.transforms, load_sample_fn=self.load_sample_fn)
        return sample

    @classmethod
    def _apply_transforms(cls, sample: PoseEstimationSample, transforms: List[AbstractKeypointTransform], load_sample_fn) -> PoseEstimationSample:
        """
        This helper method allows us to query additional samples for mixup & mosaic augmentations
        that would be also passed through augmentation pipeline. Example:

        ```
          transforms:
            - KeypointsBrightnessContrast:
                brightness_range: [ 0.8, 1.2 ]
                contrast_range: [ 0.8, 1.2 ]
                prob: 0.5
            - KeypointsHSV:
                hgain: 20
                sgain: 20
                vgain: 20
                prob: 0.5
            - KeypointsLongestMaxSize:
                max_height: ${dataset_params.image_size}
                max_width: ${dataset_params.image_size}
            - KeypointsMixup:
                prob: ${dataset_params.mixup_prob}
        ```

        In the example above all samples in mixup will be forwarded through KeypointsBrightnessContrast, KeypointsHSV,
        KeypointsLongestMaxSize and only then mixed up.

        :param sample:         Input data sample
        :param transforms:     List of transformations to apply
        :param load_sample_fn: A method to load additional samples if needed
        :return:               A data sample after applying transformations
        """
        applied_transforms_so_far = []
        for t in transforms:
            if t.additional_samples_count == 0:
                sample = t.apply_to_sample(sample)
                applied_transforms_so_far.append(t)
            else:
                additional_samples = [load_sample_fn() for _ in range(t.additional_samples_count)]
                additional_samples = [
                    cls._apply_transforms(
                        sample,
                        applied_transforms_so_far,
                        load_sample_fn=load_sample_fn,
                    )
                    for sample in additional_samples
                ]
                sample.additional_samples = additional_samples
                sample = t.apply_to_sample(sample)

        return sample

    def get_equivalent_preprocessing(self) -> List:
        preprocessing = []
        for t in self.transforms:
            preprocessing += t.get_equivalent_preprocessing()
        return preprocessing

    def __repr__(self):
        format_string = self.__class__.__name__ + "("
        for t in self.transforms:
            format_string += f"\t{repr(t)}"
        format_string += "\n)"
        return format_string

__call__(image, mask, joints, areas, bboxes)

Apply transformation to pose estimation sample passed as a tuple This method acts as a wrapper for apply_to_sample method to support old-style API.

Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_compose.py 
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
def __call__(
    self, image: np.ndarray, mask: np.ndarray, joints: np.ndarray, areas: Optional[np.ndarray], bboxes: Optional[np.ndarray]
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray], Optional[np.ndarray]]:
    """
    Apply transformation to pose estimation sample passed as a tuple
    This method acts as a wrapper for apply_to_sample method to support old-style API.
    """
    for transform in self.transforms:
        if transform.additional_samples_count > 0:
            raise RuntimeError(f"{transform.__class__.__name__} require additional samples that is not supported in old-style transforms API")

    for t in self.transforms:
        image, mask, joints, areas, bboxes = t(image, mask, joints, areas, bboxes)

    return image, mask, joints, areas, bboxes

__init__(transforms, load_sample_fn=None)

Parameters:

Name Type Description Default
transforms List[AbstractKeypointTransform]

List of keypoint-based transformations

 required
load_sample_fn

A method to load additional samples if needed (for mixup & mosaic augmentations). Default value is None, which would raise an error if additional samples are needed.

 None
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_compose.py 
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
def __init__(self, transforms: List[AbstractKeypointTransform], load_sample_fn=None):
    """

    :param transforms:         List of keypoint-based transformations
    :param load_sample_fn:     A method to load additional samples if needed (for mixup & mosaic augmentations).
                               Default value is None, which would raise an error if additional samples are needed.
    """
    for transform in transforms:
        if load_sample_fn is None and transform.additional_samples_count > 0:
            raise RuntimeError(
                f"Detected transform {transform.__class__.__name__} that require {transform.additional_samples_count} "
                f"additional samples, but load_sample_fn is None"
            )

    super().__init__()
    self.transforms = transforms
    self.load_sample_fn = load_sample_fn

apply_to_sample(sample)

Applies the series of transformations to the input sample. The function may modify the input sample inplace, so input sample should not be used after the call.

Parameters:

Name Type Description Default
sample PoseEstimationSample

Input sample

 required

Returns:

Type Description
PoseEstimationSample

Transformed sample.
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_compose.py 
48
49
50
51
52
53
54
55
56
57
58
def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
    """
    Applies the series of transformations to the input sample.
    The function may modify the input sample inplace, so input sample should not be used after the call.

    :param sample: Input sample
    :return:       Transformed sample.
    """
    sample = sample.sanitize_sample()
    sample = self._apply_transforms(sample, transforms=self.transforms, load_sample_fn=self.load_sample_fn)
    return sample

KeypointsHSV

Bases: AbstractKeypointTransform

Apply color change in HSV color space to the input image.

Parameters:

Name Type Description Default
prob float

Probability to apply the transform.

 required
hgain float

Hue gain.

 required
sgain float

Saturation gain.

 required
vgain float

Value gain.

 required
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_hsv.py 
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
@register_transform()
class KeypointsHSV(AbstractKeypointTransform):
    """
    Apply color change in HSV color space to the input image.

    :param prob:            Probability to apply the transform.
    :param hgain:           Hue gain.
    :param sgain:           Saturation gain.
    :param vgain:           Value gain.
    """

    def __init__(self, prob: float, hgain: float, sgain: float, vgain: float):
        """

        :param prob:            Probability to apply the transform.
        :param hgain:           Hue gain.
        :param sgain:           Saturation gain.
        :param vgain:           Value gain.
        """
        super().__init__()
        self.prob = prob
        self.hgain = hgain
        self.sgain = sgain
        self.vgain = vgain

    def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
        if sample.image.shape[2] != 3:
            raise ValueError("HSV transform expects image with 3 channels, got: " + str(sample.image.shape[2]))

        if random.random() < self.prob:
            image_copy = sample.image.copy()
            augment_hsv(image_copy, self.hgain, self.sgain, self.vgain, bgr_channels=(0, 1, 2))
            sample.image = image_copy
        return sample

    def get_equivalent_preprocessing(self):
        raise RuntimeError(f"{self.__class__} does not have equivalent preprocessing because it is non-deterministic.")

__init__(prob, hgain, sgain, vgain)

Parameters:

Name Type Description Default
prob float

Probability to apply the transform.

 required
hgain float

Hue gain.

 required
sgain float

Saturation gain.

 required
vgain float

Value gain.

 required
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_hsv.py 
21
22
23
24
25
26
27
28
29
30
31
32
33
def __init__(self, prob: float, hgain: float, sgain: float, vgain: float):
    """

    :param prob:            Probability to apply the transform.
    :param hgain:           Hue gain.
    :param sgain:           Saturation gain.
    :param vgain:           Value gain.
    """
    super().__init__()
    self.prob = prob
    self.hgain = hgain
    self.sgain = sgain
    self.vgain = vgain

KeypointsImageNormalize

Bases: AbstractKeypointTransform

Normalize image with mean and std using formula (image - mean) / std. Output image will allways have dtype of np.float32.

Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_image_normalize.py 
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
@register_transform(Transforms.KeypointsImageNormalize)
class KeypointsImageNormalize(AbstractKeypointTransform):
    """
    Normalize image with mean and std using formula `(image - mean) / std`.
    Output image will allways have dtype of np.float32.
    """

    def __init__(self, mean: Union[float, List[float], ListConfig], std: Union[float, List[float], ListConfig]):
        """

        :param mean: (float, List[float]) A constant bias to be subtracted from the image.
                     If it is a list, it should have the same length as the number of channels in the image.
        :param std:  (float, List[float]) A scaling factor to be applied to the image after subtracting mean.
                     If it is a list, it should have the same length as the number of channels in the image.
        """
        super().__init__()
        self.mean = np.array(list(mean)).reshape((1, 1, -1)).astype(np.float32)
        self.std = np.array(list(std)).reshape((1, 1, -1)).astype(np.float32)

    def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
        """
        Apply transformation to given pose estimation sample

        :param sample: A pose estimation sample
        :return:       Same pose estimation sample with normalized image
        """
        sample.image = np.divide(sample.image - self.mean, self.std, dtype=np.float32)
        return sample

    def __repr__(self):
        return self.__class__.__name__ + f"(mean={self.mean}, std={self.std})"

    def get_equivalent_preprocessing(self) -> List:
        return [{Processings.NormalizeImage: {"mean": self.mean, "std": self.std}}]

__init__(mean, std)

Parameters:

Name Type Description Default
mean Union[float, List[float], ListConfig]

(float, List[float]) A constant bias to be subtracted from the image. If it is a list, it should have the same length as the number of channels in the image.

 required
std Union[float, List[float], ListConfig]

(float, List[float]) A scaling factor to be applied to the image after subtracting mean. If it is a list, it should have the same length as the number of channels in the image.

 required
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_image_normalize.py 
20
21
22
23
24
25
26
27
28
29
30
def __init__(self, mean: Union[float, List[float], ListConfig], std: Union[float, List[float], ListConfig]):
    """

    :param mean: (float, List[float]) A constant bias to be subtracted from the image.
                 If it is a list, it should have the same length as the number of channels in the image.
    :param std:  (float, List[float]) A scaling factor to be applied to the image after subtracting mean.
                 If it is a list, it should have the same length as the number of channels in the image.
    """
    super().__init__()
    self.mean = np.array(list(mean)).reshape((1, 1, -1)).astype(np.float32)
    self.std = np.array(list(std)).reshape((1, 1, -1)).astype(np.float32)

apply_to_sample(sample)

Apply transformation to given pose estimation sample

Parameters:

Name Type Description Default
sample PoseEstimationSample

A pose estimation sample

 required

Returns:

Type Description
PoseEstimationSample

Same pose estimation sample with normalized image
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_image_normalize.py 
32
33
34
35
36
37
38
39
40
def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
    """
    Apply transformation to given pose estimation sample

    :param sample: A pose estimation sample
    :return:       Same pose estimation sample with normalized image
    """
    sample.image = np.divide(sample.image - self.mean, self.std, dtype=np.float32)
    return sample

KeypointsImageStandardize

Bases: AbstractKeypointTransform

Standardize image pixel values with img/max_value formula. Output image will allways have dtype of np.float32.

Parameters:

Name Type Description Default
max_value float

Current maximum value of the image pixels. (usually 255)

 255.0
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_image_standardize.py 
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
@register_transform(Transforms.KeypointsImageStandardize)
class KeypointsImageStandardize(AbstractKeypointTransform):
    """
    Standardize image pixel values with img/max_value formula.
    Output image will allways have dtype of np.float32.

    :param max_value: Current maximum value of the image pixels. (usually 255)
    """

    def __init__(self, max_value: float = 255.0):
        """

        :param max_value: A constant value to divide the image by.
        """
        super().__init__()
        self.max_value = max_value

    def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
        """
        Apply transformation to given pose estimation sample

        :param sample: A pose estimation sample
        :return:       Same pose estimation sample with standardized image
        """
        sample.image = np.divide(sample.image, self.max_value, dtype=np.float32)
        return sample

    def get_equivalent_preprocessing(self) -> List[Dict]:
        return [{Processings.StandardizeImage: {"max_value": self.max_value}}]

    def __repr__(self):
        return self.__class__.__name__ + f"(max_value={self.max_value})"

__init__(max_value=255.0)

Parameters:

Name Type Description Default
max_value float

A constant value to divide the image by.

 255.0
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_image_standardize.py 
20
21
22
23
24
25
26
def __init__(self, max_value: float = 255.0):
    """

    :param max_value: A constant value to divide the image by.
    """
    super().__init__()
    self.max_value = max_value

apply_to_sample(sample)

Apply transformation to given pose estimation sample

Parameters:

Name Type Description Default
sample PoseEstimationSample

A pose estimation sample

 required

Returns:

Type Description
PoseEstimationSample

Same pose estimation sample with standardized image
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_image_standardize.py 
28
29
30
31
32
33
34
35
36
def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
    """
    Apply transformation to given pose estimation sample

    :param sample: A pose estimation sample
    :return:       Same pose estimation sample with standardized image
    """
    sample.image = np.divide(sample.image, self.max_value, dtype=np.float32)
    return sample

KeypointsImageToTensor

Convert image from numpy array to tensor and permute axes to [C,H,W].

This transform works only for old-style transform API and will raise an exception when used in strongly-typed data samples transform API.

Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_image_to_tensor.py 
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
@register_transform(Transforms.KeypointsImageToTensor)
class KeypointsImageToTensor:
    """
    Convert image from numpy array to tensor and permute axes to [C,H,W].

    This transform works only for old-style transform API and will raise an exception when used in strongly-typed
    data samples transform API.
    """

    def __init__(self):
        self.additional_samples_count = 0

    def __call__(
        self, image: np.ndarray, mask: np.ndarray, joints: np.ndarray, areas: Optional[np.ndarray], bboxes: Optional[np.ndarray]
    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray], Optional[np.ndarray]]:
        """
        Convert image from numpy array to tensor and permute axes to [C,H,W].
        """
        image = torch.from_numpy(np.transpose(image, (2, 0, 1))).float()
        return image, mask, joints, areas, bboxes

    def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
        raise RuntimeError(
            f"{self.__class__} does not have apply_to_sample method because manual channel permutation HWC->CHW"
            f"is not needed for new data samples API. This is currently performed inside collate_fn."
        )

    def get_equivalent_preprocessing(self) -> List:
        return [
            {Processings.ImagePermute: {"permutation": (2, 0, 1)}},
        ]

    def __repr__(self):
        return self.__class__.__name__ + "()"

__call__(image, mask, joints, areas, bboxes)

Convert image from numpy array to tensor and permute axes to [C,H,W].

Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_image_to_tensor.py 
23
24
25
26
27
28
29
30
def __call__(
    self, image: np.ndarray, mask: np.ndarray, joints: np.ndarray, areas: Optional[np.ndarray], bboxes: Optional[np.ndarray]
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray], Optional[np.ndarray]]:
    """
    Convert image from numpy array to tensor and permute axes to [C,H,W].
    """
    image = torch.from_numpy(np.transpose(image, (2, 0, 1))).float()
    return image, mask, joints, areas, bboxes

KeypointsLongestMaxSize

Bases: AbstractKeypointTransform

Resize data sample to guarantee that input image dimensions is not exceeding maximum width & height

Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_longest_max_size.py 
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
@register_transform(Transforms.KeypointsLongestMaxSize)
class KeypointsLongestMaxSize(AbstractKeypointTransform):
    """
    Resize data sample to guarantee that input image dimensions is not exceeding maximum width & height
    """

    def __init__(self, max_height: int, max_width: int, interpolation: int = cv2.INTER_LINEAR, prob: float = 1.0):
        """

        :param max_height: (int) - Maximum image height
        :param max_width: (int) - Maximum image width
        :param interpolation: Used interpolation method for image
        :param prob: Probability of applying this transform
        """
        super().__init__()
        self.max_height = max_height
        self.max_width = max_width
        self.interpolation = interpolation
        self.prob = prob

    def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
        if random.random() < self.prob:
            height, width = sample.image.shape[:2]
            scale = min(self.max_height / height, self.max_width / width)
            sample.image = self.apply_to_image(sample.image, scale, cv2.INTER_LINEAR)
            sample.mask = self.apply_to_image(sample.mask, scale, cv2.INTER_NEAREST)

            if sample.image.shape[0] != self.max_height and sample.image.shape[1] != self.max_width:
                raise RuntimeError(f"Image shape is not as expected (scale={scale}, input_shape={height, width}, resized_shape={sample.image.shape[:2]})")

            if sample.image.shape[0] > self.max_height or sample.image.shape[1] > self.max_width:
                raise RuntimeError(f"Image shape is not as expected (scale={scale}, input_shape={height, width}, resized_shape={sample.image.shape[:2]}")

            sample.joints = self.apply_to_keypoints(sample.joints, scale)
            if sample.bboxes_xywh is not None:
                sample.bboxes_xywh = self.apply_to_bboxes(sample.bboxes_xywh, scale)

            if sample.areas is not None:
                sample.areas = np.multiply(sample.areas, scale**2, dtype=np.float32)

        return sample

    @classmethod
    def apply_to_image(cls, img, scale, interpolation):
        height, width = img.shape[:2]

        if scale != 1.0:
            new_height, new_width = tuple(int(dim * scale + 0.5) for dim in (height, width))
            img = cv2.resize(img, dsize=(new_width, new_height), interpolation=interpolation)
        return img

    @classmethod
    def apply_to_keypoints(cls, keypoints, scale):
        keypoints = keypoints.astype(np.float32, copy=True)
        keypoints[:, :, 0:2] *= scale
        return keypoints

    @classmethod
    def apply_to_bboxes(cls, bboxes, scale):
        return np.multiply(bboxes, scale, dtype=np.float32)

    def __repr__(self):
        return (
            self.__class__.__name__ + f"(max_height={self.max_height}, "
            f"max_width={self.max_width}, "
            f"interpolation={self.interpolation}, prob={self.prob})"
        )

    def get_equivalent_preprocessing(self) -> List:
        return [{Processings.KeypointsLongestMaxSizeRescale: {"output_shape": (self.max_height, self.max_width)}}]

__init__(max_height, max_width, interpolation=cv2.INTER_LINEAR, prob=1.0)

Parameters:

Name Type Description Default
max_height int

(int) - Maximum image height

 required
max_width int

(int) - Maximum image width

 required
interpolation int

Used interpolation method for image

 cv2.INTER_LINEAR
prob float

Probability of applying this transform

 1.0
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_longest_max_size.py 
19
20
21
22
23
24
25
26
27
28
29
30
31
def __init__(self, max_height: int, max_width: int, interpolation: int = cv2.INTER_LINEAR, prob: float = 1.0):
    """

    :param max_height: (int) - Maximum image height
    :param max_width: (int) - Maximum image width
    :param interpolation: Used interpolation method for image
    :param prob: Probability of applying this transform
    """
    super().__init__()
    self.max_height = max_height
    self.max_width = max_width
    self.interpolation = interpolation
    self.prob = prob

KeypointsMixup

Bases: AbstractKeypointTransform

Apply mixup augmentation and combine two samples into one. Images are averaged with equal weights. Targets are concatenated without any changes. This transform requires both samples have the same image size. The easiest way to achieve this is to use resize + padding before this transform:

# This will apply KeypointsLongestMaxSize and KeypointsPadIfNeeded to two samples individually
# and then apply KeypointsMixup to get a single sample.
train_dataset_params:
    transforms:
        - KeypointsLongestMaxSize:
            max_height: ${dataset_params.image_size}
            max_width: ${dataset_params.image_size}

        - KeypointsPadIfNeeded:
            min_height: ${dataset_params.image_size}
            min_width: ${dataset_params.image_size}
            image_pad_value: [127, 127, 127]
            mask_pad_value: 1
            padding_mode: center

        - KeypointsMixup:
            prob: 0.5

Parameters:

Name Type Description Default
prob float

Probability to apply the transform.

 required
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_mixup.py 
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
@register_transform()
class KeypointsMixup(AbstractKeypointTransform):
    """
    Apply mixup augmentation and combine two samples into one.
    Images are averaged with equal weights. Targets are concatenated without any changes.
    This transform requires both samples have the same image size. The easiest way to achieve this is to use resize + padding before this transform:

    ```yaml
    # This will apply KeypointsLongestMaxSize and KeypointsPadIfNeeded to two samples individually
    # and then apply KeypointsMixup to get a single sample.
    train_dataset_params:
        transforms:
            - KeypointsLongestMaxSize:
                max_height: ${dataset_params.image_size}
                max_width: ${dataset_params.image_size}

            - KeypointsPadIfNeeded:
                min_height: ${dataset_params.image_size}
                min_width: ${dataset_params.image_size}
                image_pad_value: [127, 127, 127]
                mask_pad_value: 1
                padding_mode: center

            - KeypointsMixup:
                prob: 0.5
    ```

    :param prob:            Probability to apply the transform.
    """

    def __init__(self, prob: float):
        """

        :param prob:            Probability to apply the transform.
        """
        super().__init__(additional_samples_count=1)
        self.prob = prob

    def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
        """
        Apply the transform to a single sample.

        :param sample: An input sample. It should have one additional sample in `additional_samples` field.
        :return:       A new pose estimation sample that represents the mixup sample.
        """
        if random.random() < self.prob:
            other = sample.additional_samples[0]
            if sample.image.shape != other.image.shape:
                raise RuntimeError(
                    f"KeypointsMixup requires both samples to have the same image shape. "
                    f"Got {sample.image.shape} and {other.image.shape}. "
                    f"Use KeypointsLongestMaxSize and KeypointsPadIfNeeded to resize and pad images before this transform."
                )
            sample = self._apply_mixup(sample, other)
        return sample

    def _apply_mixup(self, sample: PoseEstimationSample, other: PoseEstimationSample) -> PoseEstimationSample:
        """
        Apply mixup augmentation to a single sample.
        :param sample: First sample.
        :param other:  Second sample.
        :return:       Mixup sample.
        """
        image = (sample.image * 0.5 + other.image * 0.5).astype(sample.image.dtype)
        mask = np.logical_or(sample.mask, other.mask).astype(sample.mask.dtype)
        joints = np.concatenate([sample.joints, other.joints], axis=0)
        is_crowd = np.concatenate([sample.is_crowd, other.is_crowd], axis=0)

        bboxes = self._concatenate_arrays(sample.bboxes_xywh, other.bboxes_xywh, (0, 4))
        areas = self._concatenate_arrays(sample.areas, other.areas, (0,))
        return PoseEstimationSample(image=image, mask=mask, joints=joints, is_crowd=is_crowd, bboxes_xywh=bboxes, areas=areas, additional_samples=None)

    def _concatenate_arrays(self, arr1: Optional[np.ndarray], arr2: Optional[np.ndarray], shape_if_empty) -> Optional[np.ndarray]:
        """
        Concatenate two arrays. If one of the arrays is None, it will be replaced with array of zeros of given shape.
        This is purely utility function to simplify code of stacking arrays that may be None.
        Arrays must have same number of dims.

        :param arr1:           First array
        :param arr2:           Second array
        :param shape_if_empty: Shape of the array to create if one of the arrays is None.
        :return:               Stacked arrays along first axis. If both arrays are None, then None is returned.
        """
        if arr1 is None and arr2 is None:
            return None
        if arr1 is None:
            arr1 = np.zeros(shape_if_empty, dtype=np.float32)
        if arr2 is None:
            arr2 = np.zeros(shape_if_empty, dtype=np.float32)
        return np.concatenate([arr1, arr2], axis=0)

    def get_equivalent_preprocessing(self):
        raise RuntimeError(f"{self.__class__} does not have equivalent preprocessing because it is non-deterministic.")

__init__(prob)

Parameters:

Name Type Description Default
prob float

Probability to apply the transform.

 required
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_mixup.py 
42
43
44
45
46
47
48
def __init__(self, prob: float):
    """

    :param prob:            Probability to apply the transform.
    """
    super().__init__(additional_samples_count=1)
    self.prob = prob

apply_to_sample(sample)

Apply the transform to a single sample.

Parameters:

Name Type Description Default
sample PoseEstimationSample

An input sample. It should have one additional sample in additional_samples field.

 required

Returns:

Type Description
PoseEstimationSample

A new pose estimation sample that represents the mixup sample.
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_mixup.py 
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
    """
    Apply the transform to a single sample.

    :param sample: An input sample. It should have one additional sample in `additional_samples` field.
    :return:       A new pose estimation sample that represents the mixup sample.
    """
    if random.random() < self.prob:
        other = sample.additional_samples[0]
        if sample.image.shape != other.image.shape:
            raise RuntimeError(
                f"KeypointsMixup requires both samples to have the same image shape. "
                f"Got {sample.image.shape} and {other.image.shape}. "
                f"Use KeypointsLongestMaxSize and KeypointsPadIfNeeded to resize and pad images before this transform."
            )
        sample = self._apply_mixup(sample, other)
    return sample

KeypointsMosaic

Bases: AbstractKeypointTransform

Assemble 4 samples together to make 2x2 grid. This transform stacks input samples together to make a square with padding if necessary. This transform does not require input samples to have same size. If input samples have different sizes (H1,W1), (H2,W2), (H3,W3), (H4,W4), then resulting mosaic will have height of max(H1,H2) + max(H3,H4) and width of max(W1+W2, W2+W3), assuming the first sample is located in top left corner, second sample is in top right corner, third sample is in bottom left corner and fourth sample is in bottom right corner of mosaic.

The location of mosaic transform in the transforms list matter. It affects what transforms will be applied to all 4 samples.

In the example below, KeypointsMosaic goes after KeypointsRandomAffineTransform and KeypointsBrightnessContrast. This means that all 4 samples will be transformed with KeypointsRandomAffineTransform and KeypointsBrightnessContrast.

# This will apply KeypointsRandomAffineTransform and KeypointsBrightnessContrast to four sampls individually
# and then assemble them together in mosaic
train_dataset_params:
    transforms:
        - KeypointsRandomAffineTransform:
            min_scale: 0.75
            max_scale: 1.5

        - KeypointsBrightnessContrast:
            brightness_range: [ 0.8, 1.2 ]
            contrast_range: [ 0.8, 1.2 ]
            prob: 0.5

        - KeypointsMosaic:
            prob: 0.5

Contrary, if one puts KeypointsMosaic before KeypointsRandomAffineTransform and KeypointsBrightnessContrast, then 4 original samples will be assembled in mosaic and then transformed with KeypointsRandomAffineTransform and KeypointsBrightnessContrast:

# This will first assemble 4 samples in mosaic and then apply KeypointsRandomAffineTransform and KeypointsBrightnessContrast to the mosaic.
train_dataset_params:
    transforms:
        - KeypointsRandomAffineTransform:
            min_scale: 0.75
            max_scale: 1.5

        - KeypointsBrightnessContrast:
            brightness_range: [ 0.8, 1.2 ]
            contrast_range: [ 0.8, 1.2 ]
            prob: 0.5

        - KeypointsMosaic:
            prob: 0.5
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_mosaic.py 
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
@register_transform()
class KeypointsMosaic(AbstractKeypointTransform):
    """
    Assemble 4 samples together to make 2x2 grid.
    This transform stacks input samples together to make a square with padding if necessary.
    This transform does not require input samples to have same size.
    If input samples have different sizes (H1,W1), (H2,W2), (H3,W3), (H4,W4), then resulting mosaic will have
    height of max(H1,H2) + max(H3,H4) and width of max(W1+W2, W2+W3), assuming the first sample is located in top left corner,
    second sample is in top right corner, third sample is in bottom left corner and fourth sample is in bottom right corner of mosaic.

    The location of mosaic transform in the transforms list matter.
    It affects what transforms will be applied to all 4 samples.

    In the example below, KeypointsMosaic goes after KeypointsRandomAffineTransform and KeypointsBrightnessContrast.
    This means that all 4 samples will be transformed with KeypointsRandomAffineTransform and KeypointsBrightnessContrast.

    ```yaml
    # This will apply KeypointsRandomAffineTransform and KeypointsBrightnessContrast to four sampls individually
    # and then assemble them together in mosaic
    train_dataset_params:
        transforms:
            - KeypointsRandomAffineTransform:
                min_scale: 0.75
                max_scale: 1.5

            - KeypointsBrightnessContrast:
                brightness_range: [ 0.8, 1.2 ]
                contrast_range: [ 0.8, 1.2 ]
                prob: 0.5

            - KeypointsMosaic:
                prob: 0.5
    ```

    Contrary, if one puts KeypointsMosaic before KeypointsRandomAffineTransform and KeypointsBrightnessContrast,
    then 4 original samples will be assembled in mosaic and then transformed with KeypointsRandomAffineTransform and KeypointsBrightnessContrast:

    ```yaml
    # This will first assemble 4 samples in mosaic and then apply KeypointsRandomAffineTransform and KeypointsBrightnessContrast to the mosaic.
    train_dataset_params:
        transforms:
            - KeypointsRandomAffineTransform:
                min_scale: 0.75
                max_scale: 1.5

            - KeypointsBrightnessContrast:
                brightness_range: [ 0.8, 1.2 ]
                contrast_range: [ 0.8, 1.2 ]
                prob: 0.5

            - KeypointsMosaic:
                prob: 0.5
    ```

    """

    def __init__(self, prob: float, pad_value=(127, 127, 127)):
        """

        :param prob:     Probability to apply the transform.
        :param pad_value Value to pad the image if size of samples does not match.
        """
        super().__init__(additional_samples_count=3)
        self.prob = prob
        self.pad_value = tuple(pad_value)

    def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
        """
        Apply transformation to given estimation sample

        :param sample: A pose estimation sample. The sample must have 3 additional samples in it.
        :return:       A new pose estimation sample that represents the final mosaic.
        """
        if random.random() < self.prob:
            samples = [sample] + sample.additional_samples
            sample = self._apply_mosaic(samples)
        return sample

    def _apply_mosaic(self, samples: List[PoseEstimationSample]) -> PoseEstimationSample:
        """
        Actual method to apply mosaic to the sample.

        :param samples: List of 4 samples to make mosaic from.
        :return:        A new pose estimation sample that represents the final mosaic.
        """
        top_left, top_right, btm_left, btm_right = samples

        mosaic_sample = self._stack_samples_vertically(
            self._stack_samples_horizontally(top_left, top_right, pad_from_top=True), self._stack_samples_horizontally(btm_left, btm_right, pad_from_top=False)
        )

        return mosaic_sample

    def _pad_sample(self, sample: PoseEstimationSample, pad_top: int = 0, pad_left: int = 0, pad_right: int = 0, pad_bottom: int = 0) -> PoseEstimationSample:
        """
        Pad the sample with given padding values.

        :param sample:     Input sample. Sample is modified inplace.
        :param pad_top:    Padding in pixels from top.
        :param pad_left:   Padding in pixels from left.
        :param pad_right:  Padding in pixels from right.
        :param pad_bottom: Padding in pixels from bottom.
        :return:           Modified sample.
        """
        sample.image = cv2.copyMakeBorder(
            sample.image, top=pad_top, bottom=pad_bottom, left=pad_left, right=pad_right, borderType=cv2.BORDER_CONSTANT, value=self.pad_value
        )
        sample.mask = cv2.copyMakeBorder(sample.mask, top=pad_top, bottom=pad_bottom, left=pad_left, right=pad_right, borderType=cv2.BORDER_CONSTANT, value=1)

        sample.joints[:, :, 0] += pad_left
        sample.joints[:, :, 1] += pad_top

        sample.bboxes_xywh[:, 0] += pad_left
        sample.bboxes_xywh[:, 1] += pad_top

        return sample

    def _stack_samples_horizontally(self, left: PoseEstimationSample, right: PoseEstimationSample, pad_from_top: bool) -> PoseEstimationSample:
        """
        Stack two samples horizontally.

        :param left:         First sample (Will be located on the left side).
        :param right:        Second sample (Will be location on the right side).
        :param pad_from_top: Controls whether images should be padded from top or from bottom if they have different heights.
        :return:             A stacked sample. If first image has H1,W1 shape and second image has H2,W2 shape,
                             then resulting image will have max(H1,H2), W1+W2 shape.
        """

        max_height = max(left.image.shape[0], right.image.shape[0])
        if pad_from_top:
            left = self._pad_sample(left, pad_top=max_height - left.image.shape[0])
            right = self._pad_sample(right, pad_top=max_height - right.image.shape[0])
        else:
            left = self._pad_sample(left, pad_bottom=max_height - left.image.shape[0])
            right = self._pad_sample(right, pad_bottom=max_height - right.image.shape[0])

        image = np.concatenate([left.image, right.image], axis=1)
        mask = np.concatenate([left.mask, right.mask], axis=1)

        left_sample_width = left.image.shape[1]

        right_bboxes = right.bboxes_xywh
        if right_bboxes is None:
            right_bboxes = np.zeros((0, 4), dtype=np.float32)

        right_joints_offset = np.array([left_sample_width, 0, 0], dtype=right.joints.dtype).reshape((1, 1, 3))
        right_bboxes_offset = np.array([left_sample_width, 0, 0, 0], dtype=right_bboxes.dtype).reshape((1, 4))

        joints = np.concatenate([left.joints, right.joints + right_joints_offset], axis=0)
        bboxes = self._concatenate_arrays(left.bboxes_xywh, right_bboxes + right_bboxes_offset, shape_if_empty=(0, 4))

        is_crowd = np.concatenate([left.is_crowd, right.is_crowd], axis=0)
        areas = self._concatenate_arrays(left.areas, right.areas, shape_if_empty=(0,))
        return PoseEstimationSample(image=image, mask=mask, joints=joints, is_crowd=is_crowd, bboxes_xywh=bboxes, areas=areas, additional_samples=None)

    def _stack_samples_vertically(self, top: PoseEstimationSample, bottom: PoseEstimationSample) -> PoseEstimationSample:
        """
        Stack two samples vertically. If images have different widths, they will be padded to match the width
        of the widest image. In case padding occurs, it will be done from both sides to keep the images centered.

        :param top:    First sample (Will be located on the top).
        :param bottom: Second sample (Will be location on the bottom).
        :return:       A stacked sample. If first image has H1,W1 shape and second image has H2,W2 shape,
                       then resulting image will have H1+H2, max(W1,W2) shape.
        """
        max_width = max(top.image.shape[1], bottom.image.shape[1])

        pad_left = (max_width - top.image.shape[1]) // 2
        pad_right = max_width - top.image.shape[1] - pad_left
        top = self._pad_sample(top, pad_left=pad_left, pad_right=pad_right)

        pad_left = (max_width - bottom.image.shape[1]) // 2
        pad_right = max_width - bottom.image.shape[1] - pad_left
        bottom = self._pad_sample(bottom, pad_left=pad_left, pad_right=pad_right)

        image = np.concatenate([top.image, bottom.image], axis=0)
        mask = np.concatenate([top.mask, bottom.mask], axis=0)

        top_sample_height = top.image.shape[0]

        bottom_bboxes = bottom.bboxes_xywh
        if bottom_bboxes is None:
            bottom_bboxes = np.zeros((0, 4), dtype=np.float32)

        bottom_joints_offset = np.array([0, top_sample_height, 0], dtype=bottom.joints.dtype).reshape((1, 1, 3))
        bottom_bboxes_offset = np.array([0, top_sample_height, 0, 0], dtype=bottom_bboxes.dtype).reshape((1, 4))

        joints = np.concatenate([top.joints, bottom.joints + bottom_joints_offset], axis=0)
        bboxes = self._concatenate_arrays(top.bboxes_xywh, bottom_bboxes + bottom_bboxes_offset, shape_if_empty=(0, 4))

        is_crowd = np.concatenate([top.is_crowd, bottom.is_crowd], axis=0)
        areas = self._concatenate_arrays(top.areas, bottom.areas, shape_if_empty=(0,))
        return PoseEstimationSample(image=image, mask=mask, joints=joints, is_crowd=is_crowd, bboxes_xywh=bboxes, areas=areas, additional_samples=None)

    def _concatenate_arrays(self, arr1: Optional[np.ndarray], arr2: Optional[np.ndarray], shape_if_empty) -> Optional[np.ndarray]:
        """
        Concatenate two arrays. If one of the arrays is None, it will be replaced with array of zeros of given shape.
        This is purely utility function to simplify code of stacking arrays that may be None.
        Arrays must have same number of dims.

        :param arr1:           First array
        :param arr2:           Second array
        :param shape_if_empty: Shape of the array to create if one of the arrays is None.
        :return:               Stacked arrays along first axis. If both arrays are None, then None is returned.
        """
        if arr1 is None and arr2 is None:
            return None
        if arr1 is None:
            arr1 = np.zeros(shape_if_empty, dtype=np.float32)
        if arr2 is None:
            arr2 = np.zeros(shape_if_empty, dtype=np.float32)
        return np.concatenate([arr1, arr2], axis=0)

    def get_equivalent_preprocessing(self):
        raise RuntimeError(f"{self.__class__} does not have equivalent preprocessing because it is non-deterministic.")

__init__(prob, pad_value=(127, 127, 127))

Parameters:

Name Type Description Default
prob float

Probability to apply the transform.

 required
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_mosaic.py 
68
69
70
71
72
73
74
75
76
def __init__(self, prob: float, pad_value=(127, 127, 127)):
    """

    :param prob:     Probability to apply the transform.
    :param pad_value Value to pad the image if size of samples does not match.
    """
    super().__init__(additional_samples_count=3)
    self.prob = prob
    self.pad_value = tuple(pad_value)

apply_to_sample(sample)

Apply transformation to given estimation sample

Parameters:

Name Type Description Default
sample PoseEstimationSample

A pose estimation sample. The sample must have 3 additional samples in it.

 required

Returns:

Type Description
PoseEstimationSample

A new pose estimation sample that represents the final mosaic.
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_mosaic.py 
78
79
80
81
82
83
84
85
86
87
88
def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
    """
    Apply transformation to given estimation sample

    :param sample: A pose estimation sample. The sample must have 3 additional samples in it.
    :return:       A new pose estimation sample that represents the final mosaic.
    """
    if random.random() < self.prob:
        samples = [sample] + sample.additional_samples
        sample = self._apply_mosaic(samples)
    return sample

KeypointsPadIfNeeded

Bases: AbstractKeypointTransform

Pad image and mask to ensure that resulting image size is not less than output_size (rows, cols). Image and mask padded from right and bottom, thus joints remains unchanged.

Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_pad_if_needed.py 
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
@register_transform(Transforms.KeypointsPadIfNeeded)
class KeypointsPadIfNeeded(AbstractKeypointTransform):
    """
    Pad image and mask to ensure that resulting image size is not less than `output_size` (rows, cols).
    Image and mask padded from right and bottom, thus joints remains unchanged.
    """

    def __init__(self, min_height: int, min_width: int, image_pad_value: int, mask_pad_value: float, padding_mode: str = "bottom_right"):
        """

        :param output_size: Desired image size (rows, cols)
        :param image_pad_value: Padding value of image
        :param mask_pad_value: Padding value for mask
        """
        if padding_mode not in ("bottom_right", "center"):
            raise ValueError(f"Unknown padding mode: {padding_mode}. Supported modes: 'bottom_right', 'center'")
        super().__init__()
        self.min_height = min_height
        self.min_width = min_width
        self.image_pad_value = image_pad_value
        self.mask_pad_value = mask_pad_value
        self.padding_mode = padding_mode

    def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
        height, width = sample.image.shape[:2]
        original_dtype = sample.mask.dtype

        if self.padding_mode == "bottom_right":
            pad_left = 0
            pad_top = 0
            pad_bottom = max(0, self.min_height - height)
            pad_right = max(0, self.min_width - width)
        elif self.padding_mode == "center":
            pad_left = max(0, (self.min_width - width) // 2)
            pad_top = max(0, (self.min_height - height) // 2)
            pad_bottom = max(0, self.min_height - height - pad_top)
            pad_right = max(0, self.min_width - width - pad_left)
        else:
            raise RuntimeError(f"Unknown padding mode: {self.padding_mode}")

        image_pad_value = tuple(self.image_pad_value) if isinstance(self.image_pad_value, Iterable) else tuple([self.image_pad_value] * sample.image.shape[-1])
        sample.image = cv2.copyMakeBorder(
            sample.image, top=pad_top, bottom=pad_bottom, left=pad_left, right=pad_right, value=image_pad_value, borderType=cv2.BORDER_CONSTANT
        )

        sample.mask = cv2.copyMakeBorder(
            sample.mask.astype(np.uint8),
            top=pad_top,
            bottom=pad_bottom,
            left=pad_left,
            right=pad_right,
            value=self.mask_pad_value,
            borderType=cv2.BORDER_CONSTANT,
        ).astype(original_dtype)

        sample.joints = self.apply_to_keypoints(sample.joints, pad_left, pad_top)
        if sample.bboxes_xywh is not None:
            sample.bboxes_xywh = self.apply_to_bboxes(sample.bboxes_xywh, pad_left, pad_top)

        return sample

    def apply_to_bboxes(self, bboxes: np.ndarray, pad_left, pad_top):
        bboxes = bboxes.copy()
        bboxes[:, 0] += pad_left
        bboxes[:, 1] += pad_top
        return bboxes

    def apply_to_keypoints(self, keypoints: np.ndarray, pad_left, pad_top):
        keypoints = keypoints.copy()
        keypoints[:, :, 0] += pad_left
        keypoints[:, :, 1] += pad_top
        return keypoints

    def __repr__(self):
        return (
            self.__class__.__name__ + f"(min_height={self.min_height}, "
            f"min_width={self.min_width}, "
            f"image_pad_value={self.image_pad_value}, "
            f"mask_pad_value={self.mask_pad_value}, "
            f"padding_mode={self.padding_mode}, "
            f")"
        )

    def get_equivalent_preprocessing(self) -> List:
        if self.padding_mode == "bottom_right":
            return [{Processings.KeypointsBottomRightPadding: {"output_shape": (self.min_height, self.min_width), "pad_value": self.image_pad_value}}]
        else:
            raise RuntimeError(f"KeypointsPadIfNeeded with padding_mode={self.padding_mode} is not implemented.")

__init__(min_height, min_width, image_pad_value, mask_pad_value, padding_mode='bottom_right')

Parameters:

Name Type Description Default
output_size

Desired image size (rows, cols)

 required
image_pad_value int

Padding value of image

 required
mask_pad_value float

Padding value for mask

 required
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_pad_if_needed.py 
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
def __init__(self, min_height: int, min_width: int, image_pad_value: int, mask_pad_value: float, padding_mode: str = "bottom_right"):
    """

    :param output_size: Desired image size (rows, cols)
    :param image_pad_value: Padding value of image
    :param mask_pad_value: Padding value for mask
    """
    if padding_mode not in ("bottom_right", "center"):
        raise ValueError(f"Unknown padding mode: {padding_mode}. Supported modes: 'bottom_right', 'center'")
    super().__init__()
    self.min_height = min_height
    self.min_width = min_width
    self.image_pad_value = image_pad_value
    self.mask_pad_value = mask_pad_value
    self.padding_mode = padding_mode

KeypointsRandomAffineTransform

Bases: AbstractKeypointTransform

Apply random affine transform to image, mask and joints.

Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_random_affine.py 
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
@register_transform(Transforms.KeypointsRandomAffineTransform)
class KeypointsRandomAffineTransform(AbstractKeypointTransform):
    """
    Apply random affine transform to image, mask and joints.
    """

    def __init__(
        self,
        max_rotation: float,
        min_scale: float,
        max_scale: float,
        max_translate: float,
        image_pad_value: Union[int, float, List[int]],
        mask_pad_value: float,
        interpolation_mode: Union[int, List[int]] = cv2.INTER_LINEAR,
        prob: float = 0.5,
    ):
        """

        :param max_rotation:       Max rotation angle in degrees
        :param min_scale:          Lower bound for the scale change. For +- 20% size jitter this should be 0.8
        :param max_scale:          Lower bound for the scale change. For +- 20% size jitter this should be 1.2
        :param max_translate:      Max translation offset in percents of image size
        :param image_pad_value:    Value to pad the image during affine transform. Can be single scalar or list.
                                   If a list is provided, it should have the same length as the number of channels in the image.
        :param mask_pad_value:     Value to pad the mask during affine transform.
        :param interpolation_mode: A constant integer or list of integers, specifying the interpolation mode to use.
                                   Possible values for interpolation_mode:
                                     cv2.INTER_NEAREST = 0,
                                     cv2.INTER_LINEAR = 1,
                                     cv2.INTER_CUBIC = 2,
                                     cv2.INTER_AREA = 3,
                                     cv2.INTER_LANCZOS4 = 4
                                   To use random interpolation modes on each call, set interpolation_mode = (0,1,2,3,4)
        :param prob:               Probability to apply the transform.
        """
        super().__init__()

        self.max_rotation = max_rotation
        self.min_scale = min_scale
        self.max_scale = max_scale
        self.max_translate = max_translate
        self.image_pad_value = image_pad_value
        self.mask_pad_value = mask_pad_value
        self.prob = prob
        self.interpolation_mode = tuple(interpolation_mode) if isinstance(interpolation_mode, Iterable) else (interpolation_mode,)

    def __repr__(self):
        return (
            self.__class__.__name__ + f"(max_rotation={self.max_rotation}, "
            f"min_scale={self.min_scale}, "
            f"max_scale={self.max_scale}, "
            f"max_translate={self.max_translate}, "
            f"image_pad_value={self.image_pad_value}, "
            f"mask_pad_value={self.mask_pad_value}, "
            f"prob={self.prob})"
        )

    def _get_affine_matrix(self, img: np.ndarray, angle: float, scale: float, dx: float, dy: float) -> np.ndarray:
        """
        Compute the affine matrix that combines rotation of image around center, scaling and translation
        according to given parameters. Order of operations is: scale, rotate, translate.

        :param angle: Rotation angle in degrees
        :param scale: Scaling factor
        :param dx:    Translation in x direction
        :param dy:    Translation in y direction
        :return:      Affine matrix [2,3]
        """
        height, width = img.shape[:2]
        center = (width / 2 + dx * width, height / 2 + dy * height)
        matrix = cv2.getRotationMatrix2D(center, angle, scale)

        return matrix

    def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
        """
        Apply transformation to given pose estimation sample.
        Since this transformation apply affine transform some keypoints/bboxes may be moved outside the image.
        After applying the transform, visibility status of joints is updated to reflect the new position of joints.
        Bounding boxes are clipped to image borders.
        If sample contains areas, they are scaled according to the applied affine transform.

        :param sample: A pose estimation sample
        :return:       A transformed pose estimation sample
        """

        if random.random() < self.prob:
            angle = random.uniform(-self.max_rotation, self.max_rotation)
            scale = random.uniform(self.min_scale, self.max_scale)
            dx = random.uniform(-self.max_translate, self.max_translate)
            dy = random.uniform(-self.max_translate, self.max_translate)
            interpolation = random.choice(self.interpolation_mode)

            mat_output = self._get_affine_matrix(sample.image, angle, scale, dx, dy)
            mat_output = mat_output[:2]

            image_pad_value = (
                tuple(self.image_pad_value) if isinstance(self.image_pad_value, Iterable) else tuple([self.image_pad_value] * sample.image.shape[-1])
            )

            sample.image = self.apply_to_image(
                sample.image, mat_output, interpolation=interpolation, padding_value=image_pad_value, padding_mode=cv2.BORDER_CONSTANT
            )
            sample.mask = self.apply_to_image(
                sample.mask, mat_output, interpolation=cv2.INTER_NEAREST, padding_value=self.mask_pad_value, padding_mode=cv2.BORDER_CONSTANT
            )

            sample.joints = self.apply_to_keypoints(sample.joints, mat_output, sample.image.shape[:2])

            if sample.bboxes_xywh is not None:
                sample.bboxes_xywh = self.apply_to_bboxes(sample.bboxes_xywh, mat_output)

            if sample.areas is not None:
                sample.areas = self.apply_to_areas(sample.areas, mat_output)

            sample = sample.sanitize_sample()

        return sample

    @classmethod
    def apply_to_areas(cls, areas: np.ndarray, mat: np.ndarray) -> np.ndarray:
        """
        Apply affine transform to areas.

        :param areas: [N] Single-dimension array of areas
        :param mat:   [2,3] Affine transformation matrix
        :return:      [N] Single-dimension array of areas
        """
        det = np.linalg.det(mat[:2, :2])
        return (areas * abs(det)).astype(areas.dtype)

    @classmethod
    def apply_to_bboxes(cls, bboxes_xywh: np.ndarray, mat: np.ndarray) -> np.ndarray:
        """

        :param bboxes: (N,4) array of bboxes in XYWH format
        :param mat:    [2,3] Affine transformation matrix
        :return:       (N,4) array of bboxes in XYWH format
        """

        def bbox_shift_scale_rotate(bbox, m):
            x_min, y_min, x_max, y_max = bbox[:4]

            x = np.array([x_min, x_max, x_max, x_min])
            y = np.array([y_min, y_min, y_max, y_max])
            ones = np.ones(shape=(len(x)))
            points_ones = np.vstack([x, y, ones]).transpose()

            tr_points = m.dot(points_ones.T).T

            x_min, x_max = min(tr_points[:, 0]), max(tr_points[:, 0])
            y_min, y_max = min(tr_points[:, 1]), max(tr_points[:, 1])

            return np.array([x_min, y_min, x_max, y_max])

        if len(bboxes_xywh) == 0:
            return bboxes_xywh
        bboxes_xyxy = xywh_to_xyxy(bboxes_xywh, image_shape=None)
        bboxes_xyxy = np.array([bbox_shift_scale_rotate(box, mat) for box in bboxes_xyxy])
        return xyxy_to_xywh(bboxes_xyxy, image_shape=None).astype(bboxes_xywh.dtype)

    @classmethod
    def apply_to_keypoints(cls, keypoints: np.ndarray, mat: np.ndarray, image_shape: Tuple[int, int]) -> np.ndarray:
        """
        Apply affine transform to keypoints.

        :param keypoints:   [N,K,3] array of keypoints in (x,y,visibility) format
        :param mat:         [2,3] Affine transformation matrix
        :param image_shape: Image shape after applying affine transform (height, width).
                            Used to update visibility status of keypoints.
        :return:            [N,K,3] array of keypoints in (x,y,visibility) format
        """
        keypoints_with_visibility = keypoints.copy()
        keypoints = keypoints_with_visibility[:, :, 0:2]

        shape = keypoints.shape
        dtype = keypoints.dtype
        keypoints = keypoints.reshape(-1, 2)
        keypoints = np.dot(np.concatenate((keypoints, keypoints[:, 0:1] * 0 + 1), axis=1), mat.T).reshape(shape)

        # Update visibility status of joints that were moved outside visible area
        image_height, image_width = image_shape[:2]
        outside_left = keypoints[:, :, 0] < 0
        outside_top = keypoints[:, :, 1] < 0
        outside_right = keypoints[:, :, 0] >= image_width
        outside_bottom = keypoints[:, :, 1] >= image_height

        joints_outside_image = outside_left | outside_top | outside_right | outside_bottom

        keypoints_with_visibility[:, :, 0:2] = keypoints
        keypoints_with_visibility[joints_outside_image, 2] = 0
        return keypoints_with_visibility.astype(dtype, copy=False)

    @classmethod
    def apply_to_image(cls, image: np.ndarray, mat: np.ndarray, interpolation: int, padding_value: Union[int, float, Tuple], padding_mode: int) -> np.ndarray:
        """
        Apply affine transform to image.

        :param image:          Input image
        :param mat:            [2,3] Affine transformation matrix
        :param interpolation:  Interpolation mode. See cv2.warpAffine for details.
        :param padding_value:  Value to pad the image during affine transform. See cv2.warpAffine for details.
        :param padding_mode:   Padding mode. See cv2.warpAffine for details.
        :return:               Transformed image of the same shape as input image.
        """
        return cv2.warpAffine(
            image,
            mat,
            dsize=(image.shape[1], image.shape[0]),
            flags=interpolation,
            borderValue=padding_value,
            borderMode=padding_mode,
        )

    def get_equivalent_preprocessing(self):
        raise RuntimeError(f"{self.__class__} does not have equivalent preprocessing because it is non-deterministic.")

__init__(max_rotation, min_scale, max_scale, max_translate, image_pad_value, mask_pad_value, interpolation_mode=cv2.INTER_LINEAR, prob=0.5)

Parameters:

Name Type Description Default
max_rotation float

Max rotation angle in degrees

 required
min_scale float

Lower bound for the scale change. For +- 20% size jitter this should be 0.8

 required
max_scale float

Lower bound for the scale change. For +- 20% size jitter this should be 1.2

 required
max_translate float

Max translation offset in percents of image size

 required
image_pad_value Union[int, float, List[int]]

Value to pad the image during affine transform. Can be single scalar or list. If a list is provided, it should have the same length as the number of channels in the image.

 required
mask_pad_value float

Value to pad the mask during affine transform.

 required
interpolation_mode Union[int, List[int]]

A constant integer or list of integers, specifying the interpolation mode to use. Possible values for interpolation_mode: cv2.INTER_NEAREST = 0, cv2.INTER_LINEAR = 1, cv2.INTER_CUBIC = 2, cv2.INTER_AREA = 3, cv2.INTER_LANCZOS4 = 4 To use random interpolation modes on each call, set interpolation_mode = (0,1,2,3,4)

 cv2.INTER_LINEAR
prob float

Probability to apply the transform.

 0.5
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_random_affine.py 
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
def __init__(
    self,
    max_rotation: float,
    min_scale: float,
    max_scale: float,
    max_translate: float,
    image_pad_value: Union[int, float, List[int]],
    mask_pad_value: float,
    interpolation_mode: Union[int, List[int]] = cv2.INTER_LINEAR,
    prob: float = 0.5,
):
    """

    :param max_rotation:       Max rotation angle in degrees
    :param min_scale:          Lower bound for the scale change. For +- 20% size jitter this should be 0.8
    :param max_scale:          Lower bound for the scale change. For +- 20% size jitter this should be 1.2
    :param max_translate:      Max translation offset in percents of image size
    :param image_pad_value:    Value to pad the image during affine transform. Can be single scalar or list.
                               If a list is provided, it should have the same length as the number of channels in the image.
    :param mask_pad_value:     Value to pad the mask during affine transform.
    :param interpolation_mode: A constant integer or list of integers, specifying the interpolation mode to use.
                               Possible values for interpolation_mode:
                                 cv2.INTER_NEAREST = 0,
                                 cv2.INTER_LINEAR = 1,
                                 cv2.INTER_CUBIC = 2,
                                 cv2.INTER_AREA = 3,
                                 cv2.INTER_LANCZOS4 = 4
                               To use random interpolation modes on each call, set interpolation_mode = (0,1,2,3,4)
    :param prob:               Probability to apply the transform.
    """
    super().__init__()

    self.max_rotation = max_rotation
    self.min_scale = min_scale
    self.max_scale = max_scale
    self.max_translate = max_translate
    self.image_pad_value = image_pad_value
    self.mask_pad_value = mask_pad_value
    self.prob = prob
    self.interpolation_mode = tuple(interpolation_mode) if isinstance(interpolation_mode, Iterable) else (interpolation_mode,)

apply_to_areas(areas, mat) classmethod

Apply affine transform to areas.

Parameters:

Name Type Description Default
areas np.ndarray

[N] Single-dimension array of areas

 required
mat np.ndarray

[2,3] Affine transformation matrix

 required

Returns:

Type Description
np.ndarray

[N] Single-dimension array of areas
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_random_affine.py 
134
135
136
137
138
139
140
141
142
143
144
@classmethod
def apply_to_areas(cls, areas: np.ndarray, mat: np.ndarray) -> np.ndarray:
    """
    Apply affine transform to areas.

    :param areas: [N] Single-dimension array of areas
    :param mat:   [2,3] Affine transformation matrix
    :return:      [N] Single-dimension array of areas
    """
    det = np.linalg.det(mat[:2, :2])
    return (areas * abs(det)).astype(areas.dtype)

apply_to_bboxes(bboxes_xywh, mat) classmethod

Parameters:

Name Type Description Default
bboxes

(N,4) array of bboxes in XYWH format

 required
mat np.ndarray

[2,3] Affine transformation matrix

 required

Returns:

Type Description
np.ndarray

(N,4) array of bboxes in XYWH format
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_random_affine.py 
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
@classmethod
def apply_to_bboxes(cls, bboxes_xywh: np.ndarray, mat: np.ndarray) -> np.ndarray:
    """

    :param bboxes: (N,4) array of bboxes in XYWH format
    :param mat:    [2,3] Affine transformation matrix
    :return:       (N,4) array of bboxes in XYWH format
    """

    def bbox_shift_scale_rotate(bbox, m):
        x_min, y_min, x_max, y_max = bbox[:4]

        x = np.array([x_min, x_max, x_max, x_min])
        y = np.array([y_min, y_min, y_max, y_max])
        ones = np.ones(shape=(len(x)))
        points_ones = np.vstack([x, y, ones]).transpose()

        tr_points = m.dot(points_ones.T).T

        x_min, x_max = min(tr_points[:, 0]), max(tr_points[:, 0])
        y_min, y_max = min(tr_points[:, 1]), max(tr_points[:, 1])

        return np.array([x_min, y_min, x_max, y_max])

    if len(bboxes_xywh) == 0:
        return bboxes_xywh
    bboxes_xyxy = xywh_to_xyxy(bboxes_xywh, image_shape=None)
    bboxes_xyxy = np.array([bbox_shift_scale_rotate(box, mat) for box in bboxes_xyxy])
    return xyxy_to_xywh(bboxes_xyxy, image_shape=None).astype(bboxes_xywh.dtype)

apply_to_image(image, mat, interpolation, padding_value, padding_mode) classmethod

Apply affine transform to image.

Parameters:

Name Type Description Default
image np.ndarray

Input image

 required
mat np.ndarray

[2,3] Affine transformation matrix

 required
interpolation int

Interpolation mode. See cv2.warpAffine for details.

 required
padding_value Union[int, float, Tuple]

Value to pad the image during affine transform. See cv2.warpAffine for details.

 required
padding_mode int

Padding mode. See cv2.warpAffine for details.

 required

Returns:

Type Description
np.ndarray

Transformed image of the same shape as input image.
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_random_affine.py 
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
@classmethod
def apply_to_image(cls, image: np.ndarray, mat: np.ndarray, interpolation: int, padding_value: Union[int, float, Tuple], padding_mode: int) -> np.ndarray:
    """
    Apply affine transform to image.

    :param image:          Input image
    :param mat:            [2,3] Affine transformation matrix
    :param interpolation:  Interpolation mode. See cv2.warpAffine for details.
    :param padding_value:  Value to pad the image during affine transform. See cv2.warpAffine for details.
    :param padding_mode:   Padding mode. See cv2.warpAffine for details.
    :return:               Transformed image of the same shape as input image.
    """
    return cv2.warpAffine(
        image,
        mat,
        dsize=(image.shape[1], image.shape[0]),
        flags=interpolation,
        borderValue=padding_value,
        borderMode=padding_mode,
    )

apply_to_keypoints(keypoints, mat, image_shape) classmethod

Apply affine transform to keypoints.

Parameters:

Name Type Description Default
keypoints np.ndarray

[N,K,3] array of keypoints in (x,y,visibility) format

 required
mat np.ndarray

[2,3] Affine transformation matrix

 required
image_shape Tuple[int, int]

Image shape after applying affine transform (height, width). Used to update visibility status of keypoints.

 required

Returns:

Type Description
np.ndarray

[N,K,3] array of keypoints in (x,y,visibility) format
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_random_affine.py 
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
@classmethod
def apply_to_keypoints(cls, keypoints: np.ndarray, mat: np.ndarray, image_shape: Tuple[int, int]) -> np.ndarray:
    """
    Apply affine transform to keypoints.

    :param keypoints:   [N,K,3] array of keypoints in (x,y,visibility) format
    :param mat:         [2,3] Affine transformation matrix
    :param image_shape: Image shape after applying affine transform (height, width).
                        Used to update visibility status of keypoints.
    :return:            [N,K,3] array of keypoints in (x,y,visibility) format
    """
    keypoints_with_visibility = keypoints.copy()
    keypoints = keypoints_with_visibility[:, :, 0:2]

    shape = keypoints.shape
    dtype = keypoints.dtype
    keypoints = keypoints.reshape(-1, 2)
    keypoints = np.dot(np.concatenate((keypoints, keypoints[:, 0:1] * 0 + 1), axis=1), mat.T).reshape(shape)

    # Update visibility status of joints that were moved outside visible area
    image_height, image_width = image_shape[:2]
    outside_left = keypoints[:, :, 0] < 0
    outside_top = keypoints[:, :, 1] < 0
    outside_right = keypoints[:, :, 0] >= image_width
    outside_bottom = keypoints[:, :, 1] >= image_height

    joints_outside_image = outside_left | outside_top | outside_right | outside_bottom

    keypoints_with_visibility[:, :, 0:2] = keypoints
    keypoints_with_visibility[joints_outside_image, 2] = 0
    return keypoints_with_visibility.astype(dtype, copy=False)

apply_to_sample(sample)

Apply transformation to given pose estimation sample. Since this transformation apply affine transform some keypoints/bboxes may be moved outside the image. After applying the transform, visibility status of joints is updated to reflect the new position of joints. Bounding boxes are clipped to image borders. If sample contains areas, they are scaled according to the applied affine transform.

Parameters:

Name Type Description Default
sample PoseEstimationSample

A pose estimation sample

 required

Returns:

Type Description
PoseEstimationSample

A transformed pose estimation sample
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_random_affine.py 
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
    """
    Apply transformation to given pose estimation sample.
    Since this transformation apply affine transform some keypoints/bboxes may be moved outside the image.
    After applying the transform, visibility status of joints is updated to reflect the new position of joints.
    Bounding boxes are clipped to image borders.
    If sample contains areas, they are scaled according to the applied affine transform.

    :param sample: A pose estimation sample
    :return:       A transformed pose estimation sample
    """

    if random.random() < self.prob:
        angle = random.uniform(-self.max_rotation, self.max_rotation)
        scale = random.uniform(self.min_scale, self.max_scale)
        dx = random.uniform(-self.max_translate, self.max_translate)
        dy = random.uniform(-self.max_translate, self.max_translate)
        interpolation = random.choice(self.interpolation_mode)

        mat_output = self._get_affine_matrix(sample.image, angle, scale, dx, dy)
        mat_output = mat_output[:2]

        image_pad_value = (
            tuple(self.image_pad_value) if isinstance(self.image_pad_value, Iterable) else tuple([self.image_pad_value] * sample.image.shape[-1])
        )

        sample.image = self.apply_to_image(
            sample.image, mat_output, interpolation=interpolation, padding_value=image_pad_value, padding_mode=cv2.BORDER_CONSTANT
        )
        sample.mask = self.apply_to_image(
            sample.mask, mat_output, interpolation=cv2.INTER_NEAREST, padding_value=self.mask_pad_value, padding_mode=cv2.BORDER_CONSTANT
        )

        sample.joints = self.apply_to_keypoints(sample.joints, mat_output, sample.image.shape[:2])

        if sample.bboxes_xywh is not None:
            sample.bboxes_xywh = self.apply_to_bboxes(sample.bboxes_xywh, mat_output)

        if sample.areas is not None:
            sample.areas = self.apply_to_areas(sample.areas, mat_output)

        sample = sample.sanitize_sample()

    return sample

KeypointsRandomHorizontalFlip

Bases: AbstractKeypointTransform

Flip image, mask and joints horizontally with a given probability.

Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_random_horisontal_flip.py 
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
@register_transform(Transforms.KeypointsRandomHorizontalFlip)
class KeypointsRandomHorizontalFlip(AbstractKeypointTransform):
    """
    Flip image, mask and joints horizontally with a given probability.
    """

    def __init__(self, flip_index: List[int], prob: float = 0.5):
        """

        :param flip_index: Indexes of keypoints on the flipped image. When doing left-right flip, left hand becomes right hand.
                           So this array contains order of keypoints on the flipped image. This is dataset specific and depends on
                           how keypoints are defined in dataset.
        :param prob: Probability of flipping
        """
        super().__init__()
        self.flip_index = flip_index
        self.prob = prob

    def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
        """
        Apply transformation to given pose estimation sample

        :param sample: Input pose estimation sample.
        :return:       A new pose estimation sample.
        """
        if sample.image.shape[:2] != sample.mask.shape[:2]:
            raise RuntimeError(f"Image shape ({sample.image.shape[:2]}) does not match mask shape ({sample.mask.shape[:2]}).")

        if random.random() < self.prob:
            sample.image = self.apply_to_image(sample.image)
            sample.mask = self.apply_to_image(sample.mask)
            rows, cols = sample.image.shape[:2]
            sample.joints = self.apply_to_keypoints(sample.joints, cols)

            if sample.bboxes_xywh is not None:
                sample.bboxes_xywh = self.apply_to_bboxes(sample.bboxes_xywh, cols)

        return sample

    def apply_to_image(self, image: np.ndarray) -> np.ndarray:
        """
        Flip image horizontally

        :param image: Input image
        :return:      Horizontally flipped image
        """
        return np.ascontiguousarray(np.fliplr(image))

    def apply_to_keypoints(self, keypoints: np.ndarray, cols: int) -> np.ndarray:
        """
        Flip keypoints horizontally

        :param keypoints: Input keypoints of [N,K,3] shape
        :param cols:      Image width
        :return:          Flipped keypoints  of [N,K,3] shape
        """
        keypoints = keypoints.copy()
        keypoints = keypoints[:, self.flip_index]
        keypoints[:, :, 0] = cols - keypoints[:, :, 0] - 1
        return keypoints

    def apply_to_bboxes(self, bboxes: np.ndarray, cols: int) -> np.ndarray:
        """
        Flip boxes horizontally

        :param bboxes: Input boxes of [N,4] shape in XYWH format
        :param cols:   Image width
        :return:       Flipped boxes of [N,4] shape in XYWH format
        """

        bboxes = bboxes.copy()
        bboxes[:, 0] = cols - (bboxes[:, 0] + bboxes[:, 2])
        return bboxes

    def __repr__(self):
        return self.__class__.__name__ + f"(flip_index={self.flip_index}, prob={self.prob})"

    def get_equivalent_preprocessing(self):
        raise RuntimeError(f"{self.__class__} does not have equivalent preprocessing because it is non-deterministic.")

__init__(flip_index, prob=0.5)

Parameters:

Name Type Description Default
flip_index List[int]

Indexes of keypoints on the flipped image. When doing left-right flip, left hand becomes right hand. So this array contains order of keypoints on the flipped image. This is dataset specific and depends on how keypoints are defined in dataset.

 required
prob float

Probability of flipping

 0.5
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_random_horisontal_flip.py 
17
18
19
20
21
22
23
24
25
26
27
def __init__(self, flip_index: List[int], prob: float = 0.5):
    """

    :param flip_index: Indexes of keypoints on the flipped image. When doing left-right flip, left hand becomes right hand.
                       So this array contains order of keypoints on the flipped image. This is dataset specific and depends on
                       how keypoints are defined in dataset.
    :param prob: Probability of flipping
    """
    super().__init__()
    self.flip_index = flip_index
    self.prob = prob

apply_to_bboxes(bboxes, cols)

Flip boxes horizontally

Parameters:

Name Type Description Default
bboxes np.ndarray

Input boxes of [N,4] shape in XYWH format

 required
cols int

Image width

 required

Returns:

Type Description
np.ndarray

Flipped boxes of [N,4] shape in XYWH format
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_random_horisontal_flip.py 
72
73
74
75
76
77
78
79
80
81
82
83
def apply_to_bboxes(self, bboxes: np.ndarray, cols: int) -> np.ndarray:
    """
    Flip boxes horizontally

    :param bboxes: Input boxes of [N,4] shape in XYWH format
    :param cols:   Image width
    :return:       Flipped boxes of [N,4] shape in XYWH format
    """

    bboxes = bboxes.copy()
    bboxes[:, 0] = cols - (bboxes[:, 0] + bboxes[:, 2])
    return bboxes

apply_to_image(image)

Flip image horizontally

Parameters:

Name Type Description Default
image np.ndarray

Input image

 required

Returns:

Type Description
np.ndarray

Horizontally flipped image
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_random_horisontal_flip.py 
50
51
52
53
54
55
56
57
def apply_to_image(self, image: np.ndarray) -> np.ndarray:
    """
    Flip image horizontally

    :param image: Input image
    :return:      Horizontally flipped image
    """
    return np.ascontiguousarray(np.fliplr(image))

apply_to_keypoints(keypoints, cols)

Flip keypoints horizontally

Parameters:

Name Type Description Default
keypoints np.ndarray

Input keypoints of [N,K,3] shape

 required
cols int

Image width

 required

Returns:

Type Description
np.ndarray

Flipped keypoints of [N,K,3] shape
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_random_horisontal_flip.py 
59
60
61
62
63
64
65
66
67
68
69
70
def apply_to_keypoints(self, keypoints: np.ndarray, cols: int) -> np.ndarray:
    """
    Flip keypoints horizontally

    :param keypoints: Input keypoints of [N,K,3] shape
    :param cols:      Image width
    :return:          Flipped keypoints  of [N,K,3] shape
    """
    keypoints = keypoints.copy()
    keypoints = keypoints[:, self.flip_index]
    keypoints[:, :, 0] = cols - keypoints[:, :, 0] - 1
    return keypoints

apply_to_sample(sample)

Apply transformation to given pose estimation sample

Parameters:

Name Type Description Default
sample PoseEstimationSample

Input pose estimation sample.

 required

Returns:

Type Description
PoseEstimationSample

A new pose estimation sample.
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_random_horisontal_flip.py 
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
    """
    Apply transformation to given pose estimation sample

    :param sample: Input pose estimation sample.
    :return:       A new pose estimation sample.
    """
    if sample.image.shape[:2] != sample.mask.shape[:2]:
        raise RuntimeError(f"Image shape ({sample.image.shape[:2]}) does not match mask shape ({sample.mask.shape[:2]}).")

    if random.random() < self.prob:
        sample.image = self.apply_to_image(sample.image)
        sample.mask = self.apply_to_image(sample.mask)
        rows, cols = sample.image.shape[:2]
        sample.joints = self.apply_to_keypoints(sample.joints, cols)

        if sample.bboxes_xywh is not None:
            sample.bboxes_xywh = self.apply_to_bboxes(sample.bboxes_xywh, cols)

    return sample

KeypointsRandomRotate90

Bases: AbstractKeypointTransform

Apply 90 degree rotations to the sample.

Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_random_rotate90.py 
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
@register_transform(Transforms.KeypointsRandomRotate90)
class KeypointsRandomRotate90(AbstractKeypointTransform):
    """
    Apply 90 degree rotations to the sample.
    """

    def __init__(
        self,
        prob: float = 0.5,
    ):
        """
        Initialize transform

        :param prob (float): Probability of applying the transform
        """
        super().__init__()
        self.prob = prob

    def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
        """
        :param   sample: Input PoseEstimationSample
        :return:         Result of applying the transform
        """

        if random.random() < self.prob:
            factor = random.randint(0, 3)

            image_rows, image_cols = sample.image.shape[:2]

            sample.image = self.apply_to_image(sample.image, factor)
            sample.mask = self.apply_to_image(sample.mask, factor)
            sample.joints = self.apply_to_keypoints(sample.joints, factor, image_rows, image_cols)

            if sample.bboxes_xywh is not None:
                sample.bboxes_xywh = self.apply_to_bboxes(sample.bboxes_xywh, factor, image_rows, image_cols)

        return sample

    @classmethod
    def apply_to_image(cls, image: np.ndarray, factor: int) -> np.ndarray:
        """
        Rotate image by 90 degrees

        :param image:  Input image
        :param factor: Number of 90 degree rotations to apply. Order or rotation matches np.rot90
        :return:       Rotated image
        """
        return np.rot90(image, factor)

    @classmethod
    def apply_to_bboxes(cls, bboxes_xywh: np.ndarray, factor, rows: int, cols: int) -> np.ndarray:
        """

        :param bboxes: (N, 4) array of bboxes in XYWH format
        :param factor: Number of 90 degree rotations to apply. Order or rotation matches np.rot90
        :param rows:   Number of rows (image height) of the original (input) image
        :param cols:   Number of cols (image width) of the original (input) image
        :return:       Transformed bboxes in XYWH format
        """
        from super_gradients.training.transforms.transforms import DetectionRandomRotate90

        bboxes_xyxy = xywh_to_xyxy(bboxes_xywh, image_shape=None)
        bboxes_xyxy = DetectionRandomRotate90.xyxy_bbox_rot90(bboxes_xyxy, factor, rows, cols)
        return xyxy_to_xywh(bboxes_xyxy, image_shape=None)

    @classmethod
    def apply_to_keypoints(cls, keypoints: np.ndarray, factor, rows: int, cols: int) -> np.ndarray:
        """

        :param keypoints: Input keypoints array of [Num Instances, Num Joints, 3] shape.
                          Keypoints has format (x, y, visibility)
        :param factor:    Number of 90 degree rotations to apply. Order or rotation matches np.rot90
        :param rows:      Number of rows (image height) of the original (input) image
        :param cols:      Number of cols (image width) of the original (input) image
        :return:          Transformed keypoints array of [Num Instances, Num Joints, 3] shape.
        """
        x, y, v = keypoints[:, :, 0], keypoints[:, :, 1], keypoints[:, :, 2]

        if factor == 0:
            keypoints = x, y, v
        elif factor == 1:
            keypoints = y, cols - x - 1, v
        elif factor == 2:
            keypoints = cols - x - 1, rows - y - 1, v
        elif factor == 3:
            keypoints = rows - y - 1, x, v
        else:
            raise ValueError("Parameter n must be in set {0, 1, 2, 3}")
        return np.stack(keypoints, axis=-1)

    def __repr__(self):
        return self.__class__.__name__ + f"(prob={self.prob})"

    def get_equivalent_preprocessing(self):
        raise RuntimeError(f"{self.__class__} does not have equivalent preprocessing because it is non-deterministic.")

__init__(prob=0.5)

Initialize transform

Parameters:

Name Type Description Default
(float) prob

Probability of applying the transform

 required
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_random_rotate90.py 
17
18
19
20
21
22
23
24
25
26
27
def __init__(
    self,
    prob: float = 0.5,
):
    """
    Initialize transform

    :param prob (float): Probability of applying the transform
    """
    super().__init__()
    self.prob = prob

apply_to_bboxes(bboxes_xywh, factor, rows, cols) classmethod

Parameters:

Name Type Description Default
bboxes

(N, 4) array of bboxes in XYWH format

 required
factor

Number of 90 degree rotations to apply. Order or rotation matches np.rot90

 required
rows int

Number of rows (image height) of the original (input) image

 required
cols int

Number of cols (image width) of the original (input) image

 required

Returns:

Type Description
np.ndarray

Transformed bboxes in XYWH format
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_random_rotate90.py 
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
@classmethod
def apply_to_bboxes(cls, bboxes_xywh: np.ndarray, factor, rows: int, cols: int) -> np.ndarray:
    """

    :param bboxes: (N, 4) array of bboxes in XYWH format
    :param factor: Number of 90 degree rotations to apply. Order or rotation matches np.rot90
    :param rows:   Number of rows (image height) of the original (input) image
    :param cols:   Number of cols (image width) of the original (input) image
    :return:       Transformed bboxes in XYWH format
    """
    from super_gradients.training.transforms.transforms import DetectionRandomRotate90

    bboxes_xyxy = xywh_to_xyxy(bboxes_xywh, image_shape=None)
    bboxes_xyxy = DetectionRandomRotate90.xyxy_bbox_rot90(bboxes_xyxy, factor, rows, cols)
    return xyxy_to_xywh(bboxes_xyxy, image_shape=None)

apply_to_image(image, factor) classmethod

Rotate image by 90 degrees

Parameters:

Name Type Description Default
image np.ndarray

Input image

 required
factor int

Number of 90 degree rotations to apply. Order or rotation matches np.rot90

 required

Returns:

Type Description
np.ndarray

Rotated image
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_random_rotate90.py 
49
50
51
52
53
54
55
56
57
58
@classmethod
def apply_to_image(cls, image: np.ndarray, factor: int) -> np.ndarray:
    """
    Rotate image by 90 degrees

    :param image:  Input image
    :param factor: Number of 90 degree rotations to apply. Order or rotation matches np.rot90
    :return:       Rotated image
    """
    return np.rot90(image, factor)

apply_to_keypoints(keypoints, factor, rows, cols) classmethod

Parameters:

Name Type Description Default
keypoints np.ndarray

Input keypoints array of [Num Instances, Num Joints, 3] shape. Keypoints has format (x, y, visibility)

 required
factor

Number of 90 degree rotations to apply. Order or rotation matches np.rot90

 required
rows int

Number of rows (image height) of the original (input) image

 required
cols int

Number of cols (image width) of the original (input) image

 required

Returns:

Type Description
np.ndarray

Transformed keypoints array of [Num Instances, Num Joints, 3] shape.
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_random_rotate90.py 
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
@classmethod
def apply_to_keypoints(cls, keypoints: np.ndarray, factor, rows: int, cols: int) -> np.ndarray:
    """

    :param keypoints: Input keypoints array of [Num Instances, Num Joints, 3] shape.
                      Keypoints has format (x, y, visibility)
    :param factor:    Number of 90 degree rotations to apply. Order or rotation matches np.rot90
    :param rows:      Number of rows (image height) of the original (input) image
    :param cols:      Number of cols (image width) of the original (input) image
    :return:          Transformed keypoints array of [Num Instances, Num Joints, 3] shape.
    """
    x, y, v = keypoints[:, :, 0], keypoints[:, :, 1], keypoints[:, :, 2]

    if factor == 0:
        keypoints = x, y, v
    elif factor == 1:
        keypoints = y, cols - x - 1, v
    elif factor == 2:
        keypoints = cols - x - 1, rows - y - 1, v
    elif factor == 3:
        keypoints = rows - y - 1, x, v
    else:
        raise ValueError("Parameter n must be in set {0, 1, 2, 3}")
    return np.stack(keypoints, axis=-1)

apply_to_sample(sample)

Returns:

Type Description
PoseEstimationSample

Result of applying the transform
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_random_rotate90.py 
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
    """
    :param   sample: Input PoseEstimationSample
    :return:         Result of applying the transform
    """

    if random.random() < self.prob:
        factor = random.randint(0, 3)

        image_rows, image_cols = sample.image.shape[:2]

        sample.image = self.apply_to_image(sample.image, factor)
        sample.mask = self.apply_to_image(sample.mask, factor)
        sample.joints = self.apply_to_keypoints(sample.joints, factor, image_rows, image_cols)

        if sample.bboxes_xywh is not None:
            sample.bboxes_xywh = self.apply_to_bboxes(sample.bboxes_xywh, factor, image_rows, image_cols)

    return sample

KeypointsRandomVerticalFlip

Bases: AbstractKeypointTransform

Flip image, mask and joints vertically with a given probability.

Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_random_vertical_flip.py 
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
@register_transform(Transforms.KeypointsRandomVerticalFlip)
class KeypointsRandomVerticalFlip(AbstractKeypointTransform):
    """
    Flip image, mask and joints vertically with a given probability.
    """

    def __init__(self, prob: float = 0.5):
        """

        :param prob: Probability of flipping
        """
        super().__init__()
        self.prob = prob

    def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
        """
        Apply transformation to given pose estimation sample

        :param sample: Input pose estimation sample.
        :return:       A new pose estimation sample.
        """
        if sample.image.shape[:2] != sample.mask.shape[:2]:
            raise RuntimeError(f"Image shape ({sample.image.shape[:2]}) does not match mask shape ({sample.mask.shape[:2]}).")

        if random.random() < self.prob:
            sample.image = self.apply_to_image(sample.image)
            sample.mask = self.apply_to_image(sample.mask)
            rows, cols = sample.image.shape[:2]
            sample.joints = self.apply_to_keypoints(sample.joints, rows)

            if sample.bboxes_xywh is not None:
                sample.bboxes_xywh = self.apply_to_bboxes(sample.bboxes_xywh, rows)

        return sample

    def apply_to_image(self, image: np.ndarray) -> np.ndarray:
        """
        Flip image vertically

        :param image: Input image
        :return:      Vertically flipped image
        """
        return np.ascontiguousarray(np.flipud(image))

    def apply_to_keypoints(self, keypoints: np.ndarray, rows: int) -> np.ndarray:
        """
        Flip keypoints vertically

        :param keypoints: Input keypoints of [N,K,3] shape
        :param rows:      Image height
        :return:          Flipped keypoints  of [N,K,3] shape
        """
        keypoints = keypoints.copy()
        keypoints[:, :, 1] = rows - keypoints[:, :, 1] - 1
        return keypoints

    def apply_to_bboxes(self, bboxes: np.ndarray, rows: int) -> np.ndarray:
        """
        Flip boxes vertically

        :param bboxes: Input boxes of [N,4] shape in XYWH format
        :param rows:   Image height
        :return:       Flipped boxes of [N,4] shape in XYWH format
        """

        bboxes = bboxes.copy()
        bboxes[:, 1] = rows - (bboxes[:, 1] + bboxes[:, 3]) - 1
        return bboxes

    def __repr__(self):
        return self.__class__.__name__ + f"(prob={self.prob})"

    def get_equivalent_preprocessing(self):
        raise RuntimeError(f"{self.__class__} does not have equivalent preprocessing because it is non-deterministic.")

__init__(prob=0.5)

Parameters:

Name Type Description Default
prob float

Probability of flipping

 0.5
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_random_vertical_flip.py 
16
17
18
19
20
21
22
def __init__(self, prob: float = 0.5):
    """

    :param prob: Probability of flipping
    """
    super().__init__()
    self.prob = prob

apply_to_bboxes(bboxes, rows)

Flip boxes vertically

Parameters:

Name Type Description Default
bboxes np.ndarray

Input boxes of [N,4] shape in XYWH format

 required
rows int

Image height

 required

Returns:

Type Description
np.ndarray

Flipped boxes of [N,4] shape in XYWH format
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_random_vertical_flip.py 
66
67
68
69
70
71
72
73
74
75
76
77
def apply_to_bboxes(self, bboxes: np.ndarray, rows: int) -> np.ndarray:
    """
    Flip boxes vertically

    :param bboxes: Input boxes of [N,4] shape in XYWH format
    :param rows:   Image height
    :return:       Flipped boxes of [N,4] shape in XYWH format
    """

    bboxes = bboxes.copy()
    bboxes[:, 1] = rows - (bboxes[:, 1] + bboxes[:, 3]) - 1
    return bboxes

apply_to_image(image)

Flip image vertically

Parameters:

Name Type Description Default
image np.ndarray

Input image

 required

Returns:

Type Description
np.ndarray

Vertically flipped image
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_random_vertical_flip.py 
45
46
47
48
49
50
51
52
def apply_to_image(self, image: np.ndarray) -> np.ndarray:
    """
    Flip image vertically

    :param image: Input image
    :return:      Vertically flipped image
    """
    return np.ascontiguousarray(np.flipud(image))

apply_to_keypoints(keypoints, rows)

Flip keypoints vertically

Parameters:

Name Type Description Default
keypoints np.ndarray

Input keypoints of [N,K,3] shape

 required
rows int

Image height

 required

Returns:

Type Description
np.ndarray

Flipped keypoints of [N,K,3] shape
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_random_vertical_flip.py 
54
55
56
57
58
59
60
61
62
63
64
def apply_to_keypoints(self, keypoints: np.ndarray, rows: int) -> np.ndarray:
    """
    Flip keypoints vertically

    :param keypoints: Input keypoints of [N,K,3] shape
    :param rows:      Image height
    :return:          Flipped keypoints  of [N,K,3] shape
    """
    keypoints = keypoints.copy()
    keypoints[:, :, 1] = rows - keypoints[:, :, 1] - 1
    return keypoints

apply_to_sample(sample)

Apply transformation to given pose estimation sample

Parameters:

Name Type Description Default
sample PoseEstimationSample

Input pose estimation sample.

 required

Returns:

Type Description
PoseEstimationSample

A new pose estimation sample.
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_random_vertical_flip.py 
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
    """
    Apply transformation to given pose estimation sample

    :param sample: Input pose estimation sample.
    :return:       A new pose estimation sample.
    """
    if sample.image.shape[:2] != sample.mask.shape[:2]:
        raise RuntimeError(f"Image shape ({sample.image.shape[:2]}) does not match mask shape ({sample.mask.shape[:2]}).")

    if random.random() < self.prob:
        sample.image = self.apply_to_image(sample.image)
        sample.mask = self.apply_to_image(sample.mask)
        rows, cols = sample.image.shape[:2]
        sample.joints = self.apply_to_keypoints(sample.joints, rows)

        if sample.bboxes_xywh is not None:
            sample.bboxes_xywh = self.apply_to_bboxes(sample.bboxes_xywh, rows)

    return sample

KeypointsRemoveSmallObjects

Bases: AbstractKeypointTransform

Remove pose instances from data sample that are too small or have too few visible keypoints.

Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_remove_small_objects.py 
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
@register_transform(Transforms.KeypointsRemoveSmallObjects)
class KeypointsRemoveSmallObjects(AbstractKeypointTransform):
    """
    Remove pose instances from data sample that are too small or have too few visible keypoints.
    """

    def __init__(self, min_visible_keypoints: int = 0, min_instance_area: int = 0, min_bbox_area: int = 0):
        """

        :param min_visible_keypoints: Minimum number of visible keypoints to keep the sample.
                                      Default value is 0 which means that all samples will be kept.
        :param min_instance_area:     Minimum area of instance area to keep the sample
                                      Default value is 0 which means that all samples will be kept.
        :param min_bbox_area:         Minimum area of bounding box to keep the sample
                                      Default value is 0 which means that all samples will be kept.
        """
        super().__init__()
        self.min_visible_keypoints = min_visible_keypoints
        self.min_instance_area = min_instance_area
        self.min_bbox_area = min_bbox_area

    def __call__(
        self, image: np.ndarray, mask: np.ndarray, joints: np.ndarray, areas: Optional[np.ndarray], bboxes: Optional[np.ndarray]
    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray], Optional[np.ndarray]]:
        raise RuntimeError("This transform is not supported for old-style API")

    def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
        """
        Apply transformation to given pose estimation sample.

        :param sample: Input sample to transform.
        :return:       Filtered sample.
        """
        if self.min_visible_keypoints:
            sample = sample.filter_by_visible_joints(self.min_visible_keypoints)
        if self.min_instance_area:
            sample = sample.filter_by_pose_area(self.min_instance_area)
        if self.min_bbox_area:
            sample = sample.filter_by_bbox_area(self.min_bbox_area)
        return sample

    def __repr__(self):
        return self.__class__.__name__ + (
            f"(min_visible_keypoints={self.min_visible_keypoints}, " f"min_instance_area={self.min_instance_area}, " f"min_bbox_area={self.min_bbox_area})"
        )

    def get_equivalent_preprocessing(self) -> List:
        return []

__init__(min_visible_keypoints=0, min_instance_area=0, min_bbox_area=0)

Parameters:

Name Type Description Default
min_visible_keypoints int

Minimum number of visible keypoints to keep the sample. Default value is 0 which means that all samples will be kept.

 0
min_instance_area int

Minimum area of instance area to keep the sample Default value is 0 which means that all samples will be kept.

 0
min_bbox_area int

Minimum area of bounding box to keep the sample Default value is 0 which means that all samples will be kept.

 0
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_remove_small_objects.py 
17
18
19
20
21
22
23
24
25
26
27
28
29
30
def __init__(self, min_visible_keypoints: int = 0, min_instance_area: int = 0, min_bbox_area: int = 0):
    """

    :param min_visible_keypoints: Minimum number of visible keypoints to keep the sample.
                                  Default value is 0 which means that all samples will be kept.
    :param min_instance_area:     Minimum area of instance area to keep the sample
                                  Default value is 0 which means that all samples will be kept.
    :param min_bbox_area:         Minimum area of bounding box to keep the sample
                                  Default value is 0 which means that all samples will be kept.
    """
    super().__init__()
    self.min_visible_keypoints = min_visible_keypoints
    self.min_instance_area = min_instance_area
    self.min_bbox_area = min_bbox_area

apply_to_sample(sample)

Apply transformation to given pose estimation sample.

Parameters:

Name Type Description Default
sample PoseEstimationSample

Input sample to transform.

 required

Returns:

Type Description
PoseEstimationSample

Filtered sample.
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_remove_small_objects.py 
37
38
39
40
41
42
43
44
45
46
47
48
49
50
def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
    """
    Apply transformation to given pose estimation sample.

    :param sample: Input sample to transform.
    :return:       Filtered sample.
    """
    if self.min_visible_keypoints:
        sample = sample.filter_by_visible_joints(self.min_visible_keypoints)
    if self.min_instance_area:
        sample = sample.filter_by_pose_area(self.min_instance_area)
    if self.min_bbox_area:
        sample = sample.filter_by_bbox_area(self.min_bbox_area)
    return sample

KeypointsRescale

Bases: AbstractKeypointTransform

Resize image, mask and joints to target size without preserving aspect ratio.

Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_rescale.py 
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
@register_transform(Transforms.KeypointsRescale)
class KeypointsRescale(AbstractKeypointTransform):
    """
    Resize image, mask and joints to target size without preserving aspect ratio.
    """

    def __init__(self, height: int, width: int, interpolation: int = cv2.INTER_LINEAR, prob: float = 1.0):
        """
        :param height: Target image height
        :param width: Target image width
        :param interpolation: Used interpolation method for image. See cv2.resize for details.
        :param prob: Probability of applying this transform. Default value is 1, meaning that transform is always applied.
        """
        super().__init__()
        self.height = height
        self.width = width
        self.interpolation = interpolation
        self.prob = prob

    def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
        """
        Apply transform to sample.
        :param sample: Input sample
        :return:       Output sample
        """
        if random.random() < self.prob:
            height, width = sample.image.shape[:2]
            sy = self.height / height
            sx = self.width / width

            sample.image = self.apply_to_image(sample.image, dsize=(self.width, self.height), interpolation=self.interpolation)
            sample.mask = self.apply_to_image(sample.mask, dsize=(self.width, self.height), interpolation=cv2.INTER_NEAREST)

            sample.joints = self.apply_to_keypoints(sample.joints, sx, sy)
            if sample.bboxes_xywh is not None:
                sample.bboxes_xywh = self.apply_to_bboxes(sample.bboxes_xywh, sx, sy)

            if sample.areas is not None:
                sample.areas = np.multiply(sample.areas, sx * sy, dtype=np.float32)

        return sample

    @classmethod
    def apply_to_image(cls, img, dsize: Tuple[int, int], interpolation: int) -> np.ndarray:
        """
        Resize image to target size.
        :param img:           Input image
        :param dsize:         Target size (width, height)
        :param interpolation: OpenCV interpolation method
        :return:              Resize image
        """
        img = cv2.resize(img, dsize=dsize, interpolation=interpolation)
        return img

    @classmethod
    def apply_to_keypoints(cls, keypoints: np.ndarray, sx: float, sy: float) -> np.ndarray:
        """
        Resize keypoints to target size.
        :param keypoints: [Num Instances, Num Joints, 3] Input keypoints
        :param sx:        Scale factor along the horizontal axis
        :param sy:        Scale factor along the vertical axis
        :return:          [Num Instances, Num Joints, 3] Resized keypoints
        """
        keypoints = keypoints.astype(np.float32, copy=True)
        keypoints[:, :, 0] *= sx
        keypoints[:, :, 1] *= sy
        return keypoints

    @classmethod
    def apply_to_bboxes(cls, bboxes: np.ndarray, sx: float, sy: float) -> np.ndarray:
        """
        Resize bounding boxes to target size.

        :param bboxes: Input bounding boxes in XYWH format
        :param sx:     Scale factor along the horizontal axis
        :param sy:     Scale factor along the vertical axis
        :return:       Resized bounding boxes in XYWH format
        """
        bboxes = bboxes.astype(np.float32, copy=True)
        bboxes[:, 0::2] *= sx
        bboxes[:, 1::2] *= sy
        return bboxes

    @classmethod
    def apply_to_areas(cls, areas: np.ndarray, sx: float, sy: float) -> np.ndarray:
        """
        Resize areas to target size.
        :param areas: [N] Array of instance areas
        :param sx:    Scale factor along the horizontal axis
        :param sy:    Scale factor along the vertical axis
        :return:      [N] Array of resized instance areas
        """
        return np.multiply(areas, sx * sy, dtype=np.float32)

    def __repr__(self):
        return self.__class__.__name__ + f"(height={self.height}, " f"width={self.width}, " f"interpolation={self.interpolation}, prob={self.prob})"

    def get_equivalent_preprocessing(self) -> List:
        return [{Processings.KeypointsRescale: {"output_shape": (self.height, self.width)}}]

__init__(height, width, interpolation=cv2.INTER_LINEAR, prob=1.0)

Parameters:

Name Type Description Default
height int

Target image height

 required
width int

Target image width

 required
interpolation int

Used interpolation method for image. See cv2.resize for details.

 cv2.INTER_LINEAR
prob float

Probability of applying this transform. Default value is 1, meaning that transform is always applied.

 1.0
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_rescale.py 
18
19
20
21
22
23
24
25
26
27
28
29
def __init__(self, height: int, width: int, interpolation: int = cv2.INTER_LINEAR, prob: float = 1.0):
    """
    :param height: Target image height
    :param width: Target image width
    :param interpolation: Used interpolation method for image. See cv2.resize for details.
    :param prob: Probability of applying this transform. Default value is 1, meaning that transform is always applied.
    """
    super().__init__()
    self.height = height
    self.width = width
    self.interpolation = interpolation
    self.prob = prob

apply_to_areas(areas, sx, sy) classmethod

Resize areas to target size.

Parameters:

Name Type Description Default
areas np.ndarray

[N] Array of instance areas

 required
sx float

Scale factor along the horizontal axis

 required
sy float

Scale factor along the vertical axis

 required

Returns:

Type Description
np.ndarray

[N] Array of resized instance areas
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_rescale.py 
 95
 96
 97
 98
 99
100
101
102
103
104
@classmethod
def apply_to_areas(cls, areas: np.ndarray, sx: float, sy: float) -> np.ndarray:
    """
    Resize areas to target size.
    :param areas: [N] Array of instance areas
    :param sx:    Scale factor along the horizontal axis
    :param sy:    Scale factor along the vertical axis
    :return:      [N] Array of resized instance areas
    """
    return np.multiply(areas, sx * sy, dtype=np.float32)

apply_to_bboxes(bboxes, sx, sy) classmethod

Resize bounding boxes to target size.

Parameters:

Name Type Description Default
bboxes np.ndarray

Input bounding boxes in XYWH format

 required
sx float

Scale factor along the horizontal axis

 required
sy float

Scale factor along the vertical axis

 required

Returns:

Type Description
np.ndarray

Resized bounding boxes in XYWH format
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_rescale.py 
80
81
82
83
84
85
86
87
88
89
90
91
92
93
@classmethod
def apply_to_bboxes(cls, bboxes: np.ndarray, sx: float, sy: float) -> np.ndarray:
    """
    Resize bounding boxes to target size.

    :param bboxes: Input bounding boxes in XYWH format
    :param sx:     Scale factor along the horizontal axis
    :param sy:     Scale factor along the vertical axis
    :return:       Resized bounding boxes in XYWH format
    """
    bboxes = bboxes.astype(np.float32, copy=True)
    bboxes[:, 0::2] *= sx
    bboxes[:, 1::2] *= sy
    return bboxes

apply_to_image(img, dsize, interpolation) classmethod

Resize image to target size.

Parameters:

Name Type Description Default
img

Input image

 required
dsize Tuple[int, int]

Target size (width, height)

 required
interpolation int

OpenCV interpolation method

 required

Returns:

Type Description
np.ndarray

Resize image
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_rescale.py 
54
55
56
57
58
59
60
61
62
63
64
@classmethod
def apply_to_image(cls, img, dsize: Tuple[int, int], interpolation: int) -> np.ndarray:
    """
    Resize image to target size.
    :param img:           Input image
    :param dsize:         Target size (width, height)
    :param interpolation: OpenCV interpolation method
    :return:              Resize image
    """
    img = cv2.resize(img, dsize=dsize, interpolation=interpolation)
    return img

apply_to_keypoints(keypoints, sx, sy) classmethod

Resize keypoints to target size.

Parameters:

Name Type Description Default
keypoints np.ndarray

[Num Instances, Num Joints, 3] Input keypoints

 required
sx float

Scale factor along the horizontal axis

 required
sy float

Scale factor along the vertical axis

 required

Returns:

Type Description
np.ndarray

[Num Instances, Num Joints, 3] Resized keypoints
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_rescale.py 
66
67
68
69
70
71
72
73
74
75
76
77
78
@classmethod
def apply_to_keypoints(cls, keypoints: np.ndarray, sx: float, sy: float) -> np.ndarray:
    """
    Resize keypoints to target size.
    :param keypoints: [Num Instances, Num Joints, 3] Input keypoints
    :param sx:        Scale factor along the horizontal axis
    :param sy:        Scale factor along the vertical axis
    :return:          [Num Instances, Num Joints, 3] Resized keypoints
    """
    keypoints = keypoints.astype(np.float32, copy=True)
    keypoints[:, :, 0] *= sx
    keypoints[:, :, 1] *= sy
    return keypoints

apply_to_sample(sample)

Apply transform to sample.

Parameters:

Name Type Description Default
sample PoseEstimationSample

Input sample

 required

Returns:

Type Description
PoseEstimationSample

Output sample
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_rescale.py 
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
    """
    Apply transform to sample.
    :param sample: Input sample
    :return:       Output sample
    """
    if random.random() < self.prob:
        height, width = sample.image.shape[:2]
        sy = self.height / height
        sx = self.width / width

        sample.image = self.apply_to_image(sample.image, dsize=(self.width, self.height), interpolation=self.interpolation)
        sample.mask = self.apply_to_image(sample.mask, dsize=(self.width, self.height), interpolation=cv2.INTER_NEAREST)

        sample.joints = self.apply_to_keypoints(sample.joints, sx, sy)
        if sample.bboxes_xywh is not None:
            sample.bboxes_xywh = self.apply_to_bboxes(sample.bboxes_xywh, sx, sy)

        if sample.areas is not None:
            sample.areas = np.multiply(sample.areas, sx * sy, dtype=np.float32)

    return sample

KeypointsReverseImageChannels

Bases: AbstractKeypointTransform

Randomly reverse channel order with given probability. Given an image with RGB channels, when applied with probability 1, it returns an image with BGR channels. With probability 0.5 there is 50/50 chance to return BGR or RGB image. It usually helps to improve model's ability to generalize under different color channels order.

Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_reverse_image_channels.py 
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
@register_transform()
class KeypointsReverseImageChannels(AbstractKeypointTransform):
    """
    Randomly reverse channel order with given probability.
    Given an image with RGB channels, when applied with probability 1, it returns an image with BGR channels.
    With probability 0.5 there is 50/50 chance to return BGR or RGB image.
    It usually helps to improve model's ability to generalize under different color channels order.
    """

    def __init__(self, prob: float):
        """

        :param prob:             Probability to apply the transform.
        """
        super().__init__()
        self.prob = prob

    def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
        if random.random() < self.prob:
            sample.image = np.ascontiguousarray(sample.image[:, :, ::-1])
        return sample

    def get_equivalent_preprocessing(self) -> List:
        if self.prob < 1:
            raise RuntimeError("Cannot set preprocessing pipeline with randomness. Set prob to 1.")
        return [Processings.ReverseImageChannels]

__init__(prob)

Parameters:

Name Type Description Default
prob float

Probability to apply the transform.

 required
Source code in V3_4/src/super_gradients/training/transforms/keypoints/keypoints_reverse_image_channels.py 
21
22
23
24
25
26
27
def __init__(self, prob: float):
    """

    :param prob:             Probability to apply the transform.
    """
    super().__init__()
    self.prob = prob

DetectionHSV

Bases: DetectionTransform

Detection HSV transform.

Parameters:

Name Type Description Default
prob float

Probability to apply the transform.

 required
hgain float

Hue gain.

 0.5
sgain float

Saturation gain.

 0.5
vgain float

Value gain.

 0.5
bgr_channels

Channel indices of the BGR channels- useful for images with >3 channels, or when BGR channels are in different order.

 (0, 1, 2)
Source code in V3_4/src/super_gradients/training/transforms/transforms.py 
 985
 986
 987
 988
 989
 990
 991
 992
 993
 994
 995
 996
 997
 998
 999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
@register_transform(Transforms.DetectionHSV)
class DetectionHSV(DetectionTransform):
    """
    Detection HSV transform.

    :param prob:            Probability to apply the transform.
    :param hgain:           Hue gain.
    :param sgain:           Saturation gain.
    :param vgain:           Value gain.
    :param bgr_channels:    Channel indices of the BGR channels- useful for images with >3 channels, or when BGR channels are in different order.
    """

    def __init__(self, prob: float, hgain: float = 0.5, sgain: float = 0.5, vgain: float = 0.5, bgr_channels=(0, 1, 2)):
        super(DetectionHSV, self).__init__()
        self.prob = prob
        self.hgain = hgain
        self.sgain = sgain
        self.vgain = vgain
        self.bgr_channels = bgr_channels
        self._additional_channels_warned = False

    def __call__(self, sample: dict) -> dict:
        if sample["image"].shape[2] < 3:
            raise ValueError("HSV transform expects at least 3 channels, got: " + str(sample["image"].shape[2]))
        if sample["image"].shape[2] > 3 and not self._additional_channels_warned:
            logger.warning(
                "HSV transform received image with "
                + str(sample["image"].shape[2])
                + " channels. HSV transform will only be applied on channels: "
                + str(self.bgr_channels)
                + "."
            )
            self._additional_channels_warned = True
        if random.random() < self.prob:
            augment_hsv(sample["image"], self.hgain, self.sgain, self.vgain, self.bgr_channels)
        return sample

DetectionHorizontalFlip

Bases: DetectionTransform

Horizontal Flip for Detection

Parameters:

Name Type Description Default
prob float

Probability of applying horizontal flip

 required
Source code in V3_4/src/super_gradients/training/transforms/transforms.py 
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
@register_transform(Transforms.DetectionHorizontalFlip)
class DetectionHorizontalFlip(DetectionTransform):
    """
    Horizontal Flip for Detection

    :param prob:        Probability of applying horizontal flip
    """

    def __init__(self, prob: float, max_targets: Optional[int] = None):
        super(DetectionHorizontalFlip, self).__init__()
        _max_targets_deprication(max_targets)
        self.prob = prob

    def __call__(self, sample):
        image, targets = sample["image"], sample["target"]
        crowd_targets = sample.get("crowd_targets")
        if len(targets) == 0:
            targets = np.zeros((0, 5), dtype=np.float32)
        if random.random() < self.prob:
            image = _flip_horizontal_image(image)
            _, width, _ = image.shape
            targets[:, :4] = _flip_horizontal_boxes(targets[:, :4], width)
            if crowd_targets is not None:
                crowd_targets = _flip_horizontal_boxes(crowd_targets, width)
        sample["image"] = image
        sample["target"] = targets
        sample["crowd_targets"] = crowd_targets
        return sample

DetectionImagePermute

Bases: DetectionTransform

Permute image dims. Useful for converting image from HWC to CHW format.

Source code in V3_4/src/super_gradients/training/transforms/transforms.py 
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
@register_transform(Transforms.DetectionImagePermute)
class DetectionImagePermute(DetectionTransform):
    """
    Permute image dims. Useful for converting image from HWC to CHW format.
    """

    def __init__(self, dims: Tuple[int, int, int] = (2, 0, 1)):
        """

        :param dims: Specify new order of dims. Default value (2, 0, 1) suitable for converting from HWC to CHW format.
        """
        super().__init__()
        self.dims = tuple(dims)

    def __call__(self, sample: Dict[str, np.array]) -> dict:
        sample["image"] = np.ascontiguousarray(sample["image"].transpose(*self.dims))
        return sample

    def get_equivalent_preprocessing(self) -> List[Dict]:
        return [{Processings.ImagePermute: {"permutation": self.dims}}]

__init__(dims=(2, 0, 1))

Parameters:

Name Type Description Default
dims Tuple[int, int, int]

Specify new order of dims. Default value (2, 0, 1) suitable for converting from HWC to CHW format.

 (2, 0, 1)
Source code in V3_4/src/super_gradients/training/transforms/transforms.py 
726
727
728
729
730
731
732
def __init__(self, dims: Tuple[int, int, int] = (2, 0, 1)):
    """

    :param dims: Specify new order of dims. Default value (2, 0, 1) suitable for converting from HWC to CHW format.
    """
    super().__init__()
    self.dims = tuple(dims)

DetectionMixup

Bases: DetectionTransform

Mixup detection transform

Parameters:

Name Type Description Default
input_dim Union[int, Tuple[int, int], None]

Input dimension.

 required
mixup_scale tuple

Scale range for the additional loaded image for mixup.

 required
prob float

Probability of applying mixup.

 1.0
enable_mixup bool

Whether to apply mixup at all (regardless of prob).

 True
flip_prob float

Probability to apply horizontal flip to the additional sample.

 0.5
border_value int

Value for filling borders after applying transform.

 114
Source code in V3_4/src/super_gradients/training/transforms/transforms.py 
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
@register_transform(Transforms.DetectionMixup)
class DetectionMixup(DetectionTransform):
    """
    Mixup detection transform

    :param input_dim:        Input dimension.
    :param mixup_scale:      Scale range for the additional loaded image for mixup.
    :param prob:             Probability of applying mixup.
    :param enable_mixup:     Whether to apply mixup at all (regardless of prob).
    :param flip_prob:        Probability to apply horizontal flip to the additional sample.
    :param border_value:     Value for filling borders after applying transform.
    """

    def __init__(
        self,
        input_dim: Union[int, Tuple[int, int], None],
        mixup_scale: tuple,
        prob: float = 1.0,
        enable_mixup: bool = True,
        flip_prob: float = 0.5,
        border_value: int = 114,
    ):
        super(DetectionMixup, self).__init__(additional_samples_count=1, non_empty_targets=True)
        self.input_dim = ensure_is_tuple_of_two(input_dim)
        self.mixup_scale = mixup_scale
        self.prob = prob
        self.enable_mixup = enable_mixup
        self.flip_prob = flip_prob
        self.border_value = border_value

    def close(self):
        self.additional_samples_count = 0
        self.enable_mixup = False

    def __call__(self, sample: dict) -> dict:
        if self.enable_mixup and random.random() < self.prob:
            origin_img, origin_labels = sample["image"], sample["target"]
            target_dim = self.input_dim if self.input_dim is not None else sample["image"].shape[:2]

            cp_sample = sample["additional_samples"][0]
            img, cp_labels = cp_sample["image"], cp_sample["target"]
            cp_boxes = cp_labels[:, :4]
            if random.random() < self.prob:
                _, width, _ = img.shape
                img = _flip_horizontal_image(img)
                cp_boxes = _flip_horizontal_boxes(cp_boxes, width)
            # PLUG IN TARGET THE FLIPPED BOXES
            cp_labels[:, :4] = cp_boxes

            jit_factor = random.uniform(*self.mixup_scale)

            if len(img.shape) == 3:
                cp_img = np.ones((target_dim[0], target_dim[1], 3), dtype=np.uint8) * self.border_value
            else:
                cp_img = np.ones(target_dim, dtype=np.uint8) * self.border_value

            cp_scale_ratio = min(target_dim[0] / img.shape[0], target_dim[1] / img.shape[1])
            resized_img = cv2.resize(
                img,
                (int(img.shape[1] * cp_scale_ratio), int(img.shape[0] * cp_scale_ratio)),
                interpolation=cv2.INTER_LINEAR,
            )

            cp_img[: int(img.shape[0] * cp_scale_ratio), : int(img.shape[1] * cp_scale_ratio)] = resized_img

            cp_img = cv2.resize(
                cp_img,
                (int(cp_img.shape[1] * jit_factor), int(cp_img.shape[0] * jit_factor)),
            )
            cp_scale_ratio *= jit_factor

            origin_h, origin_w = cp_img.shape[:2]
            target_h, target_w = origin_img.shape[:2]

            if len(img.shape) == 3:
                padded_img = np.zeros((max(origin_h, target_h), max(origin_w, target_w), img.shape[2]), dtype=np.uint8)
            else:
                padded_img = np.zeros((max(origin_h, target_h), max(origin_w, target_w)), dtype=np.uint8)

            padded_img[:origin_h, :origin_w] = cp_img

            x_offset, y_offset = 0, 0
            if padded_img.shape[0] > target_h:
                y_offset = random.randint(0, padded_img.shape[0] - target_h - 1)
            if padded_img.shape[1] > target_w:
                x_offset = random.randint(0, padded_img.shape[1] - target_w - 1)
            padded_cropped_img = padded_img[y_offset : y_offset + target_h, x_offset : x_offset + target_w]

            cp_bboxes_origin_np = adjust_box_anns(cp_labels[:, :4].copy(), cp_scale_ratio, 0, 0, origin_w, origin_h)
            cp_bboxes_transformed_np = cp_bboxes_origin_np.copy()
            cp_bboxes_transformed_np[:, 0::2] = np.clip(cp_bboxes_transformed_np[:, 0::2] - x_offset, 0, target_w)
            cp_bboxes_transformed_np[:, 1::2] = np.clip(cp_bboxes_transformed_np[:, 1::2] - y_offset, 0, target_h)

            cls_labels = cp_labels[:, 4:5].copy()
            box_labels = cp_bboxes_transformed_np
            labels = np.hstack((box_labels, cls_labels))
            origin_labels = np.vstack((origin_labels, labels))
            origin_img = origin_img.astype(np.float32)
            origin_img = 0.5 * origin_img + 0.5 * padded_cropped_img.astype(np.float32)

            sample["image"], sample["target"] = origin_img.astype(np.uint8), origin_labels
        return sample

DetectionMosaic

Bases: DetectionTransform

DetectionMosaic detection transform

Parameters:

Name Type Description Default
input_dim Union[int, Tuple[int, int]]

Input dimension.

 required
prob float

Probability of applying mosaic.

 1.0
enable_mosaic bool

Whether to apply mosaic at all (regardless of prob).

 True
border_value

Value for filling borders after applying transforms.

 114
Source code in V3_4/src/super_gradients/training/transforms/transforms.py 
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
@register_transform(Transforms.DetectionMosaic)
class DetectionMosaic(DetectionTransform):
    """
    DetectionMosaic detection transform

    :param input_dim:       Input dimension.
    :param prob:            Probability of applying mosaic.
    :param enable_mosaic:   Whether to apply mosaic at all (regardless of prob).
    :param border_value:    Value for filling borders after applying transforms.
    """

    def __init__(self, input_dim: Union[int, Tuple[int, int]], prob: float = 1.0, enable_mosaic: bool = True, border_value=114):
        super(DetectionMosaic, self).__init__(additional_samples_count=3)
        self.prob = prob
        self.input_dim = ensure_is_tuple_of_two(input_dim)
        self.enable_mosaic = enable_mosaic
        self.border_value = border_value

    def close(self):
        self.additional_samples_count = 0
        self.enable_mosaic = False

    def __call__(self, sample: Union[dict, list]):
        if self.enable_mosaic and random.random() < self.prob:
            mosaic_labels = []
            mosaic_labels_seg = []
            input_h, input_w = self.input_dim[0], self.input_dim[1]

            # yc, xc = s, s  # mosaic center x, y
            yc = int(random.uniform(0.5 * input_h, 1.5 * input_h))
            xc = int(random.uniform(0.5 * input_w, 1.5 * input_w))

            # 3 additional samples, total of 4
            all_samples = [sample] + sample["additional_samples"]

            for i_mosaic, mosaic_sample in enumerate(all_samples):
                img, _labels = mosaic_sample["image"], mosaic_sample["target"]
                _labels_seg = mosaic_sample.get("target_seg")

                h0, w0 = img.shape[:2]  # orig hw
                scale = min(1.0 * input_h / h0, 1.0 * input_w / w0)
                img = cv2.resize(img, (int(w0 * scale), int(h0 * scale)), interpolation=cv2.INTER_LINEAR)
                # generate output mosaic image
                (h, w, c) = img.shape[:3]
                if i_mosaic == 0:
                    mosaic_img = np.full((input_h * 2, input_w * 2, c), self.border_value, dtype=np.uint8)

                # suffix l means large image, while s means small image in mosaic aug.
                (l_x1, l_y1, l_x2, l_y2), (s_x1, s_y1, s_x2, s_y2) = get_mosaic_coordinate(i_mosaic, xc, yc, w, h, input_h, input_w)

                mosaic_img[l_y1:l_y2, l_x1:l_x2] = img[s_y1:s_y2, s_x1:s_x2]
                padw, padh = l_x1 - s_x1, l_y1 - s_y1

                labels = _labels.copy()

                # Normalized xywh to pixel xyxy format
                if _labels.size > 0:
                    labels[:, 0] = scale * _labels[:, 0] + padw
                    labels[:, 1] = scale * _labels[:, 1] + padh
                    labels[:, 2] = scale * _labels[:, 2] + padw
                    labels[:, 3] = scale * _labels[:, 3] + padh
                mosaic_labels.append(labels)

                if _labels_seg is not None:
                    labels_seg = _labels_seg.copy()
                    if _labels.size > 0:
                        labels_seg[:, ::2] = scale * labels_seg[:, ::2] + padw
                        labels_seg[:, 1::2] = scale * labels_seg[:, 1::2] + padh
                    mosaic_labels_seg.append(labels_seg)

            if len(mosaic_labels):
                mosaic_labels = np.concatenate(mosaic_labels, 0)
                np.clip(mosaic_labels[:, 0], 0, 2 * input_w, out=mosaic_labels[:, 0])
                np.clip(mosaic_labels[:, 1], 0, 2 * input_h, out=mosaic_labels[:, 1])
                np.clip(mosaic_labels[:, 2], 0, 2 * input_w, out=mosaic_labels[:, 2])
                np.clip(mosaic_labels[:, 3], 0, 2 * input_h, out=mosaic_labels[:, 3])

            if len(mosaic_labels_seg):
                mosaic_labels_seg = np.concatenate(mosaic_labels_seg, 0)
                np.clip(mosaic_labels_seg[:, ::2], 0, 2 * input_w, out=mosaic_labels_seg[:, ::2])
                np.clip(mosaic_labels_seg[:, 1::2], 0, 2 * input_h, out=mosaic_labels_seg[:, 1::2])

            sample["image"] = mosaic_img
            sample["target"] = mosaic_labels
            sample["info"] = (mosaic_img.shape[1], mosaic_img.shape[0])
            if len(mosaic_labels_seg):
                sample["target_seg"] = mosaic_labels_seg

        return sample

DetectionNormalize

Bases: DetectionTransform

Normalize image by subtracting mean and dividing by std.

Source code in V3_4/src/super_gradients/training/transforms/transforms.py 
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
@register_transform(Transforms.DetectionNormalize)
class DetectionNormalize(DetectionTransform):
    """
    Normalize image by subtracting mean and dividing by std.
    """

    def __init__(self, mean, std):
        super().__init__()
        self.mean = np.array(list(mean)).reshape((1, 1, -1)).astype(np.float32)
        self.std = np.array(list(std)).reshape((1, 1, -1)).astype(np.float32)

    def __call__(self, sample: dict) -> dict:
        sample["image"] = (sample["image"] - self.mean) / self.std
        return sample

    def get_equivalent_preprocessing(self) -> List[Dict]:
        return [{Processings.NormalizeImage: {"mean": self.mean, "std": self.std}}]

DetectionPadToSize

Bases: DetectionTransform

Preprocessing transform to pad image and bboxes to input_dim shape (rows, cols). Transform does center padding, so that input image with bboxes located in the center of the produced image.

Note: This transformation assume that dimensions of input image is equal or less than output_size.

Source code in V3_4/src/super_gradients/training/transforms/transforms.py 
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
@register_transform(Transforms.DetectionPadToSize)
class DetectionPadToSize(DetectionTransform):
    """
    Preprocessing transform to pad image and bboxes to `input_dim` shape (rows, cols).
    Transform does center padding, so that input image with bboxes located in the center of the produced image.

    Note: This transformation assume that dimensions of input image is equal or less than `output_size`.
    """

    def __init__(self, output_size: Union[int, Tuple[int, int], None], pad_value: int):
        """
        Constructor for DetectionPadToSize transform.

        :param output_size: Output image size (rows, cols)
        :param pad_value: Padding value for image
        """
        super().__init__()
        self.output_size = ensure_is_tuple_of_two(output_size)
        self.pad_value = pad_value

    def __call__(self, sample: dict) -> dict:
        image, targets, crowd_targets = sample["image"], sample["target"], sample.get("crowd_target")
        padding_coordinates = _get_center_padding_coordinates(input_shape=image.shape, output_shape=self.output_size)

        sample["image"] = _pad_image(image=image, padding_coordinates=padding_coordinates, pad_value=self.pad_value)
        sample["target"] = _shift_bboxes(targets=targets, shift_w=padding_coordinates.left, shift_h=padding_coordinates.top)
        if crowd_targets is not None:
            sample["crowd_target"] = _shift_bboxes(targets=crowd_targets, shift_w=padding_coordinates.left, shift_h=padding_coordinates.top)
        return sample

    def get_equivalent_preprocessing(self) -> List:
        return [{Processings.DetectionCenterPadding: {"output_shape": self.output_size, "pad_value": self.pad_value}}]

__init__(output_size, pad_value)

Constructor for DetectionPadToSize transform.

Parameters:

Name Type Description Default
output_size Union[int, Tuple[int, int], None]

Output image size (rows, cols)

 required
pad_value int

Padding value for image

 required
Source code in V3_4/src/super_gradients/training/transforms/transforms.py 
751
752
753
754
755
756
757
758
759
760
def __init__(self, output_size: Union[int, Tuple[int, int], None], pad_value: int):
    """
    Constructor for DetectionPadToSize transform.

    :param output_size: Output image size (rows, cols)
    :param pad_value: Padding value for image
    """
    super().__init__()
    self.output_size = ensure_is_tuple_of_two(output_size)
    self.pad_value = pad_value

DetectionPaddedRescale

Bases: DetectionTransform

Preprocessing transform to be applied last of all transforms for validation.

Image- Rescales and pads to self.input_dim. Targets- moves the class label to first index, converts boxes format- xyxy -> cxcywh.

Parameters:

Name Type Description Default
input_dim Union[int, Tuple[int, int], None]

Final input dimension (default=(640,640))

 required
swap Tuple[int, ...]

Image axis's to be rearranged.

 (2, 0, 1)
pad_value int

Padding value for image.

 114
Source code in V3_4/src/super_gradients/training/transforms/transforms.py 
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
@register_transform(Transforms.DetectionPaddedRescale)
class DetectionPaddedRescale(DetectionTransform):
    """
    Preprocessing transform to be applied last of all transforms for validation.

    Image- Rescales and pads to self.input_dim.
    Targets- moves the class label to first index, converts boxes format- xyxy -> cxcywh.

    :param input_dim:   Final input dimension (default=(640,640))
    :param swap:        Image axis's to be rearranged.
    :param pad_value:   Padding value for image.
    """

    def __init__(
        self, input_dim: Union[int, Tuple[int, int], None], swap: Tuple[int, ...] = (2, 0, 1), max_targets: Optional[int] = None, pad_value: int = 114
    ):
        super().__init__()
        _max_targets_deprication(max_targets)
        self.swap = swap
        self.input_dim = ensure_is_tuple_of_two(input_dim)
        self.pad_value = pad_value

    def __call__(self, sample: dict) -> dict:
        img, targets, crowd_targets = sample["image"], sample["target"], sample.get("crowd_target")
        img, r = _rescale_and_pad_to_size(img, self.input_dim, self.swap, self.pad_value)

        sample["image"] = img
        sample["target"] = _rescale_xyxy_bboxes(targets, r)
        if crowd_targets is not None:
            sample["crowd_target"] = _rescale_xyxy_bboxes(crowd_targets, r)
        return sample

    def get_equivalent_preprocessing(self) -> List[Dict]:
        return [
            {Processings.DetectionLongestMaxSizeRescale: {"output_shape": self.input_dim}},
            {Processings.DetectionBottomRightPadding: {"output_shape": self.input_dim, "pad_value": self.pad_value}},
            {Processings.ImagePermute: {"permutation": self.swap}},
        ]

DetectionRGB2BGR

Bases: DetectionTransform

Detection change Red & Blue channel of the image

Parameters:

Name Type Description Default
prob float

Probability to apply the transform.

 0.5
Source code in V3_4/src/super_gradients/training/transforms/transforms.py 
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
@register_transform(Transforms.DetectionRGB2BGR)
class DetectionRGB2BGR(DetectionTransform):
    """
    Detection change Red & Blue channel of the image

    :param prob: Probability to apply the transform.
    """

    def __init__(self, prob: float = 0.5):
        super().__init__()
        self.prob = prob

    def __call__(self, sample: dict) -> dict:
        if sample["image"].shape[2] < 3:
            raise ValueError("DetectionRGB2BGR transform expects at least 3 channels, got: " + str(sample["image"].shape[2]))

        if random.random() < self.prob:
            sample["image"] = sample["image"][..., ::-1]
        return sample

    def get_equivalent_preprocessing(self) -> List:
        if self.prob < 1:
            raise RuntimeError("Cannot set preprocessing pipeline with randomness. Set prob to 1.")
        return [{Processings.ReverseImageChannels: {}}]

DetectionRandomAffine

Bases: DetectionTransform

DetectionRandomAffine detection transform

:param degrees: Degrees for random rotation, when float the random values are drawn uniformly from (-degrees, degrees) :param translate: Translate size (in pixels) for random translation, when float the random values are drawn uniformly from (center-translate, center+translate) :param scales: Values for random rescale, when float the random values are drawn uniformly from (1-scales, 1+scales) :param shear: Degrees for random shear, when float the random values are drawn uniformly from (-shear, shear) :param target_size: Desired output shape. :param filter_box_candidates: Whether to filter out transformed bboxes by edge size, area ratio, and aspect ratio (default=False). :param wh_thr: Edge size threshold when filter_box_candidates = True. Bounding oxes with edges smaller than this values will be filtered out. :param ar_thr: Aspect ratio threshold filter_box_candidates = True. Bounding boxes with aspect ratio larger than this values will be filtered out. :param area_thr: Threshold for area ratio between original image and the transformed one, when filter_box_candidates = True. Bounding boxes with such ratio smaller than this value will be filtered out. :param border_value: Value for filling borders after applying transforms.

Source code in V3_4/src/super_gradients/training/transforms/transforms.py 
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
@register_transform(Transforms.DetectionRandomAffine)
class DetectionRandomAffine(DetectionTransform):
    """
    DetectionRandomAffine detection transform

     :param degrees:                Degrees for random rotation, when float the random values are drawn uniformly from (-degrees, degrees)
     :param translate:              Translate size (in pixels) for random translation, when float the random values are drawn uniformly from
                                    (center-translate, center+translate)
     :param scales:                 Values for random rescale, when float the random values are drawn uniformly from (1-scales, 1+scales)
     :param shear:                  Degrees for random shear, when float the random values are drawn uniformly from (-shear, shear)
     :param target_size:            Desired output shape.
     :param filter_box_candidates:  Whether to filter out transformed bboxes by edge size, area ratio, and aspect ratio (default=False).
     :param wh_thr:                 Edge size threshold when filter_box_candidates = True.
                                    Bounding oxes with edges smaller than this values will be filtered out.
     :param ar_thr:                 Aspect ratio threshold filter_box_candidates = True.
                                    Bounding boxes with aspect ratio larger than this values will be filtered out.
     :param area_thr:               Threshold for area ratio between original image and the transformed one, when filter_box_candidates = True.
                                    Bounding boxes with such ratio smaller than this value will be filtered out.
     :param border_value:           Value for filling borders after applying transforms.
    """

    def __init__(
        self,
        degrees: Union[tuple, float] = 10,
        translate: Union[tuple, float] = 0.1,
        scales: Union[tuple, float] = 0.1,
        shear: Union[tuple, float] = 10,
        target_size: Union[int, Tuple[int, int], None] = (640, 640),
        filter_box_candidates: bool = False,
        wh_thr: float = 2,
        ar_thr: float = 20,
        area_thr: float = 0.1,
        border_value: int = 114,
    ):
        super(DetectionRandomAffine, self).__init__()
        self.degrees = degrees
        self.translate = translate
        self.scale = scales
        self.shear = shear
        self.target_size = ensure_is_tuple_of_two(target_size)
        self.enable = True
        self.filter_box_candidates = filter_box_candidates
        self.wh_thr = wh_thr
        self.ar_thr = ar_thr
        self.area_thr = area_thr
        self.border_value = border_value

    def close(self):
        self.enable = False

    def __call__(self, sample: dict) -> dict:
        if self.enable:
            img, target = random_affine(
                sample["image"],
                sample["target"],
                sample.get("target_seg"),
                target_size=self.target_size or tuple(reversed(sample["image"].shape[:2])),
                degrees=self.degrees,
                translate=self.translate,
                scales=self.scale,
                shear=self.shear,
                filter_box_candidates=self.filter_box_candidates,
                wh_thr=self.wh_thr,
                area_thr=self.area_thr,
                ar_thr=self.ar_thr,
                border_value=self.border_value,
            )
            sample["image"] = img
            sample["target"] = target
        return sample

DetectionRandomRotate90

Bases: DetectionTransform

Source code in V3_4/src/super_gradients/training/transforms/transforms.py 
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
@register_transform(Transforms.DetectionRandomRotate90)
class DetectionRandomRotate90(DetectionTransform):
    def __init__(self, prob: float = 0.5):
        super().__init__()
        self.prob = prob

    def __call__(self, sample: dict) -> dict:
        if random.random() < self.prob:
            k = random.randrange(0, 4)

            img, targets, crowd_targets = sample["image"], sample["target"], sample.get("crowd_target")

            sample["image"] = np.ascontiguousarray(np.rot90(img, k))
            sample["target"] = self.rotate_bboxes(targets, k, img.shape[:2])
            if crowd_targets is not None:
                sample["crowd_target"] = self.rotate_bboxes(crowd_targets, k, img.shape[:2])

        return sample

    @classmethod
    def rotate_bboxes(cls, targets, k: int, image_shape):
        if k == 0:
            return targets
        rows, cols = image_shape
        targets = targets.copy()
        targets[:, 0:4] = cls.xyxy_bbox_rot90(targets[:, 0:4], k, rows, cols)
        return targets

    @classmethod
    def xyxy_bbox_rot90(cls, bboxes: np.ndarray, factor: int, rows: int, cols: int):
        """
        Rotates a bounding box by 90 degrees CCW (see np.rot90)

        :param bboxes:  Tensor made of bounding box tuples (x_min, y_min, x_max, y_max).
        :param factor:  Number of CCW rotations. Must be in set {0, 1, 2, 3} See np.rot90.
        :param rows:    Image rows.
        :param cols:    Image cols.

        :return: A bounding box tuple (x_min, y_min, x_max, y_max).

        """
        x_min, y_min, x_max, y_max = bboxes[:, 0], bboxes[:, 1], bboxes[:, 2], bboxes[:, 3]

        if factor == 0:
            bbox = x_min, y_min, x_max, y_max
        elif factor == 1:
            bbox = y_min, cols - x_max, y_max, cols - x_min
        elif factor == 2:
            bbox = cols - x_max, rows - y_max, cols - x_min, rows - y_min
        elif factor == 3:
            bbox = rows - y_max, x_min, rows - y_min, x_max
        else:
            raise ValueError("Parameter n must be in set {0, 1, 2, 3}")
        return np.stack(bbox, axis=1)

xyxy_bbox_rot90(bboxes, factor, rows, cols) classmethod

Rotates a bounding box by 90 degrees CCW (see np.rot90)

Parameters:

Name Type Description Default
bboxes np.ndarray

Tensor made of bounding box tuples (x_min, y_min, x_max, y_max).

 required
factor int

Number of CCW rotations. Must be in set {0, 1, 2, 3} See np.rot90.

 required
rows int

Image rows.

 required
cols int

Image cols.

 required

Returns:

Type Description

A bounding box tuple (x_min, y_min, x_max, y_max).
Source code in V3_4/src/super_gradients/training/transforms/transforms.py 
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
@classmethod
def xyxy_bbox_rot90(cls, bboxes: np.ndarray, factor: int, rows: int, cols: int):
    """
    Rotates a bounding box by 90 degrees CCW (see np.rot90)

    :param bboxes:  Tensor made of bounding box tuples (x_min, y_min, x_max, y_max).
    :param factor:  Number of CCW rotations. Must be in set {0, 1, 2, 3} See np.rot90.
    :param rows:    Image rows.
    :param cols:    Image cols.

    :return: A bounding box tuple (x_min, y_min, x_max, y_max).

    """
    x_min, y_min, x_max, y_max = bboxes[:, 0], bboxes[:, 1], bboxes[:, 2], bboxes[:, 3]

    if factor == 0:
        bbox = x_min, y_min, x_max, y_max
    elif factor == 1:
        bbox = y_min, cols - x_max, y_max, cols - x_min
    elif factor == 2:
        bbox = cols - x_max, rows - y_max, cols - x_min, rows - y_min
    elif factor == 3:
        bbox = rows - y_max, x_min, rows - y_min, x_max
    else:
        raise ValueError("Parameter n must be in set {0, 1, 2, 3}")
    return np.stack(bbox, axis=1)

DetectionRescale

Bases: DetectionTransform

Resize image and bounding boxes to given image dimensions without preserving aspect ratio

Parameters:

Name Type Description Default
output_shape Union[int, Tuple[int, int]]

(rows, cols)

 required
Source code in V3_4/src/super_gradients/training/transforms/transforms.py 
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
@register_transform(Transforms.DetectionRescale)
class DetectionRescale(DetectionTransform):
    """
    Resize image and bounding boxes to given image dimensions without preserving aspect ratio

    :param output_shape: (rows, cols)
    """

    def __init__(self, output_shape: Union[int, Tuple[int, int]]):
        super().__init__()
        self.output_shape = ensure_is_tuple_of_two(output_shape)

    def __call__(self, sample: dict) -> dict:
        image, targets, crowd_targets = sample["image"], sample["target"], sample.get("crowd_target")

        sy, sx = float(self.output_shape[0]) / float(image.shape[0]), float(self.output_shape[1]) / float(image.shape[1])

        sample["image"] = _rescale_image(image=image, target_shape=self.output_shape)
        sample["target"] = _rescale_bboxes(targets, scale_factors=(sy, sx))
        if crowd_targets is not None:
            sample["crowd_target"] = _rescale_bboxes(crowd_targets, scale_factors=(sy, sx))
        return sample

    def get_equivalent_preprocessing(self) -> List[Dict]:
        return [{Processings.DetectionRescale: {"output_shape": self.output_shape}}]

DetectionStandardize

Bases: DetectionTransform

Standardize image pixel values with img/max_val

Parameters:

Name Type Description Default
max_val

Current maximum value of the image pixels. (usually 255)

 required
Source code in V3_4/src/super_gradients/training/transforms/transforms.py 
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
@register_transform(Transforms.DetectionStandardize)
class DetectionStandardize(DetectionTransform):
    """
    Standardize image pixel values with img/max_val

    :param max_val: Current maximum value of the image pixels. (usually 255)
    """

    def __init__(self, max_value: float = 255.0):
        super().__init__()
        self.max_value = max_value

    def __call__(self, sample: dict) -> dict:
        sample["image"] = (sample["image"] / self.max_value).astype(np.float32)
        return sample

    def get_equivalent_preprocessing(self) -> List[Dict]:
        return [{Processings.StandardizeImage: {"max_value": self.max_value}}]

DetectionTargetsFormatTransform

Bases: DetectionTransform

Detection targets format transform

Convert targets in input_format to output_format, filter small bboxes and pad targets.

Parameters:

Name Type Description Default
input_dim Union[int, Tuple[int, int], None]

Shape of the images to transform.

 None
input_format ConcatenatedTensorFormat

Format of the input targets. For instance [xmin, ymin, xmax, ymax, cls_id] refers to XYXY_LABEL.

 XYXY_LABEL
output_format ConcatenatedTensorFormat

Format of the output targets. For instance [xmin, ymin, xmax, ymax, cls_id] refers to XYXY_LABEL

 LABEL_CXCYWH
min_bbox_edge_size float

bboxes with edge size lower then this values will be removed.

 1
Source code in V3_4/src/super_gradients/training/transforms/transforms.py 
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
@register_transform(Transforms.DetectionTargetsFormatTransform)
class DetectionTargetsFormatTransform(DetectionTransform):
    """
    Detection targets format transform

    Convert targets in input_format to output_format, filter small bboxes and pad targets.

    :param input_dim:          Shape of the images to transform.
    :param input_format:       Format of the input targets. For instance [xmin, ymin, xmax, ymax, cls_id] refers to XYXY_LABEL.
    :param output_format:      Format of the output targets. For instance [xmin, ymin, xmax, ymax, cls_id] refers to XYXY_LABEL
    :param min_bbox_edge_size: bboxes with edge size lower then this values will be removed.
    """

    @resolve_param("input_format", ConcatenatedTensorFormatFactory())
    @resolve_param("output_format", ConcatenatedTensorFormatFactory())
    def __init__(
        self,
        input_dim: Union[int, Tuple[int, int], None] = None,
        input_format: ConcatenatedTensorFormat = XYXY_LABEL,
        output_format: ConcatenatedTensorFormat = LABEL_CXCYWH,
        min_bbox_edge_size: float = 1,
        max_targets: Optional[int] = None,
    ):
        super(DetectionTargetsFormatTransform, self).__init__()
        _max_targets_deprication(max_targets)
        if isinstance(input_format, DetectionTargetsFormat) or isinstance(output_format, DetectionTargetsFormat):
            raise TypeError(
                "DetectionTargetsFormat is not supported for input_format and output_format starting from super_gradients==3.0.7.\n"
                "You can either:\n"
                "\t - use builtin format among super_gradients.training.datasets.data_formats.default_formats.<FORMAT_NAME> (e.g. XYXY_LABEL, CXCY_LABEL, ..)\n"
                "\t - define your custom format using super_gradients.training.datasets.data_formats.formats.ConcatenatedTensorFormat\n"
            )
        self.input_format = input_format
        self.output_format = output_format
        self.min_bbox_edge_size = min_bbox_edge_size
        self.input_dim = None

        if input_dim is not None:
            input_dim = ensure_is_tuple_of_two(input_dim)
            self._setup_input_dim_related_params(input_dim)

    def _setup_input_dim_related_params(self, input_dim: tuple):
        """Setup all the parameters that are related to input_dim."""
        self.input_dim = input_dim
        self.min_bbox_edge_size = self.min_bbox_edge_size / max(input_dim) if self.output_format.bboxes_format.format.normalized else self.min_bbox_edge_size
        self.targets_format_converter = ConcatenatedTensorFormatConverter(
            input_format=self.input_format, output_format=self.output_format, image_shape=input_dim
        )

    def __call__(self, sample: dict) -> dict:

        # if self.input_dim not set yet, it will be set with first batch
        if self.input_dim is None:
            self._setup_input_dim_related_params(input_dim=sample["image"].shape[1:])

        sample["target"] = self.apply_on_targets(sample["target"])
        if "crowd_target" in sample.keys():
            sample["crowd_target"] = self.apply_on_targets(sample["crowd_target"])
        return sample

    def apply_on_targets(self, targets: np.ndarray) -> np.ndarray:
        """Convert targets in input_format to output_format, filter small bboxes and pad targets"""
        targets = self.targets_format_converter(targets)
        targets = self.filter_small_bboxes(targets)
        return np.ascontiguousarray(targets, dtype=np.float32)

    def filter_small_bboxes(self, targets: np.ndarray) -> np.ndarray:
        """Filter bboxes smaller than specified threshold."""

        def _is_big_enough(bboxes: np.ndarray) -> np.ndarray:
            return np.minimum(bboxes[:, 2], bboxes[:, 3]) > self.min_bbox_edge_size

        targets = filter_on_bboxes(fn=_is_big_enough, tensor=targets, tensor_format=self.output_format)
        return targets

    def get_equivalent_preprocessing(self) -> List:
        return []

apply_on_targets(targets)

Convert targets in input_format to output_format, filter small bboxes and pad targets

Source code in V3_4/src/super_gradients/training/transforms/transforms.py 
1102
1103
1104
1105
1106
def apply_on_targets(self, targets: np.ndarray) -> np.ndarray:
    """Convert targets in input_format to output_format, filter small bboxes and pad targets"""
    targets = self.targets_format_converter(targets)
    targets = self.filter_small_bboxes(targets)
    return np.ascontiguousarray(targets, dtype=np.float32)

filter_small_bboxes(targets)

Filter bboxes smaller than specified threshold.

Source code in V3_4/src/super_gradients/training/transforms/transforms.py 
1108
1109
1110
1111
1112
1113
1114
1115
def filter_small_bboxes(self, targets: np.ndarray) -> np.ndarray:
    """Filter bboxes smaller than specified threshold."""

    def _is_big_enough(bboxes: np.ndarray) -> np.ndarray:
        return np.minimum(bboxes[:, 2], bboxes[:, 3]) > self.min_bbox_edge_size

    targets = filter_on_bboxes(fn=_is_big_enough, tensor=targets, tensor_format=self.output_format)
    return targets

DetectionTransform

Detection transform base class.

Complex transforms that require extra data loading can use the the additional_samples_count attribute in a similar fashion to what's been done in COCODetectionDataset:

self._load_additional_inputs_for_transform(sample, transform)

after the above call, sample["additional_samples"] holds a list of additional inputs and targets.

sample = transform(sample)

Parameters:

Name Type Description Default
additional_samples_count int

Additional samples to be loaded.

 0
non_empty_targets bool

Whether the additional targets can have empty targets or not.

 False
Source code in V3_4/src/super_gradients/training/transforms/transforms.py 
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
class DetectionTransform:
    """
    Detection transform base class.

    Complex transforms that require extra data loading can use the the additional_samples_count attribute in a
     similar fashion to what's been done in COCODetectionDataset:

    self._load_additional_inputs_for_transform(sample, transform)

    # after the above call, sample["additional_samples"] holds a list of additional inputs and targets.

    sample = transform(sample)


    :param additional_samples_count:    Additional samples to be loaded.
    :param non_empty_targets:           Whether the additional targets can have empty targets or not.
    """

    def __init__(self, additional_samples_count: int = 0, non_empty_targets: bool = False):
        self.additional_samples_count = additional_samples_count
        self.non_empty_targets = non_empty_targets

    def __call__(self, sample: Union[dict, list]):
        raise NotImplementedError

    def __repr__(self):
        return self.__class__.__name__ + str(self.__dict__).replace("{", "(").replace("}", ")")

    def get_equivalent_preprocessing(self) -> List:
        raise NotImplementedError

DetectionVerticalFlip

Bases: DetectionTransform

Vertical Flip for Detection

Parameters:

Name Type Description Default
prob float

Probability of applying vertical flip

 required
Source code in V3_4/src/super_gradients/training/transforms/transforms.py 
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
@register_transform(Transforms.DetectionVerticalFlip)
class DetectionVerticalFlip(DetectionTransform):
    """
    Vertical Flip for Detection

    :param prob:        Probability of applying vertical flip
    """

    def __init__(self, prob: float, max_targets: Optional[int] = None):
        super(DetectionVerticalFlip, self).__init__()
        _max_targets_deprication(max_targets)
        self.prob = prob

    def __call__(self, sample):
        image, targets = sample["image"], sample["target"]
        crowd_targets = sample.get("crowd_targets")
        if len(targets) == 0:
            targets = np.zeros((0, 5), dtype=np.float32)
        if random.random() < self.prob:
            image = _flip_vertical_image(image)
            height, _, _ = image.shape
            targets[:, :4] = _flip_vertical_boxes(targets[:, :4], height)
            if crowd_targets is not None:
                crowd_targets[:, :4] = _flip_vertical_boxes(crowd_targets[:, :4], height)
        sample["image"] = image
        sample["target"] = targets
        sample["crowd_targets"] = crowd_targets
        return sample

SegCropImageAndMask

Bases: SegmentationTransform

Crops image and mask (synchronously). In "center" mode a center crop is performed while, in "random" mode the drop will be positioned around random coordinates.

Source code in V3_4/src/super_gradients/training/transforms/transforms.py 
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
@register_transform(Transforms.SegCropImageAndMask)
class SegCropImageAndMask(SegmentationTransform):
    """
    Crops image and mask (synchronously).
    In "center" mode a center crop is performed while, in "random" mode the drop will be positioned around
     random coordinates.
    """

    def __init__(self, crop_size: Union[float, Tuple, List], mode: str):
        """

        :param crop_size: tuple of (width, height) for the final crop size, if is scalar size is a
            square (crop_size, crop_size)
        :param mode: how to choose the center of the crop, 'center' for the center of the input image,
            'random' center the point is chosen randomally
        """

        self.crop_size = crop_size
        self.mode = mode

        self.check_valid_arguments()

    def __call__(self, sample: dict) -> dict:
        image = sample["image"]
        mask = sample["mask"]

        w, h = image.size
        if self.mode == "random":
            x1 = random.randint(0, w - self.crop_size[0])
            y1 = random.randint(0, h - self.crop_size[1])
        else:
            x1 = int(round((w - self.crop_size[0]) / 2.0))
            y1 = int(round((h - self.crop_size[1]) / 2.0))

        image = image.crop((x1, y1, x1 + self.crop_size[0], y1 + self.crop_size[1]))
        mask = mask.crop((x1, y1, x1 + self.crop_size[0], y1 + self.crop_size[1]))

        sample["image"] = image
        sample["mask"] = mask

        return sample

    def check_valid_arguments(self):
        if self.mode not in ["center", "random"]:
            raise ValueError(f"Unsupported mode: found: {self.mode}, expected: 'center' or 'random'")

        if not isinstance(self.crop_size, collections.abc.Iterable):
            self.crop_size = (self.crop_size, self.crop_size)
        if self.crop_size[0] <= 0 or self.crop_size[1] <= 0:
            raise ValueError(f"Crop size must be positive numbers, found: {self.crop_size}")

__init__(crop_size, mode)

Parameters:

Name Type Description Default
crop_size Union[float, Tuple, List]

tuple of (width, height) for the final crop size, if is scalar size is a square (crop_size, crop_size)

 required
mode str

how to choose the center of the crop, 'center' for the center of the input image, 'random' center the point is chosen randomally

 required
Source code in V3_4/src/super_gradients/training/transforms/transforms.py 
228
229
230
231
232
233
234
235
236
237
238
239
240
def __init__(self, crop_size: Union[float, Tuple, List], mode: str):
    """

    :param crop_size: tuple of (width, height) for the final crop size, if is scalar size is a
        square (crop_size, crop_size)
    :param mode: how to choose the center of the crop, 'center' for the center of the input image,
        'random' center the point is chosen randomally
    """

    self.crop_size = crop_size
    self.mode = mode

    self.check_valid_arguments()

SegPadShortToCropSize

Bases: SegmentationTransform

Pads image to 'crop_size'. Should be called only after "SegRescale" or "SegRandomRescale" in augmentations pipeline.

Source code in V3_4/src/super_gradients/training/transforms/transforms.py 
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
@register_transform(Transforms.SegPadShortToCropSize)
class SegPadShortToCropSize(SegmentationTransform):
    """
    Pads image to 'crop_size'.
    Should be called only after "SegRescale" or "SegRandomRescale" in augmentations pipeline.
    """

    def __init__(self, crop_size: Union[float, Tuple, List], fill_mask: int = 0, fill_image: Union[int, Tuple, List] = 0):
        """
        :param crop_size:   Tuple of (width, height) for the final crop size, if is scalar size is a square (crop_size, crop_size)
        :param fill_mask:   Value to fill mask labels background.
        :param fill_image:  Grey value to fill image padded background.
        """
        # CHECK IF CROP SIZE IS A ITERABLE OR SCALAR
        self.crop_size = crop_size
        self.fill_mask = fill_mask
        self.fill_image = tuple(fill_image) if isinstance(fill_image, Sequence) else fill_image

        self.check_valid_arguments()

    def __call__(self, sample: dict) -> dict:
        image = sample["image"]
        mask = sample["mask"]
        w, h = image.size

        # pad images from center symmetrically
        if w < self.crop_size[0] or h < self.crop_size[1]:
            padh = (self.crop_size[1] - h) / 2 if h < self.crop_size[1] else 0
            pad_top, pad_bottom = math.ceil(padh), math.floor(padh)
            padw = (self.crop_size[0] - w) / 2 if w < self.crop_size[0] else 0
            pad_left, pad_right = math.ceil(padw), math.floor(padw)

            image = ImageOps.expand(image, border=(pad_left, pad_top, pad_right, pad_bottom), fill=self.fill_image)
            mask = ImageOps.expand(mask, border=(pad_left, pad_top, pad_right, pad_bottom), fill=self.fill_mask)

        sample["image"] = image
        sample["mask"] = mask

        return sample

    def check_valid_arguments(self):
        if not isinstance(self.crop_size, collections.abc.Iterable):
            self.crop_size = (self.crop_size, self.crop_size)
        if self.crop_size[0] <= 0 or self.crop_size[1] <= 0:
            raise ValueError(f"Crop size must be positive numbers, found: {self.crop_size}")

        self.fill_mask, self.fill_image = _validate_fill_values_arguments(self.fill_mask, self.fill_image)

__init__(crop_size, fill_mask=0, fill_image=0)

Parameters:

Name Type Description Default
crop_size Union[float, Tuple, List]

Tuple of (width, height) for the final crop size, if is scalar size is a square (crop_size, crop_size)

 required
fill_mask int

Value to fill mask labels background.

 0
fill_image Union[int, Tuple, List]

Grey value to fill image padded background.

 0
Source code in V3_4/src/super_gradients/training/transforms/transforms.py 
302
303
304
305
306
307
308
309
310
311
312
313
def __init__(self, crop_size: Union[float, Tuple, List], fill_mask: int = 0, fill_image: Union[int, Tuple, List] = 0):
    """
    :param crop_size:   Tuple of (width, height) for the final crop size, if is scalar size is a square (crop_size, crop_size)
    :param fill_mask:   Value to fill mask labels background.
    :param fill_image:  Grey value to fill image padded background.
    """
    # CHECK IF CROP SIZE IS A ITERABLE OR SCALAR
    self.crop_size = crop_size
    self.fill_mask = fill_mask
    self.fill_image = tuple(fill_image) if isinstance(fill_image, Sequence) else fill_image

    self.check_valid_arguments()

SegRandomFlip

Bases: SegmentationTransform

Randomly flips the image and mask (synchronously) with probability 'prob'.

Source code in V3_4/src/super_gradients/training/transforms/transforms.py 
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
@register_transform(Transforms.SegRandomFlip)
class SegRandomFlip(SegmentationTransform):
    """
    Randomly flips the image and mask (synchronously) with probability 'prob'.
    """

    def __init__(self, prob: float = 0.5):
        assert 0.0 <= prob <= 1.0, f"Probability value must be between 0 and 1, found {prob}"
        self.prob = prob

    def __call__(self, sample: dict) -> dict:
        image = sample["image"]
        mask = sample["mask"]
        if random.random() < self.prob:
            image = image.transpose(Image.FLIP_LEFT_RIGHT)
            mask = mask.transpose(Image.FLIP_LEFT_RIGHT)
            sample["image"] = image
            sample["mask"] = mask

        return sample

SegRandomGaussianBlur

Bases: SegmentationTransform

Adds random Gaussian Blur to image with probability 'prob'.

Source code in V3_4/src/super_gradients/training/transforms/transforms.py 
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
@register_transform(Transforms.SegRandomGaussianBlur)
class SegRandomGaussianBlur(SegmentationTransform):
    """
    Adds random Gaussian Blur to image with probability 'prob'.
    """

    def __init__(self, prob: float = 0.5):
        assert 0.0 <= prob <= 1.0, "Probability value must be between 0 and 1"
        self.prob = prob

    def __call__(self, sample: dict) -> dict:
        image = sample["image"]
        mask = sample["mask"]

        if random.random() < self.prob:
            image = image.filter(ImageFilter.GaussianBlur(radius=random.random()))

        sample["image"] = image
        sample["mask"] = mask

        return sample

SegRandomRescale

Random rescale the image and mask (synchronously) while preserving aspect ratio. Scale factor is randomly picked between scales [min, max]

Parameters:

Name Type Description Default
scales Union[float, Tuple, List]

Scale range tuple (min, max), if scales is a float range will be defined as (1, scales) if scales > 1, otherwise (scales, 1). must be a positive number.

 (0.5, 2.0)
Source code in V3_4/src/super_gradients/training/transforms/transforms.py 
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
@register_transform(Transforms.SegRandomRescale)
class SegRandomRescale:
    """
    Random rescale the image and mask (synchronously) while preserving aspect ratio.
    Scale factor is randomly picked between scales [min, max]

    :param scales: Scale range tuple (min, max), if scales is a float range will be defined as (1, scales) if scales > 1,
            otherwise (scales, 1). must be a positive number.
    """

    def __init__(self, scales: Union[float, Tuple, List] = (0.5, 2.0)):
        self.scales = scales

        self.check_valid_arguments()

    def __call__(self, sample: dict) -> dict:
        image = sample["image"]
        mask = sample["mask"]
        w, h = image.size

        scale = random.uniform(self.scales[0], self.scales[1])
        out_size = int(scale * w), int(scale * h)

        image = image.resize(out_size, IMAGE_RESAMPLE_MODE)
        mask = mask.resize(out_size, MASK_RESAMPLE_MODE)

        sample["image"] = image
        sample["mask"] = mask

        return sample

    def check_valid_arguments(self):
        """
        Check the scale values are valid. if order is wrong, flip the order and return the right scale values.
        """
        if not isinstance(self.scales, collections.abc.Iterable):
            if self.scales <= 1:
                self.scales = (self.scales, 1)
            else:
                self.scales = (1, self.scales)

        if self.scales[0] < 0 or self.scales[1] < 0:
            raise ValueError(f"SegRandomRescale scale values must be positive numbers, found: {self.scales}")
        if self.scales[0] > self.scales[1]:
            self.scales = (self.scales[1], self.scales[0])
        return self.scales

check_valid_arguments()

Check the scale values are valid. if order is wrong, flip the order and return the right scale values.

Source code in V3_4/src/super_gradients/training/transforms/transforms.py 
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
def check_valid_arguments(self):
    """
    Check the scale values are valid. if order is wrong, flip the order and return the right scale values.
    """
    if not isinstance(self.scales, collections.abc.Iterable):
        if self.scales <= 1:
            self.scales = (self.scales, 1)
        else:
            self.scales = (1, self.scales)

    if self.scales[0] < 0 or self.scales[1] < 0:
        raise ValueError(f"SegRandomRescale scale values must be positive numbers, found: {self.scales}")
    if self.scales[0] > self.scales[1]:
        self.scales = (self.scales[1], self.scales[0])
    return self.scales

SegRandomRotate

Bases: SegmentationTransform

Randomly rotates image and mask (synchronously) between 'min_deg' and 'max_deg'.

Source code in V3_4/src/super_gradients/training/transforms/transforms.py 
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
@register_transform(Transforms.SegRandomRotate)
class SegRandomRotate(SegmentationTransform):
    """
    Randomly rotates image and mask (synchronously) between 'min_deg' and 'max_deg'.
    """

    def __init__(self, min_deg: float = -10, max_deg: float = 10, fill_mask: int = 0, fill_image: Union[int, Tuple, List] = 0):
        self.min_deg = min_deg
        self.max_deg = max_deg
        self.fill_mask = fill_mask
        # grey color in RGB mode
        self.fill_image = (fill_image, fill_image, fill_image)

        self.check_valid_arguments()

    def __call__(self, sample: dict) -> dict:
        image = sample["image"]
        mask = sample["mask"]

        deg = random.uniform(self.min_deg, self.max_deg)
        image = image.rotate(deg, resample=IMAGE_RESAMPLE_MODE, fillcolor=self.fill_image)
        mask = mask.rotate(deg, resample=MASK_RESAMPLE_MODE, fillcolor=self.fill_mask)

        sample["image"] = image
        sample["mask"] = mask

        return sample

    def check_valid_arguments(self):
        self.fill_mask, self.fill_image = _validate_fill_values_arguments(self.fill_mask, self.fill_image)

SegRescale

Bases: SegmentationTransform

Rescales the image and mask (synchronously) while preserving aspect ratio. The rescaling can be done according to scale_factor, short_size or long_size. If more than one argument is given, the rescaling mode is determined by this order: scale_factor, then short_size, then long_size.

Parameters:

Name Type Description Default
scale_factor Optional[float]

Rescaling is done by multiplying input size by scale_factor: out_size = (scale_factor * w, scale_factor * h)

 None
short_size Optional[int]

Rescaling is done by determining the scale factor by the ratio short_size / min(h, w).

 None
long_size Optional[int]

Rescaling is done by determining the scale factor by the ratio long_size / max(h, w).

 None
Source code in V3_4/src/super_gradients/training/transforms/transforms.py 
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
@register_transform(Transforms.SegRescale)
class SegRescale(SegmentationTransform):
    """
    Rescales the image and mask (synchronously) while preserving aspect ratio.
    The rescaling can be done according to scale_factor, short_size or long_size.
    If more than one argument is given, the rescaling mode is determined by this order: scale_factor, then short_size,
    then long_size.

    :param scale_factor: Rescaling is done by multiplying input size by scale_factor:
            out_size = (scale_factor * w, scale_factor * h)
    :param short_size:  Rescaling is done by determining the scale factor by the ratio short_size / min(h, w).
    :param long_size:   Rescaling is done by determining the scale factor by the ratio long_size / max(h, w).
    """

    def __init__(self, scale_factor: Optional[float] = None, short_size: Optional[int] = None, long_size: Optional[int] = None):
        self.scale_factor = scale_factor
        self.short_size = short_size
        self.long_size = long_size

        self.check_valid_arguments()

    def __call__(self, sample: dict) -> dict:
        image = sample["image"]
        mask = sample["mask"]
        w, h = image.size
        if self.scale_factor is not None:
            scale = self.scale_factor
        elif self.short_size is not None:
            short_size = min(w, h)
            scale = self.short_size / short_size
        else:
            long_size = max(w, h)
            scale = self.long_size / long_size

        out_size = int(scale * w), int(scale * h)

        image = image.resize(out_size, IMAGE_RESAMPLE_MODE)
        mask = mask.resize(out_size, MASK_RESAMPLE_MODE)

        sample["image"] = image
        sample["mask"] = mask

        return sample

    def check_valid_arguments(self):
        if self.scale_factor is None and self.short_size is None and self.long_size is None:
            raise ValueError("Must assign one rescale argument: scale_factor, short_size or long_size")

        if self.scale_factor is not None and self.scale_factor <= 0:
            raise ValueError(f"Scale factor must be a positive number, found: {self.scale_factor}")
        if self.short_size is not None and self.short_size <= 0:
            raise ValueError(f"Short size must be a positive number, found: {self.short_size}")
        if self.long_size is not None and self.long_size <= 0:
            raise ValueError(f"Long size must be a positive number, found: {self.long_size}")

Standardize

Bases: torch.nn.Module

Standardize image pixel values.

Returns:

Type Description

max_val: float, value to as described above (default=255)
Source code in V3_4/src/super_gradients/training/transforms/transforms.py 
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
@register_transform(Transforms.Standardize)
class Standardize(torch.nn.Module):
    """
    Standardize image pixel values.
    :return img/max_val

    attributes:
        max_val: float, value to as described above (default=255)
    """

    def __init__(self, max_val=255.0):
        super(Standardize, self).__init__()
        self.max_val = max_val

    def forward(self, img):
        return img / self.max_val

get_affine_matrix(input_size, target_size, degrees=10, translate=0.1, scales=0.1, shear=10)

Return a random affine transform matrix.

Parameters:

Name Type Description Default
input_size Tuple[int, int]

Input shape.

 required
target_size Tuple[int, int]

Desired output shape.

 required
degrees Union[tuple, float]

Degrees for random rotation, when float the random values are drawn uniformly from (-degrees, degrees)

 10
translate Union[tuple, float]

Translate size (in pixels) for random translation, when float the random values are drawn uniformly from (-translate, translate)

 0.1
scales Union[tuple, float]

Values for random rescale, when float the random values are drawn uniformly from (1-scales, 1+scales)

 0.1
shear Union[tuple, float]

Degrees for random shear, when float the random values are drawn uniformly from (-shear, shear)

 10

Returns:

Type Description
np.ndarray

affine_transform_matrix, drawn_scale
Source code in V3_4/src/super_gradients/training/transforms/transforms.py 
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
def get_affine_matrix(
    input_size: Tuple[int, int],
    target_size: Tuple[int, int],
    degrees: Union[tuple, float] = 10,
    translate: Union[tuple, float] = 0.1,
    scales: Union[tuple, float] = 0.1,
    shear: Union[tuple, float] = 10,
) -> np.ndarray:
    """
    Return a random affine transform matrix.

    :param input_size:      Input shape.
    :param target_size:     Desired output shape.
    :param degrees:         Degrees for random rotation, when float the random values are drawn uniformly from (-degrees, degrees)
    :param translate:       Translate size (in pixels) for random translation, when float the random values are drawn uniformly from (-translate, translate)
    :param scales:          Values for random rescale, when float the random values are drawn uniformly from (1-scales, 1+scales)
    :param shear:           Degrees for random shear, when float the random values are drawn uniformly from (-shear, shear)

    :return: affine_transform_matrix, drawn_scale
    """

    # Center in pixels
    center_m = np.eye(3)
    center = (input_size[0] // 2, input_size[1] // 2)
    center_m[0, 2] = -center[1]
    center_m[1, 2] = -center[0]

    # Rotation and scale
    rotation_m = np.eye(3)
    rotation_m[:2] = cv2.getRotationMatrix2D(angle=get_aug_params(degrees), center=(0, 0), scale=get_aug_params(scales, center=1.0))

    # Shear in degrees
    shear_m = np.eye(3)
    shear_m[0, 1] = math.tan(get_aug_params(shear) * math.pi / 180)
    shear_m[1, 0] = math.tan(get_aug_params(shear) * math.pi / 180)

    # Translation in pixels
    translation_m = np.eye(3)
    translation_m[0, 2] = get_aug_params(translate, center=0.5) * target_size[1]
    translation_m[1, 2] = get_aug_params(translate, center=0.5) * target_size[0]

    return (translation_m @ shear_m @ rotation_m @ center_m)[:2]

get_aug_params(value, center=0)

Generates a random value for augmentations as described below

Parameters:

Name Type Description Default
value Union[tuple, float]

Range of values for generation. Wen tuple-drawn uniformly between (value[0], value[1]), and (center - value, center + value) when float.

 required
center float

Center to subtract when value is float.

 0

Returns:

Type Description
float

Generated value
Source code in V3_4/src/super_gradients/training/transforms/transforms.py 
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
def get_aug_params(value: Union[tuple, float], center: float = 0) -> float:
    """
    Generates a random value for augmentations as described below

    :param value:       Range of values for generation. Wen tuple-drawn uniformly between (value[0], value[1]), and (center - value, center + value) when float.
    :param center:      Center to subtract when value is float.
    :return:            Generated value
    """
    if isinstance(value, Number):
        return random.uniform(center - float(value), center + float(value))
    elif len(value) == 2:
        return random.uniform(value[0], value[1])
    else:
        raise ValueError(
            "Affine params should be either a sequence containing two values\
                          or single float values. Got {}".format(
                value
            )
        )

random_affine(img, targets=(), targets_seg=None, target_size=(640, 640), degrees=10, translate=0.1, scales=0.1, shear=10, filter_box_candidates=False, wh_thr=2, ar_thr=20, area_thr=0.1, border_value=114)

Performs random affine transform to img, targets

Parameters:

Name Type Description Default
img np.ndarray

Input image of shape [h, w, c]

 required
targets np.ndarray

Input target

 ()
targets_seg np.ndarray

Targets derived from segmentation masks

 None
target_size tuple

Desired output shape

 (640, 640)
degrees Union[float, tuple]

Degrees for random rotation, when float the random values are drawn uniformly from (-degrees, degrees).

 10
translate Union[float, tuple]

Translate size (in pixels) for random translation, when float the random values are drawn uniformly from (-translate, translate)

 0.1
scales Union[float, tuple]

Values for random rescale, when float the random values are drawn uniformly from (0.1-scales, 0.1+scales)

 0.1
shear Union[float, tuple]

Degrees for random shear, when float the random values are drawn uniformly from (shear, shear)

 10
filter_box_candidates bool

whether to filter out transformed bboxes by edge size, area ratio, and aspect ratio.

 False
wh_thr

(float) edge size threshold when filter_box_candidates = True. Bounding oxes with edges smaller then this values will be filtered out. (default=2)

 2
ar_thr

(float) aspect ratio threshold filter_box_candidates = True. Bounding boxes with aspect ratio larger then this values will be filtered out. (default=20)

 20
area_thr

(float) threshold for area ratio between original image and the transformed one, when when filter_box_candidates = True. Bounding boxes with such ratio smaller then this value will be filtered out. (default=0.1)

 0.1
border_value

value for filling borders after applying transforms (default=114).

 114

Returns:

Type Description

Image and Target with applied random affine
Source code in V3_4/src/super_gradients/training/transforms/transforms.py 
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
def random_affine(
    img: np.ndarray,
    targets: np.ndarray = (),
    targets_seg: np.ndarray = None,
    target_size: tuple = (640, 640),
    degrees: Union[float, tuple] = 10,
    translate: Union[float, tuple] = 0.1,
    scales: Union[float, tuple] = 0.1,
    shear: Union[float, tuple] = 10,
    filter_box_candidates: bool = False,
    wh_thr=2,
    ar_thr=20,
    area_thr=0.1,
    border_value=114,
):
    """
    Performs random affine transform to img, targets
    :param img:         Input image of shape [h, w, c]
    :param targets:     Input target
    :param targets_seg: Targets derived from segmentation masks
    :param target_size: Desired output shape
    :param degrees:     Degrees for random rotation, when float the random values are drawn uniformly
                            from (-degrees, degrees).
    :param translate:   Translate size (in pixels) for random translation, when float the random values
                            are drawn uniformly from (-translate, translate)
    :param scales:      Values for random rescale, when float the random values are drawn uniformly
                            from (0.1-scales, 0.1+scales)
    :param shear:       Degrees for random shear, when float the random values are drawn uniformly
                                from (shear, shear)

    :param filter_box_candidates:    whether to filter out transformed bboxes by edge size, area ratio, and aspect ratio.
    :param wh_thr: (float) edge size threshold when filter_box_candidates = True. Bounding oxes with edges smaller
      then this values will be filtered out. (default=2)

    :param ar_thr: (float) aspect ratio threshold filter_box_candidates = True. Bounding boxes with aspect ratio larger
      then this values will be filtered out. (default=20)

    :param area_thr:(float) threshold for area ratio between original image and the transformed one, when when filter_box_candidates = True.
      Bounding boxes with such ratio smaller then this value will be filtered out. (default=0.1)

    :param border_value: value for filling borders after applying transforms (default=114).

    :return:            Image and Target with applied random affine
    """

    targets_seg = np.zeros((targets.shape[0], 0)) if targets_seg is None else targets_seg
    M = get_affine_matrix(img.shape[:2], target_size, degrees, translate, scales, shear)

    img = cv2.warpAffine(img, M, dsize=target_size, borderValue=(border_value, border_value, border_value))

    # Transform label coordinates
    if len(targets) > 0:
        targets_orig = targets.copy()
        targets = apply_affine_to_bboxes(targets, targets_seg, target_size, M)
        if filter_box_candidates:
            box_candidates_ids = _filter_box_candidates(targets_orig[:, :4], targets[:, :4], wh_thr=wh_thr, ar_thr=ar_thr, area_thr=area_thr)
            targets = targets[box_candidates_ids]
    return img, targets