Skip to content

Transforms

AbstractDepthEstimationTransform

Bases: abc.ABC

Base class for all transforms for depth estimation sample augmentation.

Source code in latest/src/super_gradients/training/transforms/depth_estimation/abstract_depth_estimation_transform.py
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
class AbstractDepthEstimationTransform(abc.ABC):
    """
    Base class for all transforms for depth estimation sample augmentation.
    """

    @abc.abstractmethod
    def __call__(self, sample: DepthEstimationSample) -> DepthEstimationSample:
        """
        Apply transformation to given depth 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()

__call__(sample) abstractmethod

Apply transformation to given depth 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 DepthEstimationSample

Input sample to transform.

required

Returns:

Type Description
DepthEstimationSample

Modified sample (It can be the same instance as input or a new object).

Source code in latest/src/super_gradients/training/transforms/depth_estimation/abstract_depth_estimation_transform.py
11
12
13
14
15
16
17
18
19
20
21
22
@abc.abstractmethod
def __call__(self, sample: DepthEstimationSample) -> DepthEstimationSample:
    """
    Apply transformation to given depth 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()

AbstractDetectionTransform

Bases: abc.ABC

Base class for all transforms for object detection sample augmentation.

Source code in latest/src/super_gradients/training/transforms/detection/abstract_detection_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
class AbstractDetectionTransform(abc.ABC):
    """
    Base class for all transforms for object detection sample augmentation.
    """

    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

    @abstractmethod
    def apply_to_sample(self, sample: DetectionSample) -> DetectionSample:
        """
        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

__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 latest/src/super_gradients/training/transforms/detection/abstract_detection_transform.py
15
16
17
18
19
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 DetectionSample

Input sample to transform.

required

Returns:

Type Description
DetectionSample

Modified sample (It can be the same instance as input or a new object).

Source code in latest/src/super_gradients/training/transforms/detection/abstract_detection_transform.py
21
22
23
24
25
26
27
28
29
30
31
32
@abstractmethod
def apply_to_sample(self, sample: DetectionSample) -> DetectionSample:
    """
    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

DetectionLongestMaxSize

Bases: AbstractDetectionTransform, LegacyDetectionTransformMixin

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

Source code in latest/src/super_gradients/training/transforms/detection/detection_longest_max_size.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
@register_transform(Transforms.DetectionLongestMaxSize)
class DetectionLongestMaxSize(AbstractDetectionTransform, LegacyDetectionTransformMixin):
    """
    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. Default: 1.0
        """
        super().__init__()
        self.max_height = int(max_height)
        self.max_width = int(max_width)
        self.interpolation = int(interpolation)
        self.prob = float(prob)

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

            sample = DetectionSample(
                image=self.apply_to_image(sample.image, scale, cv2.INTER_LINEAR),
                bboxes_xyxy=self.apply_to_bboxes(sample.bboxes_xyxy, scale),
                labels=sample.labels,
                is_crowd=sample.is_crowd,
                additional_samples=None,
            )

            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]}")

        return sample

    @classmethod
    def apply_to_image(cls, image: np.ndarray, scale: float, interpolation: int) -> np.ndarray:
        height, width = image.shape[:2]

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

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

    def get_equivalent_preprocessing(self) -> List:
        return [{Processings.DetectionLongestMaxSizeRescale: {"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. Default: 1.0

1.0
Source code in latest/src/super_gradients/training/transforms/detection/detection_longest_max_size.py
20
21
22
23
24
25
26
27
28
29
30
31
32
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. Default: 1.0
    """
    super().__init__()
    self.max_height = int(max_height)
    self.max_width = int(max_width)
    self.interpolation = int(interpolation)
    self.prob = float(prob)

DetectionPadIfNeeded

Bases: AbstractDetectionTransform, LegacyDetectionTransformMixin

Pad image and targets to ensure that resulting image size is not less than (min_width, min_height).

Source code in latest/src/super_gradients/training/transforms/detection/detection_pad_if_needed.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
@register_transform(Transforms.DetectionPadIfNeeded)
class DetectionPadIfNeeded(AbstractDetectionTransform, LegacyDetectionTransformMixin):
    """
    Pad image and targets to ensure that resulting image size is not less than (min_width, min_height).
    """

    def __init__(self, min_height: int, min_width: int, pad_value: int, padding_mode: str = "bottom_right"):
        """
        :param min_height:     Minimal height of the image.
        :param min_width:      Minimal width of the image.
        :param pad_value:      Padding value of image
        :param padding_mode:   Padding mode. Supported modes: 'bottom_right', 'center'.
        """
        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 = pad_value
        self.padding_mode = padding_mode

    def apply_to_sample(self, sample: DetectionSample) -> DetectionSample:
        """
        Apply transform to a single sample.
        :param sample: Input detection sample.
        :return:       Transformed detection sample.
        """
        height, width = sample.image.shape[:2]

        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}")

        padding_coordinates = PaddingCoordinates(top=pad_top, bottom=pad_bottom, left=pad_left, right=pad_right)

        return DetectionSample(
            image=_pad_image(sample.image, padding_coordinates, self.image_pad_value),
            bboxes_xyxy=_shift_bboxes_xyxy(sample.bboxes_xyxy, pad_left, pad_top),
            labels=sample.labels,
            is_crowd=sample.is_crowd,
            additional_samples=None,
        )

    def get_equivalent_preprocessing(self) -> List:
        if self.padding_mode == "bottom_right":
            return [{Processings.DetectionBottomRightPadding: {"output_shape": (self.min_height, self.min_width), "pad_value": self.image_pad_value}}]
        elif self.padding_mode == "center":
            return [{Processings.DetectionCenterPadding: {"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, pad_value, padding_mode='bottom_right')

Parameters:

Name Type Description Default
min_height int

Minimal height of the image.

required
min_width int

Minimal width of the image.

required
pad_value int

Padding value of image

required
padding_mode str

Padding mode. Supported modes: 'bottom_right', 'center'.

'bottom_right'
Source code in latest/src/super_gradients/training/transforms/detection/detection_pad_if_needed.py
17
18
19
20
21
22
23
24
25
26
27
28
29
30
def __init__(self, min_height: int, min_width: int, pad_value: int, padding_mode: str = "bottom_right"):
    """
    :param min_height:     Minimal height of the image.
    :param min_width:      Minimal width of the image.
    :param pad_value:      Padding value of image
    :param padding_mode:   Padding mode. Supported modes: 'bottom_right', 'center'.
    """
    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 = pad_value
    self.padding_mode = padding_mode

apply_to_sample(sample)

Apply transform to a single sample.

Parameters:

Name Type Description Default
sample DetectionSample

Input detection sample.

required

Returns:

Type Description
DetectionSample

Transformed detection sample.

Source code in latest/src/super_gradients/training/transforms/detection/detection_pad_if_needed.py
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
def apply_to_sample(self, sample: DetectionSample) -> DetectionSample:
    """
    Apply transform to a single sample.
    :param sample: Input detection sample.
    :return:       Transformed detection sample.
    """
    height, width = sample.image.shape[:2]

    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}")

    padding_coordinates = PaddingCoordinates(top=pad_top, bottom=pad_bottom, left=pad_left, right=pad_right)

    return DetectionSample(
        image=_pad_image(sample.image, padding_coordinates, self.image_pad_value),
        bboxes_xyxy=_shift_bboxes_xyxy(sample.bboxes_xyxy, pad_left, pad_top),
        labels=sample.labels,
        is_crowd=sample.is_crowd,
        additional_samples=None,
    )

LegacyDetectionTransformMixin

A mixin class to make legacy detection transforms compatible with new detection transforms that operate on DetectionSample.

Source code in latest/src/super_gradients/training/transforms/detection/legacy_detection_transform_mixin.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
class LegacyDetectionTransformMixin:
    """
    A mixin class to make legacy detection transforms compatible with new detection transforms that operate on DetectionSample.
    """

    def __call__(self, sample: Union["DetectionSample", Dict[str, Any]]) -> Union["DetectionSample", Dict[str, Any]]:
        """
        :param sample: Dict with following keys:
                        - image: numpy array of [H,W,C] or [C,H,W] format
                        - target: numpy array of [N,5] shape with bounding box of each instance (XYXY + LABEL)
                        - crowd_targets: numpy array of [N,5] shape with bounding box of each instance (XYXY + LABEL)
        """

        if isinstance(sample, DetectionSample):
            return self.apply_to_sample(sample)
        else:
            has_crowd_target = "crowd_target" in sample
            sample = self.convert_input_dict_to_detection_sample(sample)
            sample = self.apply_to_sample(sample)
            return self.convert_detection_sample_to_dict(sample, include_crowd_target=has_crowd_target)

    @classmethod
    def convert_input_dict_to_detection_sample(cls, sample_annotations: Dict[str, Union[np.ndarray, Any]]) -> DetectionSample:
        """
        Convert old-style detection sample dict to new DetectionSample dataclass.

        :param sample_annotations: Input dictionary with following keys:
                                    - image: numpy array of [H,W,C] or [C,H,W] format
                                    - target: numpy array of [N,5] shape with bounding box of each instance (XYXY + LABEL)
                                    - crowd_targets: numpy array of [N,5] shape with bounding box of each instance (XYXY + LABEL)
        :return: An instance of DetectionSample dataclass filled with data from input dictionary.
        """
        target = sample_annotations["target"]
        if len(target) == 0:
            target = np.zeros((0, 5), dtype=np.float32)

        bboxes_xyxy = target[:, 0:4].reshape(-1, 4)
        labels = target[:, 4]

        is_crowd = np.zeros_like(labels, dtype=bool)
        if "crowd_target" in sample_annotations:
            crowd_target = sample_annotations["crowd_target"]
            if len(crowd_target) == 0:
                crowd_target = np.zeros((0, 5), dtype=np.float32)

            crowd_bboxes_xyxy = crowd_target[:, 0:4].reshape(-1, 4)
            crowd_labels = crowd_target[:, 4]
            bboxes_xyxy = np.concatenate([bboxes_xyxy, crowd_bboxes_xyxy], axis=0)
            labels = np.concatenate([labels, crowd_labels], axis=0)
            is_crowd = np.concatenate([is_crowd, np.ones_like(crowd_labels, dtype=bool)], axis=0)

        return DetectionSample(
            image=sample_annotations["image"],
            bboxes_xyxy=bboxes_xyxy,
            labels=labels,
            is_crowd=is_crowd,
            additional_samples=None,
        )

    @classmethod
    def convert_detection_sample_to_dict(cls, detection_sample: DetectionSample, include_crowd_target: Union[bool, None]) -> Dict[str, Union[np.ndarray, Any]]:
        """
        Convert new DetectionSample dataclass to old-style detection sample dict.
        This is a reverse operation to convert_input_dict_to_detection_sample and used to make legacy transforms compatible with new detection transforms.
        :param detection_sample:     Input DetectionSample dataclass.
        :param include_crowd_target: A flag indicating whether to include crowd_target in output dictionary.
                                     Can be None - in this case crowd_target will be included only if crowd targets are present in input sample.
        :return:                     Output dictionary with following keys:
                                        - image: numpy array of [H,W,C] or [C,H,W] format
                                        - target: numpy array of [N,5] shape with bounding box of each instance (XYXY + LABEL)
                                        - crowd_targets: numpy array of [N,5] shape with bounding box of each instance (XYXY + LABEL)
        """
        image = detection_sample.image
        crowd_mask = detection_sample.is_crowd > 0
        crowd_labels = detection_sample.labels[crowd_mask]
        crowd_bboxes_xyxy = detection_sample.bboxes_xyxy[crowd_mask]
        crowd_target = np.concatenate([crowd_bboxes_xyxy, crowd_labels[..., None]], axis=-1)

        labels = detection_sample.labels[~crowd_mask]
        bboxes_xyxy = detection_sample.bboxes_xyxy[~crowd_mask]
        target = np.concatenate([bboxes_xyxy, labels[..., None]], axis=-1)

        sample = {
            "image": image,
            "target": target,
        }
        if include_crowd_target is None:
            include_crowd_target = crowd_mask.any()
        if include_crowd_target:
            sample["crowd_target"] = crowd_target
        return sample

__call__(sample)

Parameters:

Name Type Description Default
sample Union[DetectionSample, Dict[str, Any]]

Dict with following keys: - image: numpy array of [H,W,C] or [C,H,W] format - target: numpy array of [N,5] shape with bounding box of each instance (XYXY + LABEL) - crowd_targets: numpy array of [N,5] shape with bounding box of each instance (XYXY + LABEL)

required
Source code in latest/src/super_gradients/training/transforms/detection/legacy_detection_transform_mixin.py
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
def __call__(self, sample: Union["DetectionSample", Dict[str, Any]]) -> Union["DetectionSample", Dict[str, Any]]:
    """
    :param sample: Dict with following keys:
                    - image: numpy array of [H,W,C] or [C,H,W] format
                    - target: numpy array of [N,5] shape with bounding box of each instance (XYXY + LABEL)
                    - crowd_targets: numpy array of [N,5] shape with bounding box of each instance (XYXY + LABEL)
    """

    if isinstance(sample, DetectionSample):
        return self.apply_to_sample(sample)
    else:
        has_crowd_target = "crowd_target" in sample
        sample = self.convert_input_dict_to_detection_sample(sample)
        sample = self.apply_to_sample(sample)
        return self.convert_detection_sample_to_dict(sample, include_crowd_target=has_crowd_target)

convert_detection_sample_to_dict(detection_sample, include_crowd_target) classmethod

Convert new DetectionSample dataclass to old-style detection sample dict. This is a reverse operation to convert_input_dict_to_detection_sample and used to make legacy transforms compatible with new detection transforms.

Parameters:

Name Type Description Default
detection_sample DetectionSample

Input DetectionSample dataclass.

required
include_crowd_target Union[bool, None]

A flag indicating whether to include crowd_target in output dictionary. Can be None - in this case crowd_target will be included only if crowd targets are present in input sample.

required

Returns:

Type Description
Dict[str, Union[np.ndarray, Any]]

Output dictionary with following keys: - image: numpy array of [H,W,C] or [C,H,W] format - target: numpy array of [N,5] shape with bounding box of each instance (XYXY + LABEL) - crowd_targets: numpy array of [N,5] shape with bounding box of each instance (XYXY + LABEL)

Source code in latest/src/super_gradients/training/transforms/detection/legacy_detection_transform_mixin.py
 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
@classmethod
def convert_detection_sample_to_dict(cls, detection_sample: DetectionSample, include_crowd_target: Union[bool, None]) -> Dict[str, Union[np.ndarray, Any]]:
    """
    Convert new DetectionSample dataclass to old-style detection sample dict.
    This is a reverse operation to convert_input_dict_to_detection_sample and used to make legacy transforms compatible with new detection transforms.
    :param detection_sample:     Input DetectionSample dataclass.
    :param include_crowd_target: A flag indicating whether to include crowd_target in output dictionary.
                                 Can be None - in this case crowd_target will be included only if crowd targets are present in input sample.
    :return:                     Output dictionary with following keys:
                                    - image: numpy array of [H,W,C] or [C,H,W] format
                                    - target: numpy array of [N,5] shape with bounding box of each instance (XYXY + LABEL)
                                    - crowd_targets: numpy array of [N,5] shape with bounding box of each instance (XYXY + LABEL)
    """
    image = detection_sample.image
    crowd_mask = detection_sample.is_crowd > 0
    crowd_labels = detection_sample.labels[crowd_mask]
    crowd_bboxes_xyxy = detection_sample.bboxes_xyxy[crowd_mask]
    crowd_target = np.concatenate([crowd_bboxes_xyxy, crowd_labels[..., None]], axis=-1)

    labels = detection_sample.labels[~crowd_mask]
    bboxes_xyxy = detection_sample.bboxes_xyxy[~crowd_mask]
    target = np.concatenate([bboxes_xyxy, labels[..., None]], axis=-1)

    sample = {
        "image": image,
        "target": target,
    }
    if include_crowd_target is None:
        include_crowd_target = crowd_mask.any()
    if include_crowd_target:
        sample["crowd_target"] = crowd_target
    return sample

convert_input_dict_to_detection_sample(sample_annotations) classmethod

Convert old-style detection sample dict to new DetectionSample dataclass.

Parameters:

Name Type Description Default
sample_annotations Dict[str, Union[np.ndarray, Any]]

Input dictionary with following keys: - image: numpy array of [H,W,C] or [C,H,W] format - target: numpy array of [N,5] shape with bounding box of each instance (XYXY + LABEL) - crowd_targets: numpy array of [N,5] shape with bounding box of each instance (XYXY + LABEL)

required

Returns:

Type Description
DetectionSample

An instance of DetectionSample dataclass filled with data from input dictionary.

Source code in latest/src/super_gradients/training/transforms/detection/legacy_detection_transform_mixin.py
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
@classmethod
def convert_input_dict_to_detection_sample(cls, sample_annotations: Dict[str, Union[np.ndarray, Any]]) -> DetectionSample:
    """
    Convert old-style detection sample dict to new DetectionSample dataclass.

    :param sample_annotations: Input dictionary with following keys:
                                - image: numpy array of [H,W,C] or [C,H,W] format
                                - target: numpy array of [N,5] shape with bounding box of each instance (XYXY + LABEL)
                                - crowd_targets: numpy array of [N,5] shape with bounding box of each instance (XYXY + LABEL)
    :return: An instance of DetectionSample dataclass filled with data from input dictionary.
    """
    target = sample_annotations["target"]
    if len(target) == 0:
        target = np.zeros((0, 5), dtype=np.float32)

    bboxes_xyxy = target[:, 0:4].reshape(-1, 4)
    labels = target[:, 4]

    is_crowd = np.zeros_like(labels, dtype=bool)
    if "crowd_target" in sample_annotations:
        crowd_target = sample_annotations["crowd_target"]
        if len(crowd_target) == 0:
            crowd_target = np.zeros((0, 5), dtype=np.float32)

        crowd_bboxes_xyxy = crowd_target[:, 0:4].reshape(-1, 4)
        crowd_labels = crowd_target[:, 4]
        bboxes_xyxy = np.concatenate([bboxes_xyxy, crowd_bboxes_xyxy], axis=0)
        labels = np.concatenate([labels, crowd_labels], axis=0)
        is_crowd = np.concatenate([is_crowd, np.ones_like(crowd_labels, dtype=bool)], axis=0)

    return DetectionSample(
        image=sample_annotations["image"],
        bboxes_xyxy=bboxes_xyxy,
        labels=labels,
        is_crowd=is_crowd,
        additional_samples=None,
    )

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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 not hasattr(t, "additional_samples_count") or 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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 not hasattr(t, "additional_samples_count") or 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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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 latest/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