@@ -1,12 +1,71 @@
 import unittest
-from super_gradients.training.datasets import Cifar10DatasetInterface
+
+from super_gradients.training.datasets.dataset_interfaces.dataset_interface import PascalVOCUnifiedDetectionDatasetInterface
+from super_gradients.training.transforms.transforms import DetectionPaddedRescale, DetectionTargetsFormatTransform, DetectionMosaic, DetectionRandomAffine,\
+    DetectionHSV
+from super_gradients.training.utils.detection_utils import DetectionTargetsFormat
+from super_gradients.training.utils.detection_utils import DetectionCollateFN
+from super_gradients.training.utils import sg_model_utils
+from super_gradients.training import utils as core_utils
 
 
 class TestDatasetInterface(unittest.TestCase):
-    def test_cifar(self):
-        test_dataset_interface = Cifar10DatasetInterface()
-        cifar_dataset_sample = test_dataset_interface.get_test_sample()
-        self.assertListEqual([3, 32, 32], list(cifar_dataset_sample[0].shape))
+    def setUp(self) -> None:
+        self.root_dir = "/home/data/"
+        self.train_batch_size, self.val_batch_size = 16, 32
+        self.train_image_size, self.val_image_size = 640, 640
+        self.train_input_dim = (self.train_image_size, self.train_image_size)
+        self.val_input_dim = (self.val_image_size, self.val_image_size)
+        self.train_max_num_samples = 100
+        self.val_max_num_samples = 90
+
+    def setup_pascal_voc_interface(self):
+        """setup PascalVOCUnifiedDetectionDataSetInterfaceV2 and return dataloaders"""
+        dataset_params = {
+            "data_dir": self.root_dir + "pascal_unified_coco_format/",
+            "cache_dir": self.root_dir + "pascal_unified_coco_format/",
+            "batch_size": self.train_batch_size,
+            "val_batch_size": self.val_batch_size,
+            "train_image_size": self.train_image_size,
+            "val_image_size": self.val_image_size,
+            "train_max_num_samples": self.train_max_num_samples,
+            "val_max_num_samples": self.val_max_num_samples,
+            "train_transforms": [
+                DetectionMosaic(input_dim=self.train_input_dim, prob=1),
+                DetectionRandomAffine(degrees=0.373, translate=0.245, scales=0.898, shear=0.602, target_size=self.train_input_dim),
+                DetectionHSV(prob=1, hgain=0.0138, sgain=0.664, vgain=0.464),
+                DetectionPaddedRescale(input_dim=self.train_input_dim, max_targets=100),
+                DetectionTargetsFormatTransform(input_format=DetectionTargetsFormat.XYXY_LABEL,
+                                                output_format=DetectionTargetsFormat.LABEL_CXCYWH)],
+            "val_transforms": [
+                DetectionPaddedRescale(input_dim=self.val_input_dim),
+                DetectionTargetsFormatTransform(input_format=DetectionTargetsFormat.XYXY_LABEL,
+                                                output_format=DetectionTargetsFormat.LABEL_CXCYWH)],
+            "train_collate_fn": DetectionCollateFN(),
+            "val_collate_fn": DetectionCollateFN(),
+            "download": False,
+            "cache_train_images": False,
+            "cache_val_images": False,
+            "class_inclusion_list": ["person"]
+        }
+        dataset_interface = PascalVOCUnifiedDetectionDatasetInterface(dataset_params=dataset_params)
+        train_loader, valid_loader, _test_loader, _classes = dataset_interface.get_data_loaders()
+        return train_loader, valid_loader
+
+    def test_pascal_voc(self):
+        """Check that the dataset interface is correctly instantiated, and that the batch items are of expected size"""
+        train_loader, valid_loader = self.setup_pascal_voc_interface()
+
+        for loader, batch_size, image_size, max_num_samples in [(train_loader, self.train_batch_size, self.train_image_size, self.train_max_num_samples),
+                                                                (valid_loader, self.val_batch_size, self.val_image_size, self.val_max_num_samples)]:
+            # The dataset is at most of length max_num_samples, but can be smaller if not enough samples
+            self.assertGreaterEqual(max_num_samples, len(loader.dataset))
+
+            batch_items = next(iter(loader))
+            batch_items = core_utils.tensor_container_to_device(batch_items, 'cuda', non_blocking=True)
+
+            inputs, targets, additional_batch_items = sg_model_utils.unpack_batch_items(batch_items)
+            self.assertListEqual([batch_size, 3, image_size, image_size], list(inputs.shape))
 
 
 if __name__ == '__main__':

import unittest
import numpy as np

from super_gradients.training.datasets import DetectionDataset
from super_gradients.training.utils.detection_utils import DetectionTargetsFormat
from super_gradients.training.exceptions.dataset_exceptions import EmptyDatasetException


class DummyDetectionDataset(DetectionDataset):
    def __init__(self, input_dim, *args, **kwargs):
        """Dummy Dataset testing subclassing, designed with no annotation that includes class_2."""

        self.dummy_targets = [np.array([[0, 0, 10, 10, 0],
                                        [0, 5, 10, 15, 0],
                                        [0, 5, 15, 20, 0]]),
                              np.array([[0, 0, 10, 10, 0],
                                        [0, 5, 10, 15, 0],
                                        [0, 15, 55, 20, 1]])]

        self.image_size = input_dim
        kwargs['all_classes_list'] = ["class_0", "class_1", "class_2"]
        kwargs['original_target_format'] = DetectionTargetsFormat.XYXY_LABEL
        super().__init__(data_dir='', input_dim=input_dim, *args, **kwargs)

    def _setup_data_source(self):
        return len(self.dummy_targets)

    def _load_annotation(self, sample_id: int) -> dict:
        """Load 2 different annotations.
            - Annotation 0 is made of: 3 targets of class 0, 0 of class_1 and 0 of class_2
            - Annotation 1 is made of: 2 targets of class_0, 1 of class_1 and 0 of class_2
        """
        return {"img_path": "", "target": self.dummy_targets[sample_id]}

    # DetectionDatasetV2 will call _load_image but since we don't have any image we patch this method with
    # tensor of image shape
    def _load_image(self, index: int) -> np.ndarray:
        return np.random.random(self.image_size)


class TestDetectionDatasetSubclassing(unittest.TestCase):
    def setUp(self) -> None:
        self.config_keep_empty_annotation = [
            {
                "class_inclusion_list": ["class_0", "class_1", "class_2"],
                "expected_n_targets_after_subclass": [3, 3]
            },
            {
                "class_inclusion_list": ["class_0"],
                "expected_n_targets_after_subclass": [3, 2]
            },
            {
                "class_inclusion_list": ["class_1"],
                "expected_n_targets_after_subclass": [0, 1]
            },
            {
                "class_inclusion_list": ["class_2"],
                "expected_n_targets_after_subclass": [0, 0]
            },
        ]
        self.config_ignore_empty_annotation = [
            {
                "class_inclusion_list": ["class_0", "class_1", "class_2"],
                "expected_n_targets_after_subclass": [3, 3]
            },
            {
                "class_inclusion_list": ["class_0"],
                "expected_n_targets_after_subclass": [3, 2]
            },
            {
                "class_inclusion_list": ["class_1"],
                "expected_n_targets_after_subclass": [1]
            }
        ]

    def test_subclass_keep_empty(self):
        """Check that subclassing only keeps annotations of wanted class"""
        for config in self.config_keep_empty_annotation:
            test_dataset = DummyDetectionDataset(input_dim=(640, 512), ignore_empty_annotations=False,
                                                 class_inclusion_list=config["class_inclusion_list"])
            n_targets_after_subclass = _count_targets_after_subclass_per_index(test_dataset)
            self.assertListEqual(config["expected_n_targets_after_subclass"], n_targets_after_subclass)

    def test_subclass_drop_empty(self):
        """Check that empty annotations are not indexed (i.e. ignored) when ignore_empty_annotations=True"""
        for config in self.config_ignore_empty_annotation:
            test_dataset = DummyDetectionDataset(input_dim=(640, 512), ignore_empty_annotations=True,
                                                 class_inclusion_list=config["class_inclusion_list"])
            n_targets_after_subclass = _count_targets_after_subclass_per_index(test_dataset)
            self.assertListEqual(config["expected_n_targets_after_subclass"], n_targets_after_subclass)

        # Check last case when class_2, which should raise EmptyDatasetException because not a single image has
        # a target in class_inclusion_list
        with self.assertRaises(EmptyDatasetException):
            DummyDetectionDataset(input_dim=(640, 512), ignore_empty_annotations=True,
                                  class_inclusion_list=["class_2"])

    def test_wrong_subclass(self):
        """Check that ValueError is raised when class_inclusion_list includes a class that does not exist."""
        with self.assertRaises(ValueError):
            DummyDetectionDataset(input_dim=(640, 512), class_inclusion_list=["non_existing_class"])
        with self.assertRaises(ValueError):
            DummyDetectionDataset(input_dim=(640, 512), class_inclusion_list=["class_0", "non_existing_class"])


def _count_targets_after_subclass_per_index(test_dataset: DummyDetectionDataset):
    """Iterate through every index of the dataset and count the associated number of targets per index"""
    dataset_target_len = []
    for index in range(len(test_dataset)):
        _img, targets = test_dataset[index]
        dataset_target_len.append(len(targets))
    return dataset_target_len


if __name__ == '__main__':
    unittest.main()

import unittest
import numpy as np
import torch

from super_gradients.training.datasets import DetectionDataset
from super_gradients.training.utils.detection_utils import DetectionTargetsFormat


class DummyDetectionDataset(DetectionDataset):
    def __init__(self, dataset_size, input_dim, *args, **kwargs):
        """Dummy Dataset testing subsampling."""

        self.dataset_size = dataset_size
        self.image_size = input_dim
        kwargs['all_classes_list'] = ["class_0", "class_1", "class_2"]
        kwargs['original_target_format'] = DetectionTargetsFormat.XYXY_LABEL
        super().__init__(data_dir='', input_dim=input_dim, *args, **kwargs)

    def _setup_data_source(self):
        return self.dataset_size

    def _load_annotation(self, sample_id: int) -> dict:
        """Load dummy annotation"""
        return {"img_path": "", "target": torch.zeros(10, 6)}

    # DetectionDatasetV2 will call _load_image but since we don't have any image we patch this method with
    # tensor of image shape
    def _load_image(self, index: int) -> np.ndarray:
        return np.random.random(self.image_size)


class TestDetectionDatasetSubsampling(unittest.TestCase):

    def test_subsampling(self):
        """Check that subsampling works"""
        for max_num_samples in [1, 1_000, 1_000_000]:
            test_dataset = DummyDetectionDataset(dataset_size=100_000, input_dim=(640, 512), max_num_samples=max_num_samples)
            self.assertEqual(len(test_dataset), min(max_num_samples, 100_000))


if __name__ == '__main__':
    unittest.main()