Remove __target__ for Detection CollateFN

documentation/source/ObjectDetection.md

@@ -500,8 +500,7 @@ train_dataloader_params:
   worker_init_fn:
     _target_: super_gradients.training.utils.utils.load_func
     dotpath: super_gradients.training.datasets.datasets_utils.worker_init_reset_seed
-  collate_fn:
-    _target_: super_gradients.training.utils.detection_utils.DetectionCollateFN
+  collate_fn: DetectionCollateFN
 
 val_dataset_params:
   data_dir: ${dataset_params.root_dir}
@@ -521,8 +520,7 @@ val_dataloader_params:
   num_workers: 8
   drop_last: True
   pin_memory: True
-  collate_fn:
-    _target_: super_gradients.training.utils.detection_utils.DetectionCollateFN
+  collate_fn: DetectionCollateFN
 ```
 
 In your training recipe add/change the following lines to:

src/super_gradients/common/exceptions/dataset_exceptions.py

@@ -1,3 +1,6 @@
+from typing import Tuple, Type
+
+
 class DatasetValidationException(Exception):
     pass
 
@@ -46,3 +49,16 @@ def __init__(self, batch_items: tuple):
             "To fix this, please change the implementation of your dataset __getitem__ method, so that it would return the items defined above.\n"
         )
         super().__init__(self.message)
+
+
+class DatasetItemsException(Exception):
+    def __init__(self, data_sample: Tuple, collate_type: Type, expected_item_names: Tuple):
+        """
+        :param data_sample:         item(s) returned by a dataset
+        :param collate_type:        type of the collate that caused the exception
+        :param expected_item_names: tuple of names of items that are expected by the collate to be returned from the dataset
+        """
+        collate_type_name = collate_type.__name__
+        num_sample_items = len(data_sample) if isinstance(data_sample, tuple) else 1
+        error_msg = f"`{collate_type_name}` only supports Datasets that return a tuple {expected_item_names}, but got a tuple of len={num_sample_items}"
+        super().__init__(error_msg)

src/super_gradients/recipes/dataset_params/coco_detection_dataset_params.yaml

@@ -52,8 +52,7 @@ train_dataloader_params:
   worker_init_fn:
     _target_: super_gradients.training.utils.utils.load_func
     dotpath: super_gradients.training.datasets.datasets_utils.worker_init_reset_seed
-  collate_fn: # collate function for trainset
-    _target_: super_gradients.training.utils.detection_utils.DetectionCollateFN
+  collate_fn: DetectionCollateFN
 
 val_dataset_params:
   data_dir: /data/coco # root path to coco data
@@ -80,8 +79,7 @@ val_dataloader_params:
   num_workers: 8
   drop_last: False
   pin_memory: True
-  collate_fn: # collate function for valset
-    _target_: super_gradients.training.utils.detection_utils.CrowdDetectionCollateFN
+  collate_fn: CrowdDetectionCollateFN
 
 
 _convert_: all

src/super_gradients/recipes/dataset_params/coco_detection_ppyoloe_dataset_params.yaml

@@ -55,15 +55,15 @@ train_dataloader_params:
   worker_init_fn:
     _target_: super_gradients.training.utils.utils.load_func
     dotpath: super_gradients.training.datasets.datasets_utils.worker_init_reset_seed
-  collate_fn: # collate function for trainset
-    _target_: super_gradients.training.utils.detection_utils.PPYoloECollateFN
-    random_resize_sizes: [ 320, 352, 384, 416, 448, 480, 512, 544, 576, 608, 640, 672, 704, 736, 768 ]
-    random_resize_modes:
-      - 0 # cv::INTER_NEAREST
-      - 1 # cv::INTER_LINEAR
-      - 2 # cv::INTER_CUBIC
-      - 3 # cv::INTER_AREA
-      - 4 # cv::INTER_LANCZOS4
+  collate_fn:
+    PPYoloECollateFN:
+      random_resize_sizes: [ 320, 352, 384, 416, 448, 480, 512, 544, 576, 608, 640, 672, 704, 736, 768 ]
+      random_resize_modes:
+        - 0 # cv::INTER_NEAREST
+        - 1 # cv::INTER_LINEAR
+        - 2 # cv::INTER_CUBIC
+        - 3 # cv::INTER_AREA
+        - 4 # cv::INTER_LANCZOS4
 
 val_dataset_params:
   data_dir: /data/coco # root path to coco data
@@ -93,7 +93,6 @@ val_dataloader_params:
   drop_last: False
   shuffle: False
   pin_memory: False
-  collate_fn: # collate function for valset
-    _target_: super_gradients.training.utils.detection_utils.CrowdDetectionPPYoloECollateFN
+  collate_fn: CrowdDetectionPPYoloECollateFN
 
 _convert_: all

src/super_gradients/recipes/dataset_params/coco_detection_ssd_lite_mobilenet_v2_dataset_params.yaml

@@ -48,8 +48,7 @@ train_dataloader_params:
   worker_init_fn:
     _target_: super_gradients.training.utils.utils.load_func
     dotpath: super_gradients.training.datasets.datasets_utils.worker_init_reset_seed
-  collate_fn: # collate function for trainset
-    _target_: super_gradients.training.utils.detection_utils.DetectionCollateFN
+  collate_fn: DetectionCollateFN
 
 val_dataset_params:
   data_dir: /data/coco # root path to coco data
@@ -76,7 +75,6 @@ val_dataloader_params:
   num_workers: 8
   drop_last: False
   pin_memory: True
-  collate_fn: # collate function for valset
-    _target_: super_gradients.training.utils.detection_utils.CrowdDetectionCollateFN
+  collate_fn: CrowdDetectionCollateFN
 
 _convert_: all

src/super_gradients/recipes/dataset_params/coco_detection_yolo_format_base_dataset_params.yaml

@@ -55,8 +55,7 @@ train_dataloader_params:
   shuffle: True
   drop_last: True
   pin_memory: True
-  collate_fn:
-    _target_: super_gradients.training.utils.detection_utils.DetectionCollateFN
+  collate_fn: DetectionCollateFN
 
 val_dataset_params:
   data_dir: /data/coco # TO FILL: Where the data is stored.
@@ -88,7 +87,6 @@ val_dataloader_params:
   num_workers: 8
   drop_last: False
   pin_memory: True
-  collate_fn:
-    _target_: super_gradients.training.utils.detection_utils.DetectionCollateFN
+  collate_fn: DetectionCollateFN
 
 _convert_: all

src/super_gradients/recipes/dataset_params/coco_detection_yolo_nas_dataset_params.yaml

@@ -51,8 +51,7 @@ train_dataloader_params:
   shuffle: True
   drop_last: True
   pin_memory: True
-  collate_fn:
-    _target_: super_gradients.training.utils.detection_utils.DetectionCollateFN
+  collate_fn: DetectionCollateFN
 
 val_dataset_params:
   data_dir: /data/coco # root path to coco data
@@ -86,7 +85,6 @@ val_dataloader_params:
   drop_last: False
   shuffle: False
   pin_memory: True
-  collate_fn:
-    _target_: super_gradients.training.utils.detection_utils.CrowdDetectionCollateFN
+  collate_fn: CrowdDetectionCollateFN
 
 _convert_: all

src/super_gradients/recipes/dataset_params/pascal_voc_detection_dataset_params.yaml

@@ -39,16 +39,14 @@ train_dataloader_params:
   worker_init_fn:
     _target_: super_gradients.training.utils.utils.load_func
     dotpath: super_gradients.training.datasets.datasets_utils.worker_init_reset_seed
-  collate_fn: # collate function for trainset
-    _target_: super_gradients.training.utils.detection_utils.DetectionCollateFN
+  collate_fn: DetectionCollateFN
 
 val_dataloader_params:
   batch_size: 64
   num_workers: 8
   drop_last: False
   pin_memory: True
-  collate_fn: # collate function for trainset
-    _target_: super_gradients.training.utils.detection_utils.DetectionCollateFN
+  collate_fn: DetectionCollateFN
 
 
 _convert_: all

src/super_gradients/recipes/dataset_params/roboflow_detection_dataset_params.yaml

@@ -61,8 +61,7 @@ train_dataloader_params:
   worker_init_fn:
     _target_: super_gradients.training.utils.utils.load_func
     dotpath: super_gradients.training.datasets.datasets_utils.worker_init_reset_seed
-  collate_fn: # collate function for trainset
-    _target_: super_gradients.training.utils.detection_utils.DetectionCollateFN
+  collate_fn: DetectionCollateFN
 
 val_dataset_params:
   data_dir: ${..data_dir} # root path to Robflow datasets
@@ -94,8 +93,7 @@ val_dataloader_params:
   drop_last: False
   shuffle: False
   pin_memory: True
-  collate_fn: # collate function for valset
-    _target_: super_gradients.training.utils.detection_utils.CrowdDetectionCollateFN
+  collate_fn: CrowdDetectionCollateFN
 
 
 _convert_: all

src/super_gradients/training/utils/collate_fn/__init__.py

@@ -0,0 +1,6 @@
+from .detection_collate_fn import DetectionCollateFN
+from .ppyoloe_collate_fn import PPYoloECollateFN
+from .crowd_detection_collate_fn import CrowdDetectionCollateFN
+from .crowd_detection_ppyoloe_collate_fn import CrowdDetectionPPYoloECollateFN
+
+__all__ = ["DetectionCollateFN", "PPYoloECollateFN", "CrowdDetectionCollateFN", "CrowdDetectionPPYoloECollateFN"]

src/super_gradients/training/utils/collate_fn/crowd_detection_collate_fn.py

@@ -0,0 +1,26 @@
+from typing import Tuple, Dict
+
+import torch
+
+from super_gradients.common.registry import register_collate_function
+from super_gradients.common.exceptions.dataset_exceptions import DatasetItemsException
+from super_gradients.training.utils.collate_fn.detection_collate_fn import DetectionCollateFN
+
+
+@register_collate_function()
+class CrowdDetectionCollateFN(DetectionCollateFN):
+    """
+    Collate function for Yolox training with additional_batch_items that includes crowd targets
+    """
+
+    def __init__(self):
+        self.expected_item_names = ("image", "targets", "crowd_targets")
+        super().__init__()
+
+    def _format_batch(self, batch_data) -> Tuple[torch.Tensor, torch.Tensor, Dict[str, torch.Tensor]]:
+        try:
+            images_batch, labels_batch, crowd_labels_batch = batch_data
+        except (ValueError, TypeError):
+            raise DatasetItemsException(data_sample=batch_data, collate_type=type(self), expected_item_names=self.expected_item_names)
+
+        return self._format_images(images_batch), self._format_targets(labels_batch), {"crowd_targets": self._format_targets(crowd_labels_batch)}

src/super_gradients/training/utils/collate_fn/crowd_detection_ppyoloe_collate_fn.py

@@ -0,0 +1,30 @@
+from typing import Union, List, Tuple, Dict
+
+import torch
+
+from super_gradients.common.registry import register_collate_function
+from super_gradients.common.exceptions.dataset_exceptions import DatasetItemsException
+from super_gradients.training.utils.collate_fn.ppyoloe_collate_fn import PPYoloECollateFN
+
+
+@register_collate_function()
+class CrowdDetectionPPYoloECollateFN(PPYoloECollateFN):
+    """
+    Collate function for Yolox training with additional_batch_items that includes crowd targets
+    """
+
+    def __init__(self, random_resize_sizes: Union[List[int], None] = None, random_resize_modes: Union[List[int], None] = None):
+        super().__init__(random_resize_sizes, random_resize_modes)
+        self.expected_item_names = ("image", "targets", "crowd_targets")
+
+    def __call__(self, data) -> Tuple[torch.Tensor, torch.Tensor, Dict[str, torch.Tensor]]:
+
+        if self.random_resize_sizes is not None:
+            data = self.random_resize(data)
+
+        try:
+            images_batch, labels_batch, crowd_labels_batch = list(zip(*data))
+        except (ValueError, TypeError):
+            raise DatasetItemsException(data_sample=data[0], collate_type=type(self), expected_item_names=self.expected_item_names)
+
+        return self._format_images(images_batch), self._format_targets(labels_batch), {"crowd_targets": self._format_targets(crowd_labels_batch)}

src/super_gradients/training/utils/collate_fn/detection_collate_fn.py

@@ -0,0 +1,47 @@
+from typing import Tuple, List, Union
+
+import numpy as np
+import torch
+
+from super_gradients.common.registry import register_collate_function
+from super_gradients.common.exceptions.dataset_exceptions import DatasetItemsException
+
+
+@register_collate_function()
+class DetectionCollateFN:
+    """
+    Collate function for Yolox training
+    """
+
+    def __init__(self):
+        self.expected_item_names = ("image", "targets")
+
+    def __call__(self, data) -> Tuple[torch.Tensor, torch.Tensor]:
+        try:
+            images_batch, labels_batch = list(zip(*data))
+        except (ValueError, TypeError):
+            raise DatasetItemsException(data_sample=data[0], collate_type=type(self), expected_item_names=self.expected_item_names)
+
+        return self._format_images(images_batch), self._format_targets(labels_batch)
+
+    def _format_images(self, images_batch: List[Union[torch.Tensor, np.array]]) -> torch.Tensor:
+        images_batch = [torch.tensor(img) for img in images_batch]
+        images_batch_stack = torch.stack(images_batch, 0)
+        if images_batch_stack.shape[3] == 3:
+            images_batch_stack = torch.moveaxis(images_batch_stack, -1, 1).float()
+        return images_batch_stack
+
+    def _format_targets(self, labels_batch: List[Union[torch.Tensor, np.array]]) -> torch.Tensor:
+        """
+        Stack a batch id column to targets and concatenate
+        :param labels_batch: a list of targets per image (each of arbitrary length)
+        :return: one tensor of targets of all imahes of shape [N, 6], where N is the total number of targets in a batch
+                 and the 1st column is batch item index
+        """
+        labels_batch = [torch.tensor(labels) for labels in labels_batch]
+        labels_batch_indexed = []
+        for i, labels in enumerate(labels_batch):
+            batch_column = labels.new_ones((labels.shape[0], 1)) * i
+            labels = torch.cat((batch_column, labels), dim=-1)
+            labels_batch_indexed.append(labels)
+        return torch.cat(labels_batch_indexed, 0)

src/super_gradients/training/utils/collate_fn/ppyoloe_collate_fn.py

@@ -0,0 +1,69 @@
+import random
+from typing import Union, List, Tuple
+
+import cv2
+import numpy as np
+import torch
+
+from super_gradients.common.registry import register_collate_function
+from super_gradients.common.exceptions.dataset_exceptions import DatasetItemsException
+from super_gradients.training.utils.collate_fn.detection_collate_fn import DetectionCollateFN
+
+
+@register_collate_function()
+class PPYoloECollateFN(DetectionCollateFN):
+    """
+    Collate function for PPYoloE training
+    """
+
+    def __init__(self, random_resize_sizes: Union[List[int], None] = None, random_resize_modes: Union[List[int], None] = None):
+        """
+        :param random_resize_sizes: (rows, cols)
+        """
+        super().__init__()
+        self.random_resize_sizes = random_resize_sizes
+        self.random_resize_modes = random_resize_modes
+
+    def __repr__(self):
+        return f"PPYoloECollateFN(random_resize_sizes={self.random_resize_sizes}, random_resize_modes={self.random_resize_modes})"
+
+    def __str__(self):
+        return self.__repr__()
+
+    def __call__(self, data) -> Tuple[torch.Tensor, torch.Tensor]:
+        if self.random_resize_sizes is not None:
+            data = self.random_resize(data)
+        return super().__call__(data)
+
+    def random_resize(self, batch):
+        target_size = random.choice(self.random_resize_sizes)
+        interpolation = random.choice(self.random_resize_modes)
+        batch = [self.random_resize_sample(sample, target_size, interpolation) for sample in batch]
+        return batch
+
+    def random_resize_sample(self, sample, target_size, interpolation):
+        if len(sample) == 2:
+            image, targets = sample  # TARGETS ARE IN LABEL_CXCYWH
+            with_crowd = False
+        elif len(sample) == 3:
+            image, targets, crowd_targets = sample
+            with_crowd = True
+        else:
+            raise DatasetItemsException(data_sample=sample, collate_type=type(self), expected_item_names=self.expected_item_names)
+
+        dsize = int(target_size), int(target_size)
+        scale_factors = target_size / image.shape[0], target_size / image.shape[1]
+
+        image = cv2.resize(
+            image,
+            dsize=dsize,
+            interpolation=interpolation,
+        )
+
+        sy, sx = scale_factors
+        targets[:, 1:5] *= np.array([[sx, sy, sx, sy]], dtype=targets.dtype)
+        if with_crowd:
+            crowd_targets[:, 1:5] *= np.array([[sx, sy, sx, sy]], dtype=targets.dtype)
+            return image, targets, crowd_targets
+
+        return image, targets