Feature/sg 1488 yolo nas r integration sample and transforms

src/super_gradients/training/datasets/data_formats/obb/cxcywhr.py

@@ -0,0 +1,58 @@
+import cv2
+import numpy as np
+
+
+def cxcywhr_to_poly(boxes: np.ndarray) -> np.ndarray:
+    """
+    Convert oriented bounding boxes in CX-CY-W-H-R format to a polygon format
+    :param boxes: [N,...,5] oriented bounding boxes in CX-CY-W-H-R format
+    :return: [N,...,4, 2] oriented bounding boxes in polygon format
+    """
+    shape = boxes.shape
+    if shape[-1] != 5:
+        raise ValueError(f"Expected last dimension to be 5, got {shape[-1]}")
+
+    flat_rboxes = boxes.reshape(-1, 5).astype(np.float32)
+    polys = np.zeros((flat_rboxes.shape[0], 4, 2), dtype=np.float32)
+    for i, box in enumerate(flat_rboxes):
+        cx, cy, w, h, r = box
+        rect = ((cx, cy), (w, h), np.rad2deg(r))
+        poly = cv2.boxPoints(rect)
+        polys[i] = poly
+
+    return polys.reshape(*shape[:-1], 4, 2)
+
+
+def poly_to_cxcywhr(poly: np.ndarray) -> np.ndarray:
+    shape = poly.shape
+    if shape[-2:] != (4, 2):
+        raise ValueError(f"Expected last two dimensions to be (4, 2), got {shape[-2:]}")
+
+    flat_polys = poly.reshape(-1, 4, 2)
+    rboxes = np.zeros((flat_polys.shape[0], 5), dtype=np.float32)
+    for i, poly in enumerate(flat_polys):
+        hull = cv2.convexHull(np.reshape(poly, [-1, 2]).astype(np.float32))
+        rect = cv2.minAreaRect(hull)
+        cx, cy = rect[0]
+        w, h = rect[1]
+        angle = rect[2]
+        if h > w:
+            w, h = h, w
+            angle += 90
+        angle = np.deg2rad(angle)
+        rboxes[i] = [cx, cy, w, h, angle]
+
+    return rboxes.reshape(*shape[:-2], 5)
+
+
+def poly_to_xyxy(poly: np.ndarray) -> np.ndarray:
+    """
+    Convert oriented bounding boxes in polygon format to XYXY format
+    :param poly: [..., 4, 2]
+    :return: [..., 4]
+    """
+    x1 = poly[..., :, 0].min(axis=-1)
+    y1 = poly[..., :, 1].min(axis=-1)
+    x2 = poly[..., :, 0].max(axis=-1)
+    y2 = poly[..., :, 1].max(axis=-1)
+    return np.stack([x1, y1, x2, y2], axis=-1)

src/super_gradients/training/samples/__init__.py

@@ -2,5 +2,6 @@
 from .pose_estimation_sample import PoseEstimationSample
 from .detection_sample import DetectionSample
 from .segmentation_sample import SegmentationSample
+from .obb_sample import OBBSample
 
-__all__ = ["PoseEstimationSample", "DetectionSample", "SegmentationSample", "DepthEstimationSample"]
+__all__ = ["PoseEstimationSample", "DetectionSample", "SegmentationSample", "DepthEstimationSample", "OBBSample"]

src/super_gradients/training/samples/obb_sample.py

@@ -0,0 +1,110 @@
+import dataclasses
+from typing import Union, List, Optional
+
+import numpy as np
+import torch
+from super_gradients.training.datasets.data_formats.obb.cxcywhr import cxcywhr_to_poly, poly_to_cxcywhr
+
+
+@dataclasses.dataclass
+class OBBSample:
+    """
+    A data class describing a single object detection sample that comes from a dataset.
+    It contains both input image and target information to train an object detection model.
+
+    :param image:              Associated image with a sample. Can be in [H,W,C] or [C,H,W] format
+    :param rboxes_cxcywhr:      Numpy array of [N,5] shape with oriented bounding box of each instance (CX,CY,W,H,R)
+    :param labels:             Numpy array of [N] shape with class label for each instance
+    :param is_crowd:           (Optional) Numpy array of [N] shape with is_crowd flag for each instance
+    :param additional_samples: (Optional) List of additional samples for the same image.
+    """
+
+    __slots__ = ["image", "rboxes_cxcywhr", "labels", "is_crowd", "additional_samples"]
+
+    image: Union[np.ndarray, torch.Tensor]
+    rboxes_cxcywhr: np.ndarray
+    labels: np.ndarray
+    is_crowd: np.ndarray
+    additional_samples: Optional[List["OBBSample"]]
+
+    def __init__(
+        self,
+        image: Union[np.ndarray, torch.Tensor],
+        rboxes_cxcywhr: np.ndarray,
+        labels: np.ndarray,
+        is_crowd: Optional[np.ndarray] = None,
+        additional_samples: Optional[List["OBBSample"]] = None,
+    ):
+        if is_crowd is None:
+            is_crowd = np.zeros(len(labels), dtype=bool)
+
+        if len(rboxes_cxcywhr) != len(labels):
+            raise ValueError("Number of bounding boxes and labels must be equal. Got {len(bboxes_xyxy)} and {len(labels)} respectively")
+
+        if len(rboxes_cxcywhr) != len(is_crowd):
+            raise ValueError("Number of bounding boxes and is_crowd flags must be equal. Got {len(bboxes_xyxy)} and {len(is_crowd)} respectively")
+
+        if len(rboxes_cxcywhr.shape) != 2 or rboxes_cxcywhr.shape[1] != 5:
+            raise ValueError(f"Oriented boxes must be in [N,5] format. Shape of input bboxes is {rboxes_cxcywhr.shape}")
+
+        if len(is_crowd.shape) != 1:
+            raise ValueError(f"Number of is_crowd flags must be in [N] format. Shape of input is_crowd is {is_crowd.shape}")
+
+        if len(labels.shape) != 1:
+            raise ValueError("Labels must be in [N] format. Shape of input labels is {labels.shape}")
+
+        self.image = image
+        self.rboxes_cxcywhr = rboxes_cxcywhr
+        self.labels = labels
+        self.is_crowd = is_crowd
+        self.additional_samples = additional_samples
+
+    def sanitize_sample(self) -> "OBBSample":
+        """
+        Apply sanity checks on the detection sample, which includes clamping of rotate boxes to image boundaries
+        and removing boxes with non-positive area.
+        This method returns a new DetectionSample instance with sanitized data.
+        :return: A DetectionSample after filtering.
+        """
+        polys = cxcywhr_to_poly(self.rboxes_cxcywhr)
+        # Clamp polygons to image boundaries
+        polys[..., 0] = np.clip(polys[..., 0], 0, self.image.shape[1])
+        polys[..., 1] = np.clip(polys[..., 1], 0, self.image.shape[0])
+        rboxes_cxcywhr = poly_to_cxcywhr(polys)
+        return OBBSample(
+            image=self.image,
+            rboxes_cxcywhr=rboxes_cxcywhr,
+            labels=self.labels,
+            is_crowd=self.is_crowd,
+            additional_samples=self.additional_samples,
+        ).filter_by_bbox_area(0)
+
+    def filter_by_mask(self, mask: np.ndarray) -> "OBBSample":
+        """
+        Remove boxes & labels with respect to a given mask.
+        This method returns a new DetectionSample instance with filtered data.
+
+        :param mask:   A boolean or integer mask of samples to keep for given sample.
+        :return:       A DetectionSample after filtering.
+        """
+        return OBBSample(
+            image=self.image,
+            rboxes_cxcywhr=self.rboxes_cxcywhr[mask],
+            labels=self.labels[mask],
+            is_crowd=self.is_crowd[mask] if self.is_crowd is not None else None,
+            additional_samples=self.additional_samples,
+        )
+
+    def filter_by_bbox_area(self, min_rbox_area: Union[int, float]) -> "OBBSample":
+        """
+        Remove pose instances that has area of the corresponding bounding box less than a certain threshold.
+
+        :param min_rbox_area: Minimal rotated box area of the box to keep.
+        :return:              A OBBSample after filtering.
+        """
+        area = self.rboxes_cxcywhr[..., 2:4].prod(axis=-1)
+        keep_mask = area > min_rbox_area
+        return self.filter_by_mask(keep_mask)
+
+    def __len__(self):
+        return len(self.labels)

src/super_gradients/training/transforms/__init__.py

@@ -38,6 +38,16 @@
 from super_gradients.common.registry.albumentation import ALBUMENTATIONS_TRANSFORMS, ALBUMENTATIONS_COMP_TRANSFORMS, imported_albumentations_failure
 from super_gradients.training.transforms.detection import AbstractDetectionTransform, DetectionPadIfNeeded, DetectionLongestMaxSize
 
+from .obb import (
+    AbstractOBBDetectionTransform,
+    OBBDetectionPadIfNeeded,
+    OBBDetectionLongestMaxSize,
+    OBBDetectionStandardize,
+    OBBDetectionMixup,
+    OBBDetectionCompose,
+    OBBRemoveSmallObjects,
+)
+
 __all__ = [
     "TRANSFORMS",
     "ALBUMENTATIONS_TRANSFORMS",
@@ -76,6 +86,13 @@
     "DetectionPadIfNeeded",
     "DetectionLongestMaxSize",
     "AbstractDetectionTransform",
+    "AbstractOBBDetectionTransform",
+    "OBBDetectionPadIfNeeded",
+    "OBBDetectionLongestMaxSize",
+    "OBBDetectionStandardize",
+    "OBBDetectionMixup",
+    "OBBDetectionCompose",
+    "OBBRemoveSmallObjects",
 ]
 
 cv2.setNumThreads(0)

src/super_gradients/training/transforms/obb/__init__.py

@@ -0,0 +1,19 @@
+from super_gradients.training.samples.obb_sample import OBBSample
+from .abstract_obb_transform import AbstractOBBDetectionTransform
+from .obb_pad_if_needed import OBBDetectionPadIfNeeded
+from .obb_longest_max_size import OBBDetectionLongestMaxSize
+from .obb_standardize import OBBDetectionStandardize
+from .obb_mixup import OBBDetectionMixup
+from .obb_compose import OBBDetectionCompose
+from .obb_remove_small_objects import OBBRemoveSmallObjects
+
+__all__ = [
+    "OBBSample",
+    "AbstractOBBDetectionTransform",
+    "OBBDetectionPadIfNeeded",
+    "OBBDetectionLongestMaxSize",
+    "OBBDetectionStandardize",
+    "OBBDetectionMixup",
+    "OBBDetectionCompose",
+    "OBBRemoveSmallObjects",
+]

src/super_gradients/training/transforms/obb/abstract_obb_transform.py

@@ -0,0 +1,59 @@
+import abc
+import warnings
+from abc import abstractmethod
+from typing import List
+
+from super_gradients.training.samples.obb_sample import OBBSample
+
+__all__ = ["AbstractOBBDetectionTransform"]
+
+
+class AbstractOBBDetectionTransform(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: OBBSample) -> OBBSample:
+        """
+        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
+
+    @property
+    def additional_samples_count(self) -> int:
+        warnings.warn(
+            "This property is deprecated and will be removed in the future." "Please use `get_number_of_additional_samples` instead.", DeprecationWarning
+        )
+        return self.get_number_of_additional_samples()
+
+    def get_number_of_additional_samples(self) -> int:
+        """
+        Returns number of additional samples required. The default implementation assumes that this number is fixed and deterministic.
+        Override in case this is not the case, e.g., you randomly choose to apply MixUp, etc
+        """
+        return self._additional_samples_count
+
+    @property
+    def may_require_additional_samples(self) -> bool:
+        """
+        Indicates whether additional samples are required. The default implementation assumes that this indicator is fixed and deterministic.
+        Override in case this is not the case, e.g., you randomly choose to apply MixUp, etc
+        """
+        return self._additional_samples_count > 0
+
+    @abstractmethod
+    def get_equivalent_preprocessing(self) -> List:
+        raise NotImplementedError

src/super_gradients/training/transforms/obb/obb_compose.py

@@ -0,0 +1,93 @@
+from typing import List
+
+from .abstract_obb_transform import AbstractOBBDetectionTransform
+from super_gradients.training.samples.obb_sample import OBBSample
+
+
+class OBBDetectionCompose(AbstractOBBDetectionTransform):
+    """
+    Composes several transforms together
+    """
+
+    def __init__(self, transforms: List[AbstractOBBDetectionTransform], 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 hasattr(transform, "may_require_additional_samples") and transform.may_require_additional_samples and load_sample_fn is None:
+                raise RuntimeError(f"Transform {transform.__class__.__name__} that requires additional samples but `load_sample_fn` is None")
+
+        super().__init__()
+        self.transforms = transforms
+        self.load_sample_fn = load_sample_fn
+
+    def apply_to_sample(self, sample: OBBSample) -> OBBSample:
+        """
+        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: OBBSample, transforms: List[AbstractOBBDetectionTransform], load_sample_fn) -> OBBSample:
+        """
+        This helper method allows us to query additional samples for mixup & mosaic augmentations
+        that would be also passed through augmentation pipeline. Example:
+
+        ```
+          transforms:
+            - OBBDetectionLongestMaxSize:
+                max_height: ${dataset_params.image_size}
+                max_width: ${dataset_params.image_size}
+            - OBBDetectionMixup:
+                prob: ${dataset_params.mixup_prob}
+        ```
+
+        In the example above all samples in mixup will be forwarded through OBBDetectionLongestMaxSize,
+        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, "may_require_additional_samples") or not t.may_require_additional_samples:
+                sample = t.apply_to_sample(sample)
+                applied_transforms_so_far.append(t)
+            else:
+                additional_samples = [load_sample_fn() for _ in range(t.get_number_of_additional_samples())]
+                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