Transforms

AbstractDepthEstimationTransform

Bases: abc.ABC

Base class for all transforms for depth estimation sample augmentation.

Source code in src/super_gradients/training/transforms/depth_estimation/abstract_depth_estimation_transform.py

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 src/super_gradients/training/transforms/depth_estimation/abstract_depth_estimation_transform.py

@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 src/super_gradients/training/transforms/detection/abstract_detection_transform.py

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 src/super_gradients/training/transforms/detection/abstract_detection_transform.py

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 src/super_gradients/training/transforms/detection/abstract_detection_transform.py

@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 src/super_gradients/training/transforms/detection/detection_longest_max_size.py

@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 src/super_gradients/training/transforms/detection/detection_longest_max_size.py

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 src/super_gradients/training/transforms/detection/detection_pad_if_needed.py

@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 src/super_gradients/training/transforms/detection/detection_pad_if_needed.py

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 src/super_gradients/training/transforms/detection/detection_pad_if_needed.py

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 src/super_gradients/training/transforms/detection/legacy_detection_transform_mixin.py

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 src/super_gradients/training/transforms/detection/legacy_detection_transform_mixin.py

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 src/super_gradients/training/transforms/detection/legacy_detection_transform_mixin.py

@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 src/super_gradients/training/transforms/detection/legacy_detection_transform_mixin.py

@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 src/super_gradients/training/transforms/keypoints/abstract_keypoints_transform.py

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 src/super_gradients/training/transforms/keypoints/abstract_keypoints_transform.py

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 src/super_gradients/training/transforms/keypoints/abstract_keypoints_transform.py

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 src/super_gradients/training/transforms/keypoints/abstract_keypoints_transform.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_brightness_contrast.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_brightness_contrast.py

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 src/super_gradients/training/transforms/keypoints/keypoints_compose.py

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 src/super_gradients/training/transforms/keypoints/keypoints_compose.py

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 src/super_gradients/training/transforms/keypoints/keypoints_compose.py

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 src/super_gradients/training/transforms/keypoints/keypoints_compose.py

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 src/super_gradients/training/transforms/keypoints/keypoints_hsv.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_hsv.py

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 src/super_gradients/training/transforms/keypoints/keypoints_image_normalize.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_image_normalize.py

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 src/super_gradients/training/transforms/keypoints/keypoints_image_normalize.py

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 src/super_gradients/training/transforms/keypoints/keypoints_image_standardize.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_image_standardize.py

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 src/super_gradients/training/transforms/keypoints/keypoints_image_standardize.py

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 src/super_gradients/training/transforms/keypoints/keypoints_longest_max_size.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_longest_max_size.py

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 src/super_gradients/training/transforms/keypoints/keypoints_mixup.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_mixup.py

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 src/super_gradients/training/transforms/keypoints/keypoints_mixup.py

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 src/super_gradients/training/transforms/keypoints/keypoints_mosaic.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_mosaic.py

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 src/super_gradients/training/transforms/keypoints/keypoints_mosaic.py

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 src/super_gradients/training/transforms/keypoints/keypoints_pad_if_needed.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_pad_if_needed.py

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 src/super_gradients/training/transforms/keypoints/keypoints_random_affine.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_random_affine.py

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 src/super_gradients/training/transforms/keypoints/keypoints_random_affine.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_random_affine.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_random_affine.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_random_affine.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_random_affine.py

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 src/super_gradients/training/transforms/keypoints/keypoints_random_horisontal_flip.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_random_horisontal_flip.py

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 src/super_gradients/training/transforms/keypoints/keypoints_random_horisontal_flip.py

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 src/super_gradients/training/transforms/keypoints/keypoints_random_horisontal_flip.py

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 src/super_gradients/training/transforms/keypoints/keypoints_random_horisontal_flip.py

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 src/super_gradients/training/transforms/keypoints/keypoints_random_horisontal_flip.py

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 src/super_gradients/training/transforms/keypoints/keypoints_random_rotate90.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_random_rotate90.py

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 src/super_gradients/training/transforms/keypoints/keypoints_random_rotate90.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_random_rotate90.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_random_rotate90.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_random_rotate90.py

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 src/super_gradients/training/transforms/keypoints/keypoints_random_vertical_flip.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_random_vertical_flip.py

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 src/super_gradients/training/transforms/keypoints/keypoints_random_vertical_flip.py

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 src/super_gradients/training/transforms/keypoints/keypoints_random_vertical_flip.py

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 src/super_gradients/training/transforms/keypoints/keypoints_random_vertical_flip.py

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 src/super_gradients/training/transforms/keypoints/keypoints_random_vertical_flip.py

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 src/super_gradients/training/transforms/keypoints/keypoints_remove_small_objects.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_remove_small_objects.py

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 src/super_gradients/training/transforms/keypoints/keypoints_remove_small_objects.py

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 src/super_gradients/training/transforms/keypoints/keypoints_rescale.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_rescale.py

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 src/super_gradients/training/transforms/keypoints/keypoints_rescale.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_rescale.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_rescale.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_rescale.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_rescale.py

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 src/super_gradients/training/transforms/keypoints/abstract_keypoints_transform.py

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 src/super_gradients/training/transforms/keypoints/abstract_keypoints_transform.py

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 src/super_gradients/training/transforms/keypoints/abstract_keypoints_transform.py

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 src/super_gradients/training/transforms/keypoints/abstract_keypoints_transform.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_brightness_contrast.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_brightness_contrast.py

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 src/super_gradients/training/transforms/keypoints/keypoints_compose.py

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 src/super_gradients/training/transforms/keypoints/keypoints_compose.py

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 src/super_gradients/training/transforms/keypoints/keypoints_compose.py

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 src/super_gradients/training/transforms/keypoints/keypoints_compose.py

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 src/super_gradients/training/transforms/keypoints/keypoints_hsv.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_hsv.py

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 src/super_gradients/training/transforms/keypoints/keypoints_image_normalize.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_image_normalize.py

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 src/super_gradients/training/transforms/keypoints/keypoints_image_normalize.py

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 src/super_gradients/training/transforms/keypoints/keypoints_image_standardize.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_image_standardize.py

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 src/super_gradients/training/transforms/keypoints/keypoints_image_standardize.py

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 src/super_gradients/training/transforms/keypoints/keypoints_image_to_tensor.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_image_to_tensor.py

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 src/super_gradients/training/transforms/keypoints/keypoints_longest_max_size.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_longest_max_size.py

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 src/super_gradients/training/transforms/keypoints/keypoints_mixup.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_mixup.py

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 src/super_gradients/training/transforms/keypoints/keypoints_mixup.py

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 src/super_gradients/training/transforms/keypoints/keypoints_mosaic.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_mosaic.py

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 src/super_gradients/training/transforms/keypoints/keypoints_mosaic.py

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 src/super_gradients/training/transforms/keypoints/keypoints_pad_if_needed.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_pad_if_needed.py

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 src/super_gradients/training/transforms/keypoints/keypoints_random_affine.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_random_affine.py

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 src/super_gradients/training/transforms/keypoints/keypoints_random_affine.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_random_affine.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_random_affine.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_random_affine.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_random_affine.py

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 src/super_gradients/training/transforms/keypoints/keypoints_random_horisontal_flip.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_random_horisontal_flip.py

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 src/super_gradients/training/transforms/keypoints/keypoints_random_horisontal_flip.py

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 src/super_gradients/training/transforms/keypoints/keypoints_random_horisontal_flip.py

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 src/super_gradients/training/transforms/keypoints/keypoints_random_horisontal_flip.py

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 src/super_gradients/training/transforms/keypoints/keypoints_random_horisontal_flip.py

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 src/super_gradients/training/transforms/keypoints/keypoints_random_rotate90.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_random_rotate90.py

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 src/super_gradients/training/transforms/keypoints/keypoints_random_rotate90.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_random_rotate90.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_random_rotate90.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_random_rotate90.py

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 src/super_gradients/training/transforms/keypoints/keypoints_random_vertical_flip.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_random_vertical_flip.py

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 src/super_gradients/training/transforms/keypoints/keypoints_random_vertical_flip.py

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 src/super_gradients/training/transforms/keypoints/keypoints_random_vertical_flip.py

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 src/super_gradients/training/transforms/keypoints/keypoints_random_vertical_flip.py

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 src/super_gradients/training/transforms/keypoints/keypoints_random_vertical_flip.py

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 src/super_gradients/training/transforms/keypoints/keypoints_remove_small_objects.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_remove_small_objects.py

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 src/super_gradients/training/transforms/keypoints/keypoints_remove_small_objects.py

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 src/super_gradients/training/transforms/keypoints/keypoints_rescale.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_rescale.py

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 src/super_gradients/training/transforms/keypoints/keypoints_rescale.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_rescale.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_rescale.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_rescale.py

@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 src/super_gradients/training/transforms/keypoints/keypoints_rescale.py

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

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

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

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

    return sample

KeypointsReverseImageChannels

Bases: AbstractKeypointTransform

Randomly reverse channel order with given probability. Given an image with RGB channels, when applied with probability 1, it returns an image with BGR channels. With probability 0.5 there is 50/50 chance to return BGR or RGB image. It usually helps to improve model's ability to generalize under different color channels order.

Source code in src/super_gradients/training/transforms/keypoints/keypoints_reverse_image_channels.py

@register_transform()
class KeypointsReverseImageChannels(AbstractKeypointTransform):
    """
    Randomly reverse channel order with given probability.
    Given an image with RGB channels, when applied with probability 1, it returns an image with BGR channels.
    With probability 0.5 there is 50/50 chance to return BGR or RGB image.
    It usually helps to improve model's ability to generalize under different color channels order.
    """

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

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

    def apply_to_sample(self, sample: PoseEstimationSample) -> PoseEstimationSample:
        if random.random() < self.prob:
            sample.image = np.ascontiguousarray(sample.image[:, :, ::-1])
        return sample

    def get_equivalent_preprocessing(self) -> List:
        if self.prob < 1:
            raise RuntimeError("Cannot set preprocessing pipeline with randomness. Set prob to 1.")
        return [Processings.ReverseImageChannels]

__init__(prob)

Parameters:

Name	Type	Description	Default
prob	float	Probability to apply the transform.	required

Source code in src/super_gradients/training/transforms/keypoints/keypoints_reverse_image_channels.py

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

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

AbstractSegmentationTransform

Bases: abc.ABC

Base class for all transforms for object detection sample augmentation.

Source code in src/super_gradients/training/transforms/segmentation/abstract_segmentation_transform.py

class AbstractSegmentationTransform(abc.ABC):
    """
    Base class for all transforms for object detection sample augmentation.
    """

    @abstractmethod
    def apply_to_sample(self, sample: SegmentationSample) -> SegmentationSample:
        """
        Apply transformation to given segmentation 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

apply_to_sample(sample) abstractmethod

Apply transformation to given segmentation 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	SegmentationSample	Input sample to transform.	required

Returns:

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

Source code in src/super_gradients/training/transforms/segmentation/abstract_segmentation_transform.py

@abstractmethod
def apply_to_sample(self, sample: SegmentationSample) -> SegmentationSample:
    """
    Apply transformation to given segmentation 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

LegacySegmentationTransformMixin

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

Source code in src/super_gradients/training/transforms/segmentation/legacy_segmentation_transform_mixin.py

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

    def __call__(self, sample: Union["SegmentationSample", Dict[str, Any]]) -> Union["SegmentationSample", 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, SegmentationSample):
            return self.apply_to_sample(sample)
        else:
            sample, image_is_pil = self.convert_input_dict_to_segmentation_sample(sample)
            sample = self.apply_to_sample(sample)
            return self.convert_segmentation_sample_to_dict(sample, image_is_pil)

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

        :param sample_annotations: Input dictionary with following keys:
            image:              Associated image with sample, in [H,W,C] (or H,W for greyscale) format.
            mask:               Associated segmentation mask with sample, in [H,W]

        :return: A tuple of SegmentationSample and a boolean value indicating whether original input dict has images as PIL Image
                 An instance of SegmentationSample dataclass filled with data from input dictionary.
        """

        image_is_pil = isinstance(sample_annotations["image"], Image.Image)
        return SegmentationSample(image=sample_annotations["image"], mask=sample_annotations["mask"]), image_is_pil

    @classmethod
    def convert_segmentation_sample_to_dict(cls, sample: SegmentationSample, image_is_pil: bool) -> Dict[str, Union[np.ndarray, Any]]:
        """
        Convert new SegmentationSample 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 segmentation
        transforms.
        :param sample:       Transformed sample
        :param image_is_pil: A boolean value indicating whether original input dict has images as PIL Image
                             If True, output dict will also have images as PIL Image, otherwise as numpy array.
        """

        image = sample.image
        mask = sample.mask

        if image_is_pil:
            if isinstance(image, np.ndarray):
                image = Image.fromarray(image)
            if isinstance(mask, np.ndarray):
                mask = Image.fromarray(mask)
        return {"image": image, "mask": mask}

__call__(sample)

Parameters:

Name	Type	Description	Default
sample	Union[SegmentationSample, 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 src/super_gradients/training/transforms/segmentation/legacy_segmentation_transform_mixin.py

def __call__(self, sample: Union["SegmentationSample", Dict[str, Any]]) -> Union["SegmentationSample", 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, SegmentationSample):
        return self.apply_to_sample(sample)
    else:
        sample, image_is_pil = self.convert_input_dict_to_segmentation_sample(sample)
        sample = self.apply_to_sample(sample)
        return self.convert_segmentation_sample_to_dict(sample, image_is_pil)

convert_input_dict_to_segmentation_sample(sample_annotations) classmethod

Convert old-style segmentation 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: Associated image with sample, in [H,W,C] (or H,W for greyscale) format. mask: Associated segmentation mask with sample, in [H,W]	required

Returns:

Type	Description
Tuple[SegmentationSample, bool]	A tuple of SegmentationSample and a boolean value indicating whether original input dict has images as PIL Image An instance of SegmentationSample dataclass filled with data from input dictionary.

Source code in src/super_gradients/training/transforms/segmentation/legacy_segmentation_transform_mixin.py

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

    :param sample_annotations: Input dictionary with following keys:
        image:              Associated image with sample, in [H,W,C] (or H,W for greyscale) format.
        mask:               Associated segmentation mask with sample, in [H,W]

    :return: A tuple of SegmentationSample and a boolean value indicating whether original input dict has images as PIL Image
             An instance of SegmentationSample dataclass filled with data from input dictionary.
    """

    image_is_pil = isinstance(sample_annotations["image"], Image.Image)
    return SegmentationSample(image=sample_annotations["image"], mask=sample_annotations["mask"]), image_is_pil

convert_segmentation_sample_to_dict(sample, image_is_pil) classmethod

Convert new SegmentationSample 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 segmentation transforms.

Parameters:

Name	Type	Description	Default
sample	SegmentationSample	Transformed sample	required
image_is_pil	bool	A boolean value indicating whether original input dict has images as PIL Image If True, output dict will also have images as PIL Image, otherwise as numpy array.	required

Source code in src/super_gradients/training/transforms/segmentation/legacy_segmentation_transform_mixin.py

@classmethod
def convert_segmentation_sample_to_dict(cls, sample: SegmentationSample, image_is_pil: bool) -> Dict[str, Union[np.ndarray, Any]]:
    """
    Convert new SegmentationSample 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 segmentation
    transforms.
    :param sample:       Transformed sample
    :param image_is_pil: A boolean value indicating whether original input dict has images as PIL Image
                         If True, output dict will also have images as PIL Image, otherwise as numpy array.
    """

    image = sample.image
    mask = sample.mask

    if image_is_pil:
        if isinstance(image, np.ndarray):
            image = Image.fromarray(image)
        if isinstance(mask, np.ndarray):
            mask = Image.fromarray(mask)
    return {"image": image, "mask": mask}

DetectionHSV

Bases: AbstractDetectionTransform, LegacyDetectionTransformMixin

Detection HSV transform.

Parameters:

Name	Type	Description	Default
prob	float	Probability to apply the transform.	required
hgain	float	Hue gain.	0.5
sgain	float	Saturation gain.	0.5
vgain	float	Value gain.	0.5
bgr_channels		Channel indices of the BGR channels- useful for images with >3 channels, or when BGR channels are in different order.	(0, 1, 2)

Source code in src/super_gradients/training/transforms/transforms.py

@register_transform(Transforms.DetectionHSV)
class DetectionHSV(AbstractDetectionTransform, LegacyDetectionTransformMixin):
    """
    Detection HSV transform.

    :param prob:            Probability to apply the transform.
    :param hgain:           Hue gain.
    :param sgain:           Saturation gain.
    :param vgain:           Value gain.
    :param bgr_channels:    Channel indices of the BGR channels- useful for images with >3 channels, or when BGR channels are in different order.
    """

    def __init__(self, prob: float, hgain: float = 0.5, sgain: float = 0.5, vgain: float = 0.5, bgr_channels=(0, 1, 2)):
        super(DetectionHSV, self).__init__()
        self.prob = prob
        self.hgain = hgain
        self.sgain = sgain
        self.vgain = vgain
        self.bgr_channels = bgr_channels
        self._additional_channels_warned = False

    def apply_to_sample(self, sample: DetectionSample) -> DetectionSample:
        if sample.image.shape[2] < 3:
            raise ValueError("HSV transform expects at least 3 channels, got: " + str(sample.image.shape[2]))
        if sample.image.shape[2] > 3 and not self._additional_channels_warned:
            logger.warning(
                "HSV transform received image with "
                + str(sample.image.shape[2])
                + " channels. HSV transform will only be applied on channels: "
                + str(self.bgr_channels)
                + "."
            )
            self._additional_channels_warned = True

        if random.random() < self.prob:
            sample = DetectionSample(
                image=self.apply_to_image(sample.image),
                bboxes_xyxy=sample.bboxes_xyxy,
                labels=sample.labels,
                is_crowd=sample.is_crowd,
                additional_samples=None,
            )
        return sample

    def apply_to_image(self, image: np.ndarray) -> np.ndarray:
        return augment_hsv(image.copy(), self.hgain, self.sgain, self.vgain, self.bgr_channels)

    def get_equivalent_preprocessing(self) -> List[Dict]:
        raise NotImplementedError("get_equivalent_preprocessing is not implemented for non-deterministic transforms.")

DetectionHorizontalFlip

Bases: AbstractDetectionTransform, LegacyDetectionTransformMixin

Horizontal Flip for Detection

Parameters:

Name	Type	Description	Default
prob	float	Probability of applying horizontal flip	required

Source code in src/super_gradients/training/transforms/transforms.py

@register_transform(Transforms.DetectionHorizontalFlip)
class DetectionHorizontalFlip(AbstractDetectionTransform, LegacyDetectionTransformMixin):
    """
    Horizontal Flip for Detection

    :param prob:        Probability of applying horizontal flip
    """

    def __init__(self, prob: float):
        super(DetectionHorizontalFlip, self).__init__()
        self.prob = float(prob)

    def apply_to_sample(self, sample: DetectionSample) -> DetectionSample:
        """
        Apply horizontal flip to sample
        :param sample: Input detection sample
        :return:       Transformed detection sample
        """
        if random.random() < self.prob:
            sample = DetectionSample(
                image=_flip_horizontal_image(sample.image),
                bboxes_xyxy=_flip_horizontal_boxes_xyxy(sample.bboxes_xyxy, sample.image.shape[1]),
                labels=sample.labels,
                is_crowd=sample.is_crowd,
                additional_samples=None,
            )
        return sample

    def get_equivalent_preprocessing(self) -> List[Dict]:
        raise NotImplementedError("get_equivalent_preprocessing is not implemented for non-deterministic transforms.")

apply_to_sample(sample)

Apply horizontal flip to sample

Parameters:

Name	Type	Description	Default
sample	DetectionSample	Input detection sample	required

Returns:

Type	Description
DetectionSample	Transformed detection sample

Source code in src/super_gradients/training/transforms/transforms.py

def apply_to_sample(self, sample: DetectionSample) -> DetectionSample:
    """
    Apply horizontal flip to sample
    :param sample: Input detection sample
    :return:       Transformed detection sample
    """
    if random.random() < self.prob:
        sample = DetectionSample(
            image=_flip_horizontal_image(sample.image),
            bboxes_xyxy=_flip_horizontal_boxes_xyxy(sample.bboxes_xyxy, sample.image.shape[1]),
            labels=sample.labels,
            is_crowd=sample.is_crowd,
            additional_samples=None,
        )
    return sample

DetectionImagePermute

Bases: AbstractDetectionTransform, LegacyDetectionTransformMixin

Permute image dims. Useful for converting image from HWC to CHW format.

Source code in src/super_gradients/training/transforms/transforms.py

@register_transform(Transforms.DetectionImagePermute)
class DetectionImagePermute(AbstractDetectionTransform, LegacyDetectionTransformMixin):
    """
    Permute image dims. Useful for converting image from HWC to CHW format.
    """

    def __init__(self, dims: Tuple[int, int, int] = (2, 0, 1)):
        """

        :param dims: Specify new order of dims. Default value (2, 0, 1) suitable for converting from HWC to CHW format.
        """
        super().__init__()
        self.dims = tuple(dims)

    def apply_to_sample(self, sample: DetectionSample) -> DetectionSample:
        sample.image = np.ascontiguousarray(sample.image.transpose(*self.dims))
        return sample

    def get_equivalent_preprocessing(self) -> List[Dict]:
        return [{Processings.ImagePermute: {"permutation": self.dims}}]

__init__(dims=(2, 0, 1))

Parameters:

Name	Type	Description	Default
dims	Tuple[int, int, int]	Specify new order of dims. Default value (2, 0, 1) suitable for converting from HWC to CHW format.	(2, 0, 1)

Source code in src/super_gradients/training/transforms/transforms.py

def __init__(self, dims: Tuple[int, int, int] = (2, 0, 1)):
    """

    :param dims: Specify new order of dims. Default value (2, 0, 1) suitable for converting from HWC to CHW format.
    """
    super().__init__()
    self.dims = tuple(dims)

DetectionMixup

Bases: AbstractDetectionTransform, LegacyDetectionTransformMixin

Mixup detection transform

Parameters:

Name	Type	Description	Default
input_dim	Union[int, Tuple[int, int], None]	Input dimension.	required
mixup_scale	tuple	Scale range for the additional loaded image for mixup.	required
prob	float	Probability of applying mixup.	1.0
enable_mixup	bool	Whether to apply mixup at all (regardless of prob).	True
flip_prob	float	Probability to apply horizontal flip to the additional sample.	0.5
border_value	int	Value for filling borders after applying transform.	114

Source code in src/super_gradients/training/transforms/transforms.py

@register_transform(Transforms.DetectionMixup)
class DetectionMixup(AbstractDetectionTransform, LegacyDetectionTransformMixin):
    """
    Mixup detection transform

    :param input_dim:        Input dimension.
    :param mixup_scale:      Scale range for the additional loaded image for mixup.
    :param prob:             Probability of applying mixup.
    :param enable_mixup:     Whether to apply mixup at all (regardless of prob).
    :param flip_prob:        Probability to apply horizontal flip to the additional sample.
    :param border_value:     Value for filling borders after applying transform.
    """

    def __init__(
        self,
        input_dim: Union[int, Tuple[int, int], None],
        mixup_scale: tuple,
        prob: float = 1.0,
        enable_mixup: bool = True,
        flip_prob: float = 0.5,
        border_value: int = 114,
    ):
        super(DetectionMixup, self).__init__(additional_samples_count=1)
        self.input_dim = ensure_is_tuple_of_two(input_dim)
        self.mixup_scale = mixup_scale
        self.prob = prob
        self.enable_mixup = enable_mixup
        self.flip_prob = flip_prob
        self.border_value = border_value
        self.non_empty_targets = True
        self.maybe_flip = DetectionHorizontalFlip(prob=flip_prob)

    def close(self):
        self.additional_samples_count = 0
        self.enable_mixup = False

    def apply_to_sample(self, sample: DetectionSample) -> DetectionSample:
        if self.enable_mixup and random.random() < self.prob:
            (cp_sample,) = sample.additional_samples
            target_dim = self.input_dim if self.input_dim is not None else sample.image.shape[:2]

            cp_sample = self.maybe_flip.apply_to_sample(cp_sample)

            jit_factor = random.uniform(*self.mixup_scale)

            if len(sample.image.shape) == 3:
                cp_img = np.ones((target_dim[0], target_dim[1], sample.image.shape[2]), dtype=np.uint8) * self.border_value
            else:
                cp_img = np.ones(target_dim, dtype=np.uint8) * self.border_value

            cp_scale_ratio = min(target_dim[0] / cp_sample.image.shape[0], target_dim[1] / cp_sample.image.shape[1])
            resized_img = cv2.resize(
                cp_sample.image,
                (int(cp_sample.image.shape[1] * cp_scale_ratio), int(cp_sample.image.shape[0] * cp_scale_ratio)),
                interpolation=cv2.INTER_LINEAR,
            )

            cp_img[: resized_img.shape[0], : resized_img.shape[1]] = resized_img

            cp_img = cv2.resize(
                cp_img,
                (int(cp_img.shape[1] * jit_factor), int(cp_img.shape[0] * jit_factor)),
                interpolation=cv2.INTER_LINEAR,
            )
            cp_scale_ratio *= jit_factor

            origin_h, origin_w = cp_img.shape[:2]
            target_h, target_w = sample.image.shape[:2]

            if len(cp_img.shape) == 3:
                padded_img = np.zeros((max(origin_h, target_h), max(origin_w, target_w), cp_img.shape[2]), dtype=np.uint8)
            else:
                padded_img = np.zeros((max(origin_h, target_h), max(origin_w, target_w)), dtype=np.uint8)

            padded_img[:origin_h, :origin_w] = cp_img

            x_offset, y_offset = 0, 0
            if padded_img.shape[0] > target_h:
                y_offset = random.randint(0, padded_img.shape[0] - target_h - 1)
            if padded_img.shape[1] > target_w:
                x_offset = random.randint(0, padded_img.shape[1] - target_w - 1)
            padded_cropped_img = padded_img[y_offset : y_offset + target_h, x_offset : x_offset + target_w]

            cp_bboxes_origin_np = adjust_box_anns(cp_sample.bboxes_xyxy[:, :4].copy(), cp_scale_ratio, 0, 0, origin_w, origin_h)
            cp_bboxes_transformed_np = cp_bboxes_origin_np.copy()
            cp_bboxes_transformed_np[:, [0, 2]] = cp_bboxes_transformed_np[:, [0, 2]] - x_offset
            cp_bboxes_transformed_np[:, [1, 3]] = cp_bboxes_transformed_np[:, [1, 3]] - y_offset
            cp_bboxes_transformed_np = change_bbox_bounds_for_image_size_inplace(cp_bboxes_transformed_np, (target_h, target_w))

            mixup_boxes = np.concatenate([sample.bboxes_xyxy, cp_bboxes_transformed_np], axis=0)
            mixup_labels = np.concatenate([sample.labels, cp_sample.labels], axis=0)
            mixup_crowds = np.concatenate([sample.is_crowd, cp_sample.is_crowd], axis=0)

            mixup_image = (0.5 * sample.image + 0.5 * padded_cropped_img).astype(sample.image.dtype)
            sample = DetectionSample(
                image=mixup_image,
                bboxes_xyxy=mixup_boxes,
                labels=mixup_labels,
                is_crowd=mixup_crowds,
                additional_samples=None,
            )

        return sample

    def get_equivalent_preprocessing(self):
        raise NotImplementedError("get_equivalent_preprocessing is not implemented for non-deterministic transforms.")

DetectionMosaic

Bases: AbstractDetectionTransform, LegacyDetectionTransformMixin

DetectionMosaic detection transform

Parameters:

Name	Type	Description	Default
input_dim	Union[int, Tuple[int, int]]	Input dimension.	required
prob	float	Probability of applying mosaic.	1.0
enable_mosaic	bool	Whether to apply mosaic at all (regardless of prob).	True
border_value		Value for filling borders after applying transforms.	114

Source code in src/super_gradients/training/transforms/transforms.py

@register_transform(Transforms.DetectionMosaic)
class DetectionMosaic(AbstractDetectionTransform, LegacyDetectionTransformMixin):
    """
    DetectionMosaic detection transform

    :param input_dim:       Input dimension.
    :param prob:            Probability of applying mosaic.
    :param enable_mosaic:   Whether to apply mosaic at all (regardless of prob).
    :param border_value:    Value for filling borders after applying transforms.
    """

    def __init__(self, input_dim: Union[int, Tuple[int, int]], prob: float = 1.0, enable_mosaic: bool = True, border_value=114):
        super(DetectionMosaic, self).__init__(additional_samples_count=3)
        self.prob = prob
        self.input_dim = ensure_is_tuple_of_two(input_dim)
        self.enable_mosaic = enable_mosaic
        self.border_value = border_value

    def close(self):
        self.additional_samples_count = 0
        self.enable_mosaic = False

    def apply_to_sample(self, sample: DetectionSample) -> DetectionSample:
        if self.enable_mosaic and random.random() < self.prob:
            mosaic_labels = []
            mosaic_bboxes = []
            mosaic_iscrowd = []

            input_h, input_w = self.input_dim[0], self.input_dim[1]

            # yc, xc = s, s  # mosaic center x, y
            yc = int(random.uniform(0.5 * input_h, 1.5 * input_h))
            xc = int(random.uniform(0.5 * input_w, 1.5 * input_w))

            # 3 additional samples, total of 4
            all_samples: List[DetectionSample] = [sample] + sample.additional_samples

            for i_mosaic, sample in enumerate(all_samples):
                img = sample.image

                h0, w0 = img.shape[:2]  # orig hw
                scale = min(1.0 * input_h / h0, 1.0 * input_w / w0)
                img = cv2.resize(img, (int(w0 * scale), int(h0 * scale)), interpolation=cv2.INTER_LINEAR)
                # generate output mosaic image
                (h, w, c) = img.shape[:3]
                if i_mosaic == 0:
                    mosaic_img = np.full((input_h * 2, input_w * 2, c), self.border_value, dtype=np.uint8)

                # suffix l means large image, while s means small image in mosaic aug.
                (l_x1, l_y1, l_x2, l_y2), (s_x1, s_y1, s_x2, s_y2) = get_mosaic_coordinate(i_mosaic, xc, yc, w, h, input_h, input_w)

                mosaic_img[l_y1:l_y2, l_x1:l_x2] = img[s_y1:s_y2, s_x1:s_x2]
                padw, padh = l_x1 - s_x1, l_y1 - s_y1

                bboxes = sample.bboxes_xyxy * scale + np.array([[padw, padh, padw, padh]], dtype=np.float32)

                mosaic_labels.append(sample.labels)
                mosaic_iscrowd.append(sample.is_crowd)
                mosaic_bboxes.append(bboxes)

            mosaic_iscrowd = np.concatenate(mosaic_iscrowd, 0)
            mosaic_labels = np.concatenate(mosaic_labels, 0)
            mosaic_bboxes = np.concatenate(mosaic_bboxes, 0)

            # No need to adjust bboxes for image size since DetectionSample constructor will do this anyway
            sample = DetectionSample(
                image=mosaic_img,
                bboxes_xyxy=mosaic_bboxes,
                labels=mosaic_labels,
                is_crowd=mosaic_iscrowd,
                additional_samples=None,
            )

        return sample

    def get_equivalent_preprocessing(self):
        raise NotImplementedError("get_equivalent_preprocessing is not implemented for non-deterministic transforms.")

DetectionNormalize

Bases: AbstractDetectionTransform, LegacyDetectionTransformMixin

Normalize image by subtracting mean and dividing by std.

Source code in src/super_gradients/training/transforms/transforms.py

@register_transform(Transforms.DetectionNormalize)
class DetectionNormalize(AbstractDetectionTransform, LegacyDetectionTransformMixin):
    """
    Normalize image by subtracting mean and dividing by std.
    """

    def __init__(self, mean, std):
        super().__init__()
        self.mean = np.array(list(mean)).reshape((1, 1, -1)).astype(np.float32)
        self.std = np.array(list(std)).reshape((1, 1, -1)).astype(np.float32)

    def apply_to_sample(self, sample: DetectionSample) -> DetectionSample:
        return DetectionSample(
            image=self.apply_to_image(sample.image),
            bboxes_xyxy=sample.bboxes_xyxy,
            labels=sample.labels,
            is_crowd=sample.is_crowd,
            additional_samples=None,
        )

    def apply_to_image(self, image: np.ndarray) -> np.ndarray:
        return (image - self.mean) / self.std

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

DetectionPadToSize

Bases: DetectionPadIfNeeded

Preprocessing transform to pad image and bboxes to input_dim shape (rows, cols). Transform does center padding, so that input image with bboxes located in the center of the produced image.

Note: This transformation assume that dimensions of input image is equal or less than output_size. This class exists for backward compatibility with previous versions of the library. Use DetectionPadIfNeeded instead.

Source code in src/super_gradients/training/transforms/transforms.py

@register_transform(Transforms.DetectionPadToSize)
class DetectionPadToSize(DetectionPadIfNeeded):
    """
    Preprocessing transform to pad image and bboxes to `input_dim` shape (rows, cols).
    Transform does center padding, so that input image with bboxes located in the center of the produced image.

    Note: This transformation assume that dimensions of input image is equal or less than `output_size`.
    This class exists for backward compatibility with previous versions of the library.
    Use `DetectionPadIfNeeded` instead.
    """

    def __init__(self, output_size: Union[int, Tuple[int, int], None], pad_value: Union[int, Tuple[int, ...]]):
        """
        Constructor for DetectionPadToSize transform.

        :param output_size: Output image size (rows, cols)
        :param pad_value: Padding value for image
        """
        min_height, min_width = ensure_is_tuple_of_two(output_size)

        super().__init__(min_height=min_height, min_width=min_width, pad_value=pad_value, padding_mode="center")
        self.output_size = ensure_is_tuple_of_two(output_size)
        self.pad_value = pad_value

__init__(output_size, pad_value)

Constructor for DetectionPadToSize transform.

Parameters:

Name	Type	Description	Default
output_size	Union[int, Tuple[int, int], None]	Output image size (rows, cols)	required
pad_value	Union[int, Tuple[int, ...]]	Padding value for image	required

Source code in src/super_gradients/training/transforms/transforms.py

def __init__(self, output_size: Union[int, Tuple[int, int], None], pad_value: Union[int, Tuple[int, ...]]):
    """
    Constructor for DetectionPadToSize transform.

    :param output_size: Output image size (rows, cols)
    :param pad_value: Padding value for image
    """
    min_height, min_width = ensure_is_tuple_of_two(output_size)

    super().__init__(min_height=min_height, min_width=min_width, pad_value=pad_value, padding_mode="center")
    self.output_size = ensure_is_tuple_of_two(output_size)
    self.pad_value = pad_value

DetectionPaddedRescale

Bases: AbstractDetectionTransform, LegacyDetectionTransformMixin

Preprocessing transform to be applied last of all transforms for validation.

Image- Rescales and pads to self.input_dim. Targets- moves the class label to first index, converts boxes format- xyxy -> cxcywh.

Parameters:

Name	Type	Description	Default
input_dim	Union[int, Tuple[int, int], None]	Final input dimension (default=(640,640))	required
swap	Tuple[int, ...]	Image axis's to be rearranged.	(2, 0, 1)
pad_value	int	Padding value for image.	114

Source code in src/super_gradients/training/transforms/transforms.py

@register_transform(Transforms.DetectionPaddedRescale)
class DetectionPaddedRescale(AbstractDetectionTransform, LegacyDetectionTransformMixin):
    """
    Preprocessing transform to be applied last of all transforms for validation.

    Image- Rescales and pads to self.input_dim.
    Targets- moves the class label to first index, converts boxes format- xyxy -> cxcywh.

    :param input_dim:   Final input dimension (default=(640,640))
    :param swap:        Image axis's to be rearranged.
    :param pad_value:   Padding value for image.
    """

    def __init__(
        self, input_dim: Union[int, Tuple[int, int], None], swap: Tuple[int, ...] = (2, 0, 1), max_targets: Optional[int] = None, pad_value: int = 114
    ):
        super().__init__()
        _max_targets_deprication(max_targets)
        self.swap = swap
        self.input_dim = ensure_is_tuple_of_two(input_dim)
        self.pad_value = pad_value

    def apply_to_sample(self, sample: DetectionSample) -> DetectionSample:
        sample.image, r = _rescale_and_pad_to_size(sample.image, self.input_dim, self.swap, self.pad_value)
        sample.bboxes_xyxy = _rescale_xyxy_bboxes(sample.bboxes_xyxy, r)
        return sample

    def get_equivalent_preprocessing(self) -> List[Dict]:
        return [
            {Processings.DetectionLongestMaxSizeRescale: {"output_shape": self.input_dim}},
            {Processings.DetectionBottomRightPadding: {"output_shape": self.input_dim, "pad_value": self.pad_value}},
            {Processings.ImagePermute: {"permutation": self.swap}},
        ]

DetectionRGB2BGR

Bases: AbstractDetectionTransform, LegacyDetectionTransformMixin

Detection change Red & Blue channel of the image

Parameters:

Name	Type	Description	Default
prob	float	Probability to apply the transform.	0.5

Source code in src/super_gradients/training/transforms/transforms.py

@register_transform(Transforms.DetectionRGB2BGR)
class DetectionRGB2BGR(AbstractDetectionTransform, LegacyDetectionTransformMixin):
    """
    Detection change Red & Blue channel of the image

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

    def __init__(self, prob: float = 0.5):
        super().__init__()
        self.prob = float(prob)

    def apply_to_sample(self, sample: DetectionSample) -> DetectionSample:
        if len(sample.image.shape) != 3 or sample.image.shape[2] < 3:
            raise ValueError("DetectionRGB2BGR transform expects image to have 3 channels, got input image shape: " + str(sample.image.shape))
        if random.random() < self.prob:
            sample = DetectionSample(
                image=np.ascontiguousarray(sample.image[..., ::-1]),
                bboxes_xyxy=sample.bboxes_xyxy,
                labels=sample.labels,
                is_crowd=sample.is_crowd,
                additional_samples=None,
            )
        return sample

    def get_equivalent_preprocessing(self) -> List:
        if self.prob < 1:
            raise RuntimeError("Cannot set preprocessing pipeline with randomness. Set prob to 1.")
        return [{Processings.ReverseImageChannels: {}}]

DetectionRandomAffine

Bases: AbstractDetectionTransform, LegacyDetectionTransformMixin

DetectionRandomAffine detection transform

:param degrees: Degrees for random rotation, when float the random values are drawn uniformly from (-degrees, degrees) :param translate: Translate size (in pixels) for random translation, when float the random values are drawn uniformly from (center-translate, center+translate) :param scales: Values for random rescale, when float the random values are drawn uniformly from (1-scales, 1+scales) :param shear: Degrees for random shear, when float the random values are drawn uniformly from (-shear, shear) :param target_size: Desired output shape. :param filter_box_candidates: Whether to filter out transformed bboxes by edge size, area ratio, and aspect ratio (default=False). :param wh_thr: Edge size threshold when filter_box_candidates = True. Bounding oxes with edges smaller than this values will be filtered out. :param ar_thr: Aspect ratio threshold filter_box_candidates = True. Bounding boxes with aspect ratio larger than this values will be filtered out. :param area_thr: Threshold for area ratio between original image and the transformed one, when filter_box_candidates = True. Bounding boxes with such ratio smaller than this value will be filtered out. :param border_value: Value for filling borders after applying transforms.

Source code in src/super_gradients/training/transforms/transforms.py

@register_transform(Transforms.DetectionRandomAffine)
class DetectionRandomAffine(AbstractDetectionTransform, LegacyDetectionTransformMixin):
    """
    DetectionRandomAffine detection transform

     :param degrees:                Degrees for random rotation, when float the random values are drawn uniformly from (-degrees, degrees)
     :param translate:              Translate size (in pixels) for random translation, when float the random values are drawn uniformly from
                                    (center-translate, center+translate)
     :param scales:                 Values for random rescale, when float the random values are drawn uniformly from (1-scales, 1+scales)
     :param shear:                  Degrees for random shear, when float the random values are drawn uniformly from (-shear, shear)
     :param target_size:            Desired output shape.
     :param filter_box_candidates:  Whether to filter out transformed bboxes by edge size, area ratio, and aspect ratio (default=False).
     :param wh_thr:                 Edge size threshold when filter_box_candidates = True.
                                    Bounding oxes with edges smaller than this values will be filtered out.
     :param ar_thr:                 Aspect ratio threshold filter_box_candidates = True.
                                    Bounding boxes with aspect ratio larger than this values will be filtered out.
     :param area_thr:               Threshold for area ratio between original image and the transformed one, when filter_box_candidates = True.
                                    Bounding boxes with such ratio smaller than this value will be filtered out.
     :param border_value:           Value for filling borders after applying transforms.
    """

    def __init__(
        self,
        degrees: Union[tuple, float] = 10,
        translate: Union[tuple, float] = 0.1,
        scales: Union[tuple, float] = 0.1,
        shear: Union[tuple, float] = 10,
        target_size: Union[int, Tuple[int, int], None] = (640, 640),
        filter_box_candidates: bool = False,
        wh_thr: float = 2,
        ar_thr: float = 20,
        area_thr: float = 0.1,
        border_value: int = 114,
    ):
        super(DetectionRandomAffine, self).__init__()
        self.degrees = degrees
        self.translate = translate
        self.scale = scales
        self.shear = shear
        self.target_size = ensure_is_tuple_of_two(target_size)
        self.enable = True
        self.filter_box_candidates = filter_box_candidates
        self.wh_thr = wh_thr
        self.ar_thr = ar_thr
        self.area_thr = area_thr
        self.border_value = border_value

    def close(self):
        self.enable = False

    def apply_to_sample(self, sample: DetectionSample) -> DetectionSample:
        if self.enable:
            crowd_mask = sample.is_crowd > 0
            crowd_targets = np.concatenate([sample.bboxes_xyxy[crowd_mask], sample.labels[crowd_mask, None]], axis=1)
            targets = np.concatenate([sample.bboxes_xyxy[~crowd_mask], sample.labels[~crowd_mask, None]], axis=1)

            img, targets, crowd_targets = random_affine(
                sample.image,
                targets=targets,
                targets_seg=None,
                crowd_targets=crowd_targets,
                target_size=self.target_size or tuple(reversed(sample.image.shape[:2])),
                degrees=self.degrees,
                translate=self.translate,
                scales=self.scale,
                shear=self.shear,
                filter_box_candidates=self.filter_box_candidates,
                wh_thr=self.wh_thr,
                area_thr=self.area_thr,
                ar_thr=self.ar_thr,
                border_value=self.border_value,
            )

            is_crowd = np.array([0] * len(targets) + [1] * len(crowd_targets), dtype=bool)
            bboxes_xyxy = np.concatenate([targets[:, 0:4], crowd_targets[:, 0:4]], axis=0, dtype=sample.bboxes_xyxy.dtype)
            labels = np.concatenate([targets[:, 4], crowd_targets[:, 4]], axis=0, dtype=sample.labels.dtype)

            sample = DetectionSample(
                image=img,
                bboxes_xyxy=bboxes_xyxy,
                labels=labels,
                is_crowd=is_crowd,
                additional_samples=None,
            )
        return sample

    def get_equivalent_preprocessing(self):
        raise NotImplementedError("get_equivalent_preprocessing is not implemented for non-deterministic transforms.")

DetectionRandomRotate90

Bases: AbstractDetectionTransform, LegacyDetectionTransformMixin

Source code in src/super_gradients/training/transforms/transforms.py

@register_transform(Transforms.DetectionRandomRotate90)
class DetectionRandomRotate90(AbstractDetectionTransform, LegacyDetectionTransformMixin):
    def __init__(self, prob: float = 0.5):
        super().__init__()
        self.prob = prob

    def apply_to_sample(self, sample: DetectionSample) -> DetectionSample:
        if random.random() < self.prob:
            k = random.randrange(0, 4)
            image_shape = sample.image.shape[:2]
            sample = DetectionSample(
                image=self.apply_to_image(sample.image, k),
                bboxes_xyxy=self.apply_to_bboxes(sample.bboxes_xyxy, k, image_shape),
                labels=sample.labels,
                is_crowd=sample.is_crowd,
                additional_samples=None,
            )
        return sample

    def apply_to_image(self, image: np.ndarray, factor: int) -> np.ndarray:
        """
        Apply a `factor` number of 90-degree rotation to image.

        :param image:  Input image (HWC).
        :param factor: Number of CCW rotations. Must be in set {0, 1, 2, 3} See np.rot90.
        :return:       Rotated image (HWC).
        """
        return np.ascontiguousarray(np.rot90(image, factor))

    def apply_to_bboxes(self, bboxes: np.ndarray, factor: int, image_shape: Tuple[int, int]):
        """
        Apply a `factor` number of 90-degree rotation to bounding boxes.

        :param bboxes:       Input bounding boxes in XYXY format.
        :param factor:       Number of CCW rotations. Must be in set {0, 1, 2, 3} See np.rot90.
        :param image_shape:  Original image shape
        :return:             Rotated bounding boxes in XYXY format.
        """
        rows, cols = image_shape
        bboxes_rotated = self.xyxy_bbox_rot90(bboxes, factor, rows, cols)
        return bboxes_rotated

    @classmethod
    def xyxy_bbox_rot90(cls, bboxes: np.ndarray, factor: int, rows: int, cols: int):
        """
        Rotates a bounding box by 90 degrees CCW (see np.rot90)

        :param bboxes:  Tensor made of bounding box tuples (x_min, y_min, x_max, y_max).
        :param factor:  Number of CCW rotations. Must be in set {0, 1, 2, 3} See np.rot90.
        :param rows:    Image rows of the original image.
        :param cols:    Image cols of the original image.

        :return: A bounding box tuple (x_min, y_min, x_max, y_max).

        """
        x_min, y_min, x_max, y_max = bboxes[:, 0], bboxes[:, 1], bboxes[:, 2], bboxes[:, 3]

        if factor == 0:
            bbox = x_min, y_min, x_max, y_max
        elif factor == 1:
            bbox = y_min, cols - x_max, y_max, cols - x_min
        elif factor == 2:
            bbox = cols - x_max, rows - y_max, cols - x_min, rows - y_min
        elif factor == 3:
            bbox = rows - y_max, x_min, rows - y_min, x_max
        else:
            raise ValueError("Parameter n must be in set {0, 1, 2, 3}")
        return np.stack(bbox, axis=1)

    def get_equivalent_preprocessing(self) -> List[Dict]:
        raise NotImplementedError("get_equivalent_preprocessing is not implemented for non-deterministic transforms.")

apply_to_bboxes(bboxes, factor, image_shape)

Apply a factor number of 90-degree rotation to bounding boxes.

Parameters:

Name	Type	Description	Default
bboxes	np.ndarray	Input bounding boxes in XYXY format.	required
factor	int	Number of CCW rotations. Must be in set {0, 1, 2, 3} See np.rot90.	required
image_shape	Tuple[int, int]	Original image shape	required

Returns:

Type	Description
	Rotated bounding boxes in XYXY format.

Source code in src/super_gradients/training/transforms/transforms.py

def apply_to_bboxes(self, bboxes: np.ndarray, factor: int, image_shape: Tuple[int, int]):
    """
    Apply a `factor` number of 90-degree rotation to bounding boxes.

    :param bboxes:       Input bounding boxes in XYXY format.
    :param factor:       Number of CCW rotations. Must be in set {0, 1, 2, 3} See np.rot90.
    :param image_shape:  Original image shape
    :return:             Rotated bounding boxes in XYXY format.
    """
    rows, cols = image_shape
    bboxes_rotated = self.xyxy_bbox_rot90(bboxes, factor, rows, cols)
    return bboxes_rotated

apply_to_image(image, factor)

Apply a factor number of 90-degree rotation to image.

Parameters:

Name	Type	Description	Default
image	np.ndarray	Input image (HWC).	required
factor	int	Number of CCW rotations. Must be in set {0, 1, 2, 3} See np.rot90.	required

Returns:

Type	Description
np.ndarray	Rotated image (HWC).

Source code in src/super_gradients/training/transforms/transforms.py

def apply_to_image(self, image: np.ndarray, factor: int) -> np.ndarray:
    """
    Apply a `factor` number of 90-degree rotation to image.

    :param image:  Input image (HWC).
    :param factor: Number of CCW rotations. Must be in set {0, 1, 2, 3} See np.rot90.
    :return:       Rotated image (HWC).
    """
    return np.ascontiguousarray(np.rot90(image, factor))

xyxy_bbox_rot90(bboxes, factor, rows, cols) classmethod

Rotates a bounding box by 90 degrees CCW (see np.rot90)

Parameters:

Name	Type	Description	Default
bboxes	np.ndarray	Tensor made of bounding box tuples (x_min, y_min, x_max, y_max).	required
factor	int	Number of CCW rotations. Must be in set {0, 1, 2, 3} See np.rot90.	required
rows	int	Image rows of the original image.	required
cols	int	Image cols of the original image.	required

Returns:

Type	Description
	A bounding box tuple (x_min, y_min, x_max, y_max).

Source code in src/super_gradients/training/transforms/transforms.py

@classmethod
def xyxy_bbox_rot90(cls, bboxes: np.ndarray, factor: int, rows: int, cols: int):
    """
    Rotates a bounding box by 90 degrees CCW (see np.rot90)

    :param bboxes:  Tensor made of bounding box tuples (x_min, y_min, x_max, y_max).
    :param factor:  Number of CCW rotations. Must be in set {0, 1, 2, 3} See np.rot90.
    :param rows:    Image rows of the original image.
    :param cols:    Image cols of the original image.

    :return: A bounding box tuple (x_min, y_min, x_max, y_max).

    """
    x_min, y_min, x_max, y_max = bboxes[:, 0], bboxes[:, 1], bboxes[:, 2], bboxes[:, 3]

    if factor == 0:
        bbox = x_min, y_min, x_max, y_max
    elif factor == 1:
        bbox = y_min, cols - x_max, y_max, cols - x_min
    elif factor == 2:
        bbox = cols - x_max, rows - y_max, cols - x_min, rows - y_min
    elif factor == 3:
        bbox = rows - y_max, x_min, rows - y_min, x_max
    else:
        raise ValueError("Parameter n must be in set {0, 1, 2, 3}")
    return np.stack(bbox, axis=1)

DetectionRescale

Bases: AbstractDetectionTransform, LegacyDetectionTransformMixin

Resize image and bounding boxes to given image dimensions without preserving aspect ratio

Parameters:

Name	Type	Description	Default
output_shape	Union[int, Tuple[int, int]]	(rows, cols)	required

Source code in src/super_gradients/training/transforms/transforms.py

@register_transform(Transforms.DetectionRescale)
class DetectionRescale(AbstractDetectionTransform, LegacyDetectionTransformMixin):
    """
    Resize image and bounding boxes to given image dimensions without preserving aspect ratio

    :param output_shape: (rows, cols)
    """

    def __init__(self, output_shape: Union[int, Tuple[int, int]]):
        super().__init__()
        self.output_shape = ensure_is_tuple_of_two(output_shape)

    @classmethod
    def apply_to_image(self, image: np.ndarray, output_width: int, output_height: int) -> np.ndarray:
        return _rescale_image(image=image, target_shape=(output_height, output_width))

    @classmethod
    def apply_to_bboxes(self, bboxes: np.ndarray, sx: float, sy: float) -> np.ndarray:
        return _rescale_bboxes(bboxes, scale_factors=(sy, sx))

    def apply_to_sample(self, sample: DetectionSample) -> DetectionSample:
        image_height, image_width = sample.image.shape[:2]
        output_height, output_width = self.output_shape
        sx = output_width / image_width
        sy = output_height / image_height

        return DetectionSample(
            image=self.apply_to_image(sample.image, output_width=output_width, output_height=output_height),
            bboxes_xyxy=self.apply_to_bboxes(sample.bboxes_xyxy, sx=sx, sy=sy),
            labels=sample.labels,
            is_crowd=sample.is_crowd,
            additional_samples=None,
        )

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

DetectionStandardize

Bases: AbstractDetectionTransform, LegacyDetectionTransformMixin

Standardize image pixel values with img/max_val

Parameters:

Name	Type	Description	Default
max_val		Current maximum value of the image pixels. (usually 255)	required

Source code in src/super_gradients/training/transforms/transforms.py

@register_transform(Transforms.DetectionStandardize)
class DetectionStandardize(AbstractDetectionTransform, LegacyDetectionTransformMixin):
    """
    Standardize image pixel values with img/max_val

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

    def __init__(self, max_value: float = 255.0):
        super().__init__()
        self.max_value = float(max_value)

    @classmethod
    def apply_to_image(self, image: np.ndarray, max_value: float) -> np.ndarray:
        return (image / max_value).astype(np.float32)

    def apply_to_sample(self, sample: DetectionSample) -> DetectionSample:
        sample.image = self.apply_to_image(sample.image, max_value=self.max_value)
        return sample

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

DetectionTargetsFormatTransform

Bases: AbstractDetectionTransform, LegacyDetectionTransformMixin

Detection targets format transform

Convert targets in input_format to output_format, filter small bboxes and pad targets.

Parameters:

Name	Type	Description	Default
input_dim	Union[int, Tuple[int, int], None]	Shape of the images to transform.	None
input_format	ConcatenatedTensorFormat	Format of the input targets. For instance [xmin, ymin, xmax, ymax, cls_id] refers to XYXY_LABEL.	XYXY_LABEL
output_format	ConcatenatedTensorFormat	Format of the output targets. For instance [xmin, ymin, xmax, ymax, cls_id] refers to XYXY_LABEL	LABEL_CXCYWH
min_bbox_edge_size	float	bboxes with edge size lower than this values will be removed.	1

Source code in src/super_gradients/training/transforms/transforms.py

@register_transform(Transforms.DetectionTargetsFormatTransform)
class DetectionTargetsFormatTransform(AbstractDetectionTransform, LegacyDetectionTransformMixin):
    """
    Detection targets format transform

    Convert targets in input_format to output_format, filter small bboxes and pad targets.

    :param input_dim:          Shape of the images to transform.
    :param input_format:       Format of the input targets. For instance [xmin, ymin, xmax, ymax, cls_id] refers to XYXY_LABEL.
    :param output_format:      Format of the output targets. For instance [xmin, ymin, xmax, ymax, cls_id] refers to XYXY_LABEL
    :param min_bbox_edge_size: bboxes with edge size lower than this values will be removed.
    """

    @resolve_param("input_format", ConcatenatedTensorFormatFactory())
    @resolve_param("output_format", ConcatenatedTensorFormatFactory())
    def __init__(
        self,
        input_dim: Union[int, Tuple[int, int], None] = None,
        input_format: ConcatenatedTensorFormat = XYXY_LABEL,
        output_format: ConcatenatedTensorFormat = LABEL_CXCYWH,
        min_bbox_edge_size: float = 1,
        max_targets: Optional[int] = None,
    ):
        super(DetectionTargetsFormatTransform, self).__init__()
        _max_targets_deprication(max_targets)
        if isinstance(input_format, DetectionTargetsFormat) or isinstance(output_format, DetectionTargetsFormat):
            raise TypeError(
                "DetectionTargetsFormat is not supported for input_format and output_format starting from super_gradients==3.0.7.\n"
                "You can either:\n"
                "\t - use builtin format among super_gradients.training.datasets.data_formats.default_formats.<FORMAT_NAME> (e.g. XYXY_LABEL, CXCY_LABEL, ..)\n"
                "\t - define your custom format using super_gradients.training.datasets.data_formats.formats.ConcatenatedTensorFormat\n"
            )
        self.input_format = input_format
        self.output_format = output_format
        self.min_bbox_edge_size = min_bbox_edge_size
        self.input_dim = None

        if input_dim is not None:
            input_dim = ensure_is_tuple_of_two(input_dim)
            self._setup_input_dim_related_params(input_dim)

    def _setup_input_dim_related_params(self, input_dim: tuple):
        """Setup all the parameters that are related to input_dim."""
        self.input_dim = input_dim
        self.min_bbox_edge_size = self.min_bbox_edge_size / max(input_dim) if self.output_format.bboxes_format.format.normalized else self.min_bbox_edge_size
        self.targets_format_converter = ConcatenatedTensorFormatConverter(
            input_format=self.input_format, output_format=self.output_format, image_shape=input_dim
        )

    def __call__(self, sample: Union[dict, DetectionSample]) -> dict:
        if isinstance(sample, DetectionSample):
            pass
        else:
            # if self.input_dim not set yet, it will be set with first batch
            if self.input_dim is None:
                self._setup_input_dim_related_params(input_dim=sample["image"].shape[1:])

            sample["target"] = self.apply_on_targets(sample["target"])
            if "crowd_target" in sample.keys():
                sample["crowd_target"] = self.apply_on_targets(sample["crowd_target"])
        return sample

    def apply_to_sample(self, sample: DetectionSample):
        # DIRTY HACK: No-op if a detection sample is passed
        # DIRTY HACK: As a workaround we will do this transform in dataset class for now
        return sample

    def apply_on_targets(self, targets: np.ndarray) -> np.ndarray:
        """Convert targets in input_format to output_format, filter small bboxes and pad targets"""
        targets = self.filter_small_bboxes(targets)
        targets = self.targets_format_converter(targets)
        return np.ascontiguousarray(targets, dtype=np.float32)

    def filter_small_bboxes(self, targets: np.ndarray) -> np.ndarray:
        """Filter bboxes smaller than specified threshold."""

        def _is_big_enough(bboxes: np.ndarray) -> np.ndarray:
            bboxes_xywh = xyxy_to_xywh(bboxes, image_shape=None)
            return np.minimum(bboxes_xywh[:, 2], bboxes_xywh[:, 3]) > self.min_bbox_edge_size

        targets = filter_on_bboxes(fn=_is_big_enough, tensor=targets, tensor_format=self.input_format)
        return targets

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

apply_on_targets(targets)

Convert targets in input_format to output_format, filter small bboxes and pad targets

Source code in src/super_gradients/training/transforms/transforms.py

def apply_on_targets(self, targets: np.ndarray) -> np.ndarray:
    """Convert targets in input_format to output_format, filter small bboxes and pad targets"""
    targets = self.filter_small_bboxes(targets)
    targets = self.targets_format_converter(targets)
    return np.ascontiguousarray(targets, dtype=np.float32)

filter_small_bboxes(targets)

Filter bboxes smaller than specified threshold.

Source code in src/super_gradients/training/transforms/transforms.py

def filter_small_bboxes(self, targets: np.ndarray) -> np.ndarray:
    """Filter bboxes smaller than specified threshold."""

    def _is_big_enough(bboxes: np.ndarray) -> np.ndarray:
        bboxes_xywh = xyxy_to_xywh(bboxes, image_shape=None)
        return np.minimum(bboxes_xywh[:, 2], bboxes_xywh[:, 3]) > self.min_bbox_edge_size

    targets = filter_on_bboxes(fn=_is_big_enough, tensor=targets, tensor_format=self.input_format)
    return targets

DetectionTransform

Detection transform base class. Complex transforms that require extra data loading can use the the additional_samples_count attribute in a similar fashion to what's been done in COCODetectionDataset: self._load_additional_inputs_for_transform(sample, transform)

after the above call, sample["additional_samples"] holds a list of additional inputs and targets.

sample = transform(sample)

Parameters:

Name	Type	Description	Default
additional_samples_count	int	Additional samples to be loaded.	0
non_empty_targets	bool	Whether the additional targets can have empty targets or not.	False

Source code in src/super_gradients/training/transforms/transforms.py

class DetectionTransform:
    """
    Detection transform base class.
    Complex transforms that require extra data loading can use the the additional_samples_count attribute in a
     similar fashion to what's been done in COCODetectionDataset:
    self._load_additional_inputs_for_transform(sample, transform)
    # after the above call, sample["additional_samples"] holds a list of additional inputs and targets.
    sample = transform(sample)
    :param additional_samples_count:    Additional samples to be loaded.
    :param non_empty_targets:           Whether the additional targets can have empty targets or not.
    """

    def __init__(self, additional_samples_count: int = 0, non_empty_targets: bool = False):
        self.additional_samples_count = additional_samples_count
        self.non_empty_targets = non_empty_targets
        warnings.warn(
            "Inheriting from DetectionTransform is deprecated. "
            "If you have a custom detection transform please change the base class to "
            "AbstractDetectionTransform and implement apply_to_sample() method instead of __call__.",
            DeprecationWarning,
        )

    def __call__(self, sample: Union[dict, list]):
        raise NotImplementedError

    def apply_to_sample(self, sample: DetectionSample) -> DetectionSample:
        """
        Apply transformation to the input detection sample.
        This method exists here for compatibility reasons to ensure a custom transform that inherits from DetectionSample
        would still work.

        :param sample: Input detection sample.
        :return:       Output detection sample.
        """
        sample_dict = LegacyDetectionTransformMixin.convert_detection_sample_to_dict(sample, include_crowd_target=sample.is_crowd.any())
        sample_dict = self(sample_dict)
        return LegacyDetectionTransformMixin.convert_input_dict_to_detection_sample(sample_dict)

    def __repr__(self):
        return self.__class__.__name__ + str(self.__dict__).replace("{", "(").replace("}", ")")

    def get_equivalent_preprocessing(self) -> List:
        raise NotImplementedError

apply_to_sample(sample)

Apply transformation to the input detection sample. This method exists here for compatibility reasons to ensure a custom transform that inherits from DetectionSample would still work.

Parameters:

Name	Type	Description	Default
sample	DetectionSample	Input detection sample.	required

Returns:

Type	Description
DetectionSample	Output detection sample.

Source code in src/super_gradients/training/transforms/transforms.py

def apply_to_sample(self, sample: DetectionSample) -> DetectionSample:
    """
    Apply transformation to the input detection sample.
    This method exists here for compatibility reasons to ensure a custom transform that inherits from DetectionSample
    would still work.

    :param sample: Input detection sample.
    :return:       Output detection sample.
    """
    sample_dict = LegacyDetectionTransformMixin.convert_detection_sample_to_dict(sample, include_crowd_target=sample.is_crowd.any())
    sample_dict = self(sample_dict)
    return LegacyDetectionTransformMixin.convert_input_dict_to_detection_sample(sample_dict)

DetectionVerticalFlip

Bases: AbstractDetectionTransform, LegacyDetectionTransformMixin

Vertical Flip for Detection

Parameters:

Name	Type	Description	Default
prob	float	Probability of applying vertical flip	required

Source code in src/super_gradients/training/transforms/transforms.py

@register_transform(Transforms.DetectionVerticalFlip)
class DetectionVerticalFlip(AbstractDetectionTransform, LegacyDetectionTransformMixin):
    """
    Vertical Flip for Detection

    :param prob:        Probability of applying vertical flip
    """

    def __init__(self, prob: float):
        super(DetectionVerticalFlip, self).__init__()
        self.prob = float(prob)

    def apply_to_sample(self, sample: DetectionSample) -> DetectionSample:
        if random.random() < self.prob:
            sample = DetectionSample(
                image=_flip_vertical_image(sample.image),
                bboxes_xyxy=_flip_vertical_boxes_xyxy(sample.bboxes_xyxy, sample.image.shape[0]),
                labels=sample.labels,
                is_crowd=sample.is_crowd,
                additional_samples=None,
            )
        return sample

    def get_equivalent_preprocessing(self) -> List[Dict]:
        raise NotImplementedError("get_equivalent_preprocessing is not implemented for non-deterministic transforms.")

SegConvertToTensor

Bases: AbstractSegmentationTransform, LegacySegmentationTransformMixin

Converts SegmentationSample images and masks to PyTorch tensors.

Parameters:

Name	Type	Description	Default
(Optional[str])	mask_output_dtype	The desired output data type for the mask tensor.	required
(bool)	add_mask_dummy_dim	Whether to add a dummy channels dimension to the mask tensor.	required

Source code in src/super_gradients/training/transforms/transforms.py

@register_transform(Transforms.SegConvertToTensor)
class SegConvertToTensor(AbstractSegmentationTransform, LegacySegmentationTransformMixin):
    """
    Converts SegmentationSample images and masks to PyTorch tensors.

    :param mask_output_dtype (Optional[str]): The desired output data type for the mask tensor.
    :param add_mask_dummy_dim (bool): Whether to add a dummy channels dimension to the mask tensor.
    """

    @resolve_param("image_output_dtype", TorchDtypeFactory())
    @resolve_param("mask_output_dtype", TorchDtypeFactory())
    def __init__(self, image_output_dtype=torch.float32, mask_output_dtype: Optional[torch.dtype] = None, add_mask_dummy_dim: bool = False):
        self.image_output_dtype = image_output_dtype
        self.mask_output_dtype = mask_output_dtype
        self.add_mask_dummy_dim = add_mask_dummy_dim

    def apply_to_sample(self, sample: SegmentationSample) -> SegmentationSample:
        sample.image = torch.from_numpy(sample.image).permute(2, 0, 1).to(self.image_output_dtype)
        sample.mask = torch.from_numpy(np.array(sample.mask))

        # Convert mask to torch tensor with specified dtype
        if self.mask_output_dtype is not None:
            sample.mask = torch.from_numpy(np.array(sample.mask)).to(dtype=self.mask_output_dtype)

        # Add dummy channels dimension if needed
        if self.add_mask_dummy_dim and len(sample.mask.shape) == 2:
            sample.mask = sample.mask.unsqueeze(0)

        return sample

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

SegCropImageAndMask

Bases: AbstractSegmentationTransform, LegacySegmentationTransformMixin

Crops image and mask (synchronously). In "center" mode a center crop is performed while, in "random" mode the drop will be positioned around random coordinates.

Source code in src/super_gradients/training/transforms/transforms.py

@register_transform(Transforms.SegCropImageAndMask)
class SegCropImageAndMask(AbstractSegmentationTransform, LegacySegmentationTransformMixin):
    """
    Crops image and mask (synchronously).
    In "center" mode a center crop is performed while, in "random" mode the drop will be positioned around
     random coordinates.
    """

    def __init__(self, crop_size: Union[float, Tuple, List], mode: str):
        """

        :param crop_size: tuple of (width, height) for the final crop size, if is scalar size is a
            square (crop_size, crop_size)
        :param mode: how to choose the center of the crop, 'center' for the center of the input image,
            'random' center the point is chosen randomally
        """

        self.crop_size = crop_size
        self.mode = mode

        self.check_valid_arguments()

    def apply_to_sample(self, sample: SegmentationSample) -> SegmentationSample:
        image = Image.fromarray(sample.image)
        mask = Image.fromarray(sample.mask)

        w, h = image.size
        if self.mode == "random":
            x1 = random.randint(0, w - self.crop_size[0])
            y1 = random.randint(0, h - self.crop_size[1])
        else:
            x1 = int(round((w - self.crop_size[0]) / 2.0))
            y1 = int(round((h - self.crop_size[1]) / 2.0))

        image = image.crop((x1, y1, x1 + self.crop_size[0], y1 + self.crop_size[1]))
        mask = mask.crop((x1, y1, x1 + self.crop_size[0], y1 + self.crop_size[1]))

        return SegmentationSample(image=image, mask=mask)

    def check_valid_arguments(self):
        if self.mode not in ["center", "random"]:
            raise ValueError(f"Unsupported mode: found: {self.mode}, expected: 'center' or 'random'")

        if not isinstance(self.crop_size, Iterable):
            self.crop_size = (self.crop_size, self.crop_size)
        if self.crop_size[0] <= 0 or self.crop_size[1] <= 0:
            raise ValueError(f"Crop size must be positive numbers, found: {self.crop_size}")

    def get_equivalent_preprocessing(self) -> List[Dict]:
        raise NotImplementedError("get_equivalent_preprocessing is not implemented for non-deterministic transforms.")

__init__(crop_size, mode)

Parameters:

Name	Type	Description	Default
crop_size	Union[float, Tuple, List]	tuple of (width, height) for the final crop size, if is scalar size is a square (crop_size, crop_size)	required
mode	str	how to choose the center of the crop, 'center' for the center of the input image, 'random' center the point is chosen randomally	required

Source code in src/super_gradients/training/transforms/transforms.py

def __init__(self, crop_size: Union[float, Tuple, List], mode: str):
    """

    :param crop_size: tuple of (width, height) for the final crop size, if is scalar size is a
        square (crop_size, crop_size)
    :param mode: how to choose the center of the crop, 'center' for the center of the input image,
        'random' center the point is chosen randomally
    """

    self.crop_size = crop_size
    self.mode = mode

    self.check_valid_arguments()

SegNormalize

Bases: AbstractSegmentationTransform, LegacySegmentationTransformMixin

Normalization to be applied on the segmentation sample's image.

Parameters:

Name	Type	Description	Default
(sequence)	mean	Sequence of means for each channel.	required

Source code in src/super_gradients/training/transforms/transforms.py

@register_transform(Transforms.SegNormalize)
class SegNormalize(AbstractSegmentationTransform, LegacySegmentationTransformMixin):
    """
    Normalization to be applied on the segmentation sample's image.

    :param mean (sequence): Sequence of means for each channel.
    :param std (sequence): Sequence of standard deviations for each channel.
    """

    def __init__(self, mean: Sequence[float], std: Sequence[float]):
        self.mean = np.array(mean).reshape(1, 1, -1).astype(np.float32)
        self.std = np.array(std).reshape(1, 1, -1).astype(np.float32)

    def apply_to_sample(self, sample: SegmentationSample) -> SegmentationSample:
        if not isinstance(sample.image, np.ndarray):
            sample.image = np.array(sample.image)
        sample.image = (sample.image - self.mean) / self.std
        return sample

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

SegPadShortToCropSize

Bases: AbstractSegmentationTransform, LegacySegmentationTransformMixin

Pads image to 'crop_size'. Should be called only after "SegRescale" or "SegRandomRescale" in augmentations pipeline. Please note that if input image size > crop size no change will be made to the image. This transform only pads the image and mask into "crop_size" if it's larger than image size

Source code in src/super_gradients/training/transforms/transforms.py

@register_transform(Transforms.SegPadShortToCropSize)
class SegPadShortToCropSize(AbstractSegmentationTransform, LegacySegmentationTransformMixin):
    """
    Pads image to 'crop_size'.
    Should be called only after "SegRescale" or "SegRandomRescale" in augmentations pipeline.
    Please note that if input image size > crop size no change will be made to the image.
    This transform only pads the image and mask into "crop_size" if it's larger than image size
    """

    def __init__(self, crop_size: Union[float, Tuple, List], fill_mask: int = 0, fill_image: Union[int, Tuple, List] = 0):
        """
        :param crop_size:   Tuple of (width, height) for the final crop size, if is scalar size is a square (crop_size, crop_size)
        :param fill_mask:   Value to fill mask labels background.
        :param fill_image:  Grey value to fill image padded background.
        """
        # CHECK IF CROP SIZE IS A ITERABLE OR SCALAR
        self.crop_size = crop_size
        self.fill_mask = fill_mask
        self.fill_image = tuple(fill_image) if isinstance(fill_image, Sequence) else fill_image

        self.check_valid_arguments()

    def apply_to_sample(self, sample: SegmentationSample) -> SegmentationSample:
        image = Image.fromarray(sample.image)
        mask = Image.fromarray(sample.mask)
        w, h = image.size

        # pad images from center symmetrically
        if w < self.crop_size[0] or h < self.crop_size[1]:
            padh = (self.crop_size[1] - h) / 2 if h < self.crop_size[1] else 0
            pad_top, pad_bottom = math.ceil(padh), math.floor(padh)
            padw = (self.crop_size[0] - w) / 2 if w < self.crop_size[0] else 0
            pad_left, pad_right = math.ceil(padw), math.floor(padw)

            image = ImageOps.expand(image, border=(pad_left, pad_top, pad_right, pad_bottom), fill=self.fill_image)
            mask = ImageOps.expand(mask, border=(pad_left, pad_top, pad_right, pad_bottom), fill=self.fill_mask)

        return SegmentationSample(image=image, mask=mask)

    def check_valid_arguments(self):
        if not isinstance(self.crop_size, Iterable):
            self.crop_size = (self.crop_size, self.crop_size)
        if self.crop_size[0] <= 0 or self.crop_size[1] <= 0:
            raise ValueError(f"Crop size must be positive numbers, found: {self.crop_size}")

        self.fill_mask, self.fill_image = _validate_fill_values_arguments(self.fill_mask, self.fill_image)

    def get_equivalent_preprocessing(self) -> List[Dict]:
        return [{Processings.SegmentationPadShortToCropSize: {"crop_size": self.crop_size, "fill_image": self.fill_image}}]

__init__(crop_size, fill_mask=0, fill_image=0)

Parameters:

Name	Type	Description	Default
crop_size	Union[float, Tuple, List]	Tuple of (width, height) for the final crop size, if is scalar size is a square (crop_size, crop_size)	required
fill_mask	int	Value to fill mask labels background.	0
fill_image	Union[int, Tuple, List]	Grey value to fill image padded background.	0

Source code in src/super_gradients/training/transforms/transforms.py

def __init__(self, crop_size: Union[float, Tuple, List], fill_mask: int = 0, fill_image: Union[int, Tuple, List] = 0):
    """
    :param crop_size:   Tuple of (width, height) for the final crop size, if is scalar size is a square (crop_size, crop_size)
    :param fill_mask:   Value to fill mask labels background.
    :param fill_image:  Grey value to fill image padded background.
    """
    # CHECK IF CROP SIZE IS A ITERABLE OR SCALAR
    self.crop_size = crop_size
    self.fill_mask = fill_mask
    self.fill_image = tuple(fill_image) if isinstance(fill_image, Sequence) else fill_image

    self.check_valid_arguments()

SegRandomFlip

Bases: AbstractSegmentationTransform, LegacySegmentationTransformMixin

Randomly flips the image and mask (synchronously) with probability 'prob'.

Source code in src/super_gradients/training/transforms/transforms.py

@register_transform(Transforms.SegRandomFlip)
class SegRandomFlip(AbstractSegmentationTransform, LegacySegmentationTransformMixin):
    """
    Randomly flips the image and mask (synchronously) with probability 'prob'.
    """

    def __init__(self, prob: float = 0.5):
        assert 0.0 <= prob <= 1.0, f"Probability value must be between 0 and 1, found {prob}"
        self.prob = prob

    def apply_to_sample(self, sample: SegmentationSample) -> SegmentationSample:
        image = Image.fromarray(sample.image)
        mask = Image.fromarray(sample.mask)
        if random.random() < self.prob:
            image = image.transpose(Image.FLIP_LEFT_RIGHT)
            mask = mask.transpose(Image.FLIP_LEFT_RIGHT)
            sample = SegmentationSample(image=image, mask=mask)
        return sample

    def get_equivalent_preprocessing(self) -> List[Dict]:
        raise NotImplementedError("get_equivalent_preprocessing is not implemented for non-deterministic transforms.")

SegRandomGaussianBlur

Bases: AbstractSegmentationTransform, LegacySegmentationTransformMixin

Adds random Gaussian Blur to image with probability 'prob'.

Source code in src/super_gradients/training/transforms/transforms.py

@register_transform(Transforms.SegRandomGaussianBlur)
class SegRandomGaussianBlur(AbstractSegmentationTransform, LegacySegmentationTransformMixin):
    """
    Adds random Gaussian Blur to image with probability 'prob'.
    """

    def __init__(self, prob: float = 0.5):
        assert 0.0 <= prob <= 1.0, "Probability value must be between 0 and 1"
        self.prob = prob

    def apply_to_sample(self, sample: SegmentationSample) -> SegmentationSample:
        image = Image.fromarray(sample.image)

        if random.random() < self.prob:
            image = image.filter(ImageFilter.GaussianBlur(radius=random.random()))

        return SegmentationSample(image=image, mask=sample.mask)

    def get_equivalent_preprocessing(self) -> List[Dict]:
        raise NotImplementedError("get_equivalent_preprocessing is not implemented for non-deterministic transforms.")

SegRandomRescale

Bases: AbstractSegmentationTransform, LegacySegmentationTransformMixin

Random rescale the image and mask (synchronously) while preserving aspect ratio. Scale factor is randomly picked between scales [min, max]

Parameters:

Name	Type	Description	Default
scales	Union[float, Tuple, List]	Scale range tuple (min, max), if scales is a float range will be defined as (1, scales) if scales > 1, otherwise (scales, 1). must be a positive number.	(0.5, 2.0)

Source code in src/super_gradients/training/transforms/transforms.py

@register_transform(Transforms.SegRandomRescale)
class SegRandomRescale(AbstractSegmentationTransform, LegacySegmentationTransformMixin):
    """
    Random rescale the image and mask (synchronously) while preserving aspect ratio.
    Scale factor is randomly picked between scales [min, max]

    :param scales: Scale range tuple (min, max), if scales is a float range will be defined as (1, scales) if scales > 1,
            otherwise (scales, 1). must be a positive number.
    """

    def __init__(self, scales: Union[float, Tuple, List] = (0.5, 2.0)):
        self.scales = scales

        self.check_valid_arguments()

    def apply_to_sample(self, sample: SegmentationSample) -> SegmentationSample:
        image = Image.fromarray(sample.image)
        mask = Image.fromarray(sample.mask)
        w, h = image.size

        scale = random.uniform(self.scales[0], self.scales[1])
        out_size = int(scale * w), int(scale * h)
        image = image.resize(out_size, IMAGE_RESAMPLE_MODE)
        mask = mask.resize(out_size, MASK_RESAMPLE_MODE)

        return SegmentationSample(image=image, mask=mask)

    def check_valid_arguments(self):
        """
        Check the scale values are valid. if order is wrong, flip the order and return the right scale values.
        """
        if not isinstance(self.scales, Iterable):
            if self.scales <= 1:
                self.scales = (self.scales, 1)
            else:
                self.scales = (1, self.scales)

        if self.scales[0] < 0 or self.scales[1] < 0:
            raise ValueError(f"SegRandomRescale scale values must be positive numbers, found: {self.scales}")
        if self.scales[0] > self.scales[1]:
            self.scales = (self.scales[1], self.scales[0])
        return self.scales

    def get_equivalent_preprocessing(self) -> List[Dict]:
        raise NotImplementedError("get_equivalent_preprocessing is not implemented for non-deterministic transforms.")

check_valid_arguments()

Check the scale values are valid. if order is wrong, flip the order and return the right scale values.

Source code in src/super_gradients/training/transforms/transforms.py

def check_valid_arguments(self):
    """
    Check the scale values are valid. if order is wrong, flip the order and return the right scale values.
    """
    if not isinstance(self.scales, Iterable):
        if self.scales <= 1:
            self.scales = (self.scales, 1)
        else:
            self.scales = (1, self.scales)

    if self.scales[0] < 0 or self.scales[1] < 0:
        raise ValueError(f"SegRandomRescale scale values must be positive numbers, found: {self.scales}")
    if self.scales[0] > self.scales[1]:
        self.scales = (self.scales[1], self.scales[0])
    return self.scales

SegRandomRotate

Bases: AbstractSegmentationTransform, LegacySegmentationTransformMixin

Randomly rotates image and mask (synchronously) between 'min_deg' and 'max_deg'.

Source code in src/super_gradients/training/transforms/transforms.py

@register_transform(Transforms.SegRandomRotate)
class SegRandomRotate(AbstractSegmentationTransform, LegacySegmentationTransformMixin):
    """
    Randomly rotates image and mask (synchronously) between 'min_deg' and 'max_deg'.
    """

    def __init__(self, min_deg: float = -10, max_deg: float = 10, fill_mask: int = 0, fill_image: Union[int, Tuple, List] = 0):
        self.min_deg = min_deg
        self.max_deg = max_deg
        self.fill_mask = fill_mask
        # grey color in RGB mode
        self.fill_image = (fill_image, fill_image, fill_image)

        self.check_valid_arguments()

    def apply_to_sample(self, sample: SegmentationSample) -> SegmentationSample:
        image = Image.fromarray(sample.image)
        mask = Image.fromarray(sample.mask)

        deg = random.uniform(self.min_deg, self.max_deg)
        image = image.rotate(deg, resample=IMAGE_RESAMPLE_MODE, fillcolor=self.fill_image)
        mask = mask.rotate(deg, resample=MASK_RESAMPLE_MODE, fillcolor=self.fill_mask)

        return SegmentationSample(image=image, mask=mask)

    def check_valid_arguments(self):
        self.fill_mask, self.fill_image = _validate_fill_values_arguments(self.fill_mask, self.fill_image)

    def get_equivalent_preprocessing(self) -> List[Dict]:
        raise NotImplementedError("get_equivalent_preprocessing is not implemented for non-deterministic transforms.")

SegRescale

Bases: AbstractSegmentationTransform, LegacySegmentationTransformMixin

Rescales the image and mask (synchronously) while preserving aspect ratio. The rescaling can be done according to scale_factor, short_size or long_size. If more than one argument is given, the rescaling mode is determined by this order: scale_factor, then short_size, then long_size.

Parameters:

Name	Type	Description	Default
scale_factor	Optional[float]	Rescaling is done by multiplying input size by scale_factor: out_size = (scale_factor * w, scale_factor * h)	None
short_size	Optional[int]	Rescaling is done by determining the scale factor by the ratio short_size / min(h, w).	None
long_size	Optional[int]	Rescaling is done by determining the scale factor by the ratio long_size / max(h, w).	None

Source code in src/super_gradients/training/transforms/transforms.py

@register_transform(Transforms.SegRescale)
class SegRescale(AbstractSegmentationTransform, LegacySegmentationTransformMixin):
    """
    Rescales the image and mask (synchronously) while preserving aspect ratio.
    The rescaling can be done according to scale_factor, short_size or long_size.
    If more than one argument is given, the rescaling mode is determined by this order: scale_factor, then short_size,
    then long_size.

    :param scale_factor: Rescaling is done by multiplying input size by scale_factor:
            out_size = (scale_factor * w, scale_factor * h)
    :param short_size:  Rescaling is done by determining the scale factor by the ratio short_size / min(h, w).
    :param long_size:   Rescaling is done by determining the scale factor by the ratio long_size / max(h, w).
    """

    def __init__(self, scale_factor: Optional[float] = None, short_size: Optional[int] = None, long_size: Optional[int] = None):
        self.scale_factor = scale_factor
        self.short_size = short_size
        self.long_size = long_size

        self.check_valid_arguments()

    def apply_to_sample(self, sample: SegmentationSample) -> SegmentationSample:
        image = Image.fromarray(sample.image)
        mask = Image.fromarray(sample.mask)
        w, h = image.size
        scale = _compute_scale_factor(self.scale_factor, self.short_size, self.long_size, w, h)

        out_size = int(scale * w), int(scale * h)
        image = image.resize(out_size, IMAGE_RESAMPLE_MODE)
        mask = mask.resize(out_size, MASK_RESAMPLE_MODE)

        return SegmentationSample(
            image=image,
            mask=mask,
        )

    def check_valid_arguments(self):
        if self.scale_factor is None and self.short_size is None and self.long_size is None:
            raise ValueError("Must assign one rescale argument: scale_factor, short_size or long_size")

        if self.scale_factor is not None and self.scale_factor <= 0:
            raise ValueError(f"Scale factor must be a positive number, found: {self.scale_factor}")
        if self.short_size is not None and self.short_size <= 0:
            raise ValueError(f"Short size must be a positive number, found: {self.short_size}")
        if self.long_size is not None and self.long_size <= 0:
            raise ValueError(f"Long size must be a positive number, found: {self.long_size}")

    def get_equivalent_preprocessing(self) -> List[Dict]:
        return [{Processings.SegmentationRescale: {"scale_factor": self.scale_factor, "short_size": self.short_size, "long_size": self.long_size}}]

SegStandardize

Bases: AbstractSegmentationTransform, LegacySegmentationTransformMixin

Standardize image pixel values with img/max_val

Parameters:

Name	Type	Description	Default
max_value		Current maximum value of the image pixels. (usually 255)	255

Source code in src/super_gradients/training/transforms/transforms.py

@register_transform(Transforms.SegStandardize)
class SegStandardize(AbstractSegmentationTransform, LegacySegmentationTransformMixin):
    """
    Standardize image pixel values with img/max_val

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

    def __init__(self, max_value=255):
        self.max_value = max_value

    def apply_to_sample(self, sample: SegmentationSample) -> SegmentationSample:
        if not isinstance(sample.image, np.ndarray):
            sample.image = np.array(sample.image)
        sample.image = (sample.image / self.max_value).astype(np.float32)
        return sample

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

Standardize

Bases: torch.nn.Module

Standardize image pixel values.

Returns:

Type	Description
	max_val: float, value to as described above (default=255)

Source code in src/super_gradients/training/transforms/transforms.py

@register_transform(Transforms.Standardize)
class Standardize(torch.nn.Module):
    """
    Standardize image pixel values.
    :return img/max_val

    attributes:
        max_val: float, value to as described above (default=255)
    """

    def __init__(self, max_val=255.0):
        super(Standardize, self).__init__()
        self.max_val = max_val

    def forward(self, img):
        return img / self.max_val

get_affine_matrix(input_size, target_size, degrees=10, translate=0.1, scales=0.1, shear=10)

Return a random affine transform matrix.

Parameters:

Name	Type	Description	Default
input_size	Tuple[int, int]	Input shape.	required
target_size	Tuple[int, int]	Desired output shape.	required
degrees	Union[tuple, float]	Degrees for random rotation, when float the random values are drawn uniformly from (-degrees, degrees)	10
translate	Union[tuple, float]	Translate size (in pixels) for random translation, when float the random values are drawn uniformly from (-translate, translate)	0.1
scales	Union[tuple, float]	Values for random rescale, when float the random values are drawn uniformly from (1-scales, 1+scales)	0.1
shear	Union[tuple, float]	Degrees for random shear, when float the random values are drawn uniformly from (-shear, shear)	10

Returns:

Type	Description
np.ndarray	affine_transform_matrix, drawn_scale

Source code in src/super_gradients/training/transforms/transforms.py

def get_affine_matrix(
    input_size: Tuple[int, int],
    target_size: Tuple[int, int],
    degrees: Union[tuple, float] = 10,
    translate: Union[tuple, float] = 0.1,
    scales: Union[tuple, float] = 0.1,
    shear: Union[tuple, float] = 10,
) -> np.ndarray:
    """
    Return a random affine transform matrix.

    :param input_size:      Input shape.
    :param target_size:     Desired output shape.
    :param degrees:         Degrees for random rotation, when float the random values are drawn uniformly from (-degrees, degrees)
    :param translate:       Translate size (in pixels) for random translation, when float the random values are drawn uniformly from (-translate, translate)
    :param scales:          Values for random rescale, when float the random values are drawn uniformly from (1-scales, 1+scales)
    :param shear:           Degrees for random shear, when float the random values are drawn uniformly from (-shear, shear)

    :return: affine_transform_matrix, drawn_scale
    """

    # Center in pixels
    center_m = np.eye(3)
    center = (input_size[0] // 2, input_size[1] // 2)
    center_m[0, 2] = -center[1]
    center_m[1, 2] = -center[0]

    # Rotation and scale
    rotation_m = np.eye(3)
    rotation_m[:2] = cv2.getRotationMatrix2D(angle=get_aug_params(degrees), center=(0, 0), scale=get_aug_params(scales, center=1.0))

    # Shear in degrees
    shear_m = np.eye(3)
    shear_m[0, 1] = math.tan(get_aug_params(shear) * math.pi / 180)
    shear_m[1, 0] = math.tan(get_aug_params(shear) * math.pi / 180)

    # Translation in pixels
    translation_m = np.eye(3)
    translation_m[0, 2] = get_aug_params(translate, center=0.5) * target_size[1]
    translation_m[1, 2] = get_aug_params(translate, center=0.5) * target_size[0]

    return (translation_m @ shear_m @ rotation_m @ center_m)[:2]

get_aug_params(value, center=0)

Generates a random value for augmentations as described below

Parameters:

Name	Type	Description	Default
value	Union[tuple, float]	Range of values for generation. Wen tuple-drawn uniformly between (value[0], value[1]), and (center - value, center + value) when float.	required
center	float	Center to subtract when value is float.	0

Returns:

Type	Description
float	Generated value

Source code in src/super_gradients/training/transforms/transforms.py

def get_aug_params(value: Union[tuple, float], center: float = 0) -> float:
    """
    Generates a random value for augmentations as described below

    :param value:       Range of values for generation. Wen tuple-drawn uniformly between (value[0], value[1]), and (center - value, center + value) when float.
    :param center:      Center to subtract when value is float.
    :return:            Generated value
    """
    if isinstance(value, Number):
        return random.uniform(center - float(value), center + float(value))
    elif len(value) == 2:
        return random.uniform(value[0], value[1])
    else:
        raise ValueError(
            "Affine params should be either a sequence containing two values\
                          or single float values. Got {}".format(
                value
            )
        )

random_affine(img, targets=(), targets_seg=None, target_size=(640, 640), degrees=10, translate=0.1, scales=0.1, shear=10, filter_box_candidates=False, wh_thr=2, ar_thr=20, area_thr=0.1, border_value=114, crowd_targets=None)

Performs random affine transform to img, targets

Parameters:

Name	Type	Description	Default
img	np.ndarray	Input image of shape [h, w, c]	required
targets	np.ndarray	Input target	()
targets_seg	np.ndarray	Targets derived from segmentation masks	None
target_size	tuple	Desired output shape	(640, 640)
degrees	Union[float, tuple]	Degrees for random rotation, when float the random values are drawn uniformly from (-degrees, degrees).	10
translate	Union[float, tuple]	Translate size (in pixels) for random translation, when float the random values are drawn uniformly from (-translate, translate)	0.1
scales	Union[float, tuple]	Values for random rescale, when float the random values are drawn uniformly from (0.1-scales, 0.1+scales)	0.1
shear	Union[float, tuple]	Degrees for random shear, when float the random values are drawn uniformly from (shear, shear)	10
filter_box_candidates	bool	whether to filter out transformed bboxes by edge size, area ratio, and aspect ratio.	False
wh_thr		(float) edge size threshold when filter_box_candidates = True. Bounding oxes with edges smaller then this values will be filtered out. (default=2)	2
ar_thr		(float) aspect ratio threshold filter_box_candidates = True. Bounding boxes with aspect ratio larger then this values will be filtered out. (default=20)	20
area_thr		(float) threshold for area ratio between original image and the transformed one, when when filter_box_candidates = True. Bounding boxes with such ratio smaller then this value will be filtered out. (default=0.1)	0.1
border_value		value for filling borders after applying transforms (default=114).	114
crowd_targets	np.ndarray	Optional array of crowd annotations. If provided, it will be transformed in the same way as targets.	None

Returns:

Type	Description
	Image and Target with applied random affine

Source code in src/super_gradients/training/transforms/transforms.py

def random_affine(
    img: np.ndarray,
    targets: np.ndarray = (),
    targets_seg: np.ndarray = None,
    target_size: tuple = (640, 640),
    degrees: Union[float, tuple] = 10,
    translate: Union[float, tuple] = 0.1,
    scales: Union[float, tuple] = 0.1,
    shear: Union[float, tuple] = 10,
    filter_box_candidates: bool = False,
    wh_thr=2,
    ar_thr=20,
    area_thr=0.1,
    border_value=114,
    crowd_targets: np.ndarray = None,
):
    """
    Performs random affine transform to img, targets
    :param img:         Input image of shape [h, w, c]
    :param targets:     Input target
    :param targets_seg: Targets derived from segmentation masks
    :param target_size: Desired output shape
    :param degrees:     Degrees for random rotation, when float the random values are drawn uniformly
                            from (-degrees, degrees).
    :param translate:   Translate size (in pixels) for random translation, when float the random values
                            are drawn uniformly from (-translate, translate)
    :param scales:      Values for random rescale, when float the random values are drawn uniformly
                            from (0.1-scales, 0.1+scales)
    :param shear:       Degrees for random shear, when float the random values are drawn uniformly
                                from (shear, shear)

    :param filter_box_candidates:    whether to filter out transformed bboxes by edge size, area ratio, and aspect ratio.
    :param wh_thr: (float) edge size threshold when filter_box_candidates = True. Bounding oxes with edges smaller
      then this values will be filtered out. (default=2)

    :param ar_thr: (float) aspect ratio threshold filter_box_candidates = True. Bounding boxes with aspect ratio larger
      then this values will be filtered out. (default=20)

    :param area_thr:(float) threshold for area ratio between original image and the transformed one, when when filter_box_candidates = True.
      Bounding boxes with such ratio smaller then this value will be filtered out. (default=0.1)

    :param border_value: value for filling borders after applying transforms (default=114).
    :param crowd_targets: Optional array of crowd annotations. If provided, it will be transformed in the same way as targets.
    :return:            Image and Target with applied random affine
    """

    M = get_affine_matrix(img.shape[:2], target_size, degrees, translate, scales, shear)

    img = cv2.warpAffine(img, M, dsize=target_size, borderValue=(border_value, border_value, border_value))

    # Transform label coordinates
    if len(targets) > 0:
        targets_seg = np.zeros((targets.shape[0], 0)) if targets_seg is None else targets_seg
        targets_orig = targets.copy()
        targets = apply_affine_to_bboxes(targets, targets_seg, target_size, M)
        if filter_box_candidates:
            box_candidates_ids = _filter_box_candidates(targets_orig[:, :4], targets[:, :4], wh_thr=wh_thr, ar_thr=ar_thr, area_thr=area_thr)
            targets = targets[box_candidates_ids]

    if crowd_targets is not None:
        if len(crowd_targets) > 0:
            crowd_targets_seg = np.zeros((crowd_targets.shape[0], 0))
            crowd_targets_orig = crowd_targets.copy()
            crowd_targets = apply_affine_to_bboxes(crowd_targets, crowd_targets_seg, target_size, M)
            if filter_box_candidates:
                box_candidates_ids = _filter_box_candidates(crowd_targets_orig[:, :4], crowd_targets[:, :4], wh_thr=wh_thr, ar_thr=ar_thr, area_thr=area_thr)
                crowd_targets = crowd_targets[box_candidates_ids]
        return img, targets, crowd_targets

    return img, targets