@@ -1,7 +1,7 @@
 import unittest
 
 import super_gradients
-from super_gradients.training.datasets import PascalVOCDetectionDataset
+from super_gradients.training.datasets import PascalVOCDetectionDataset, COCODetectionDataset
 from super_gradients.training.transforms import DetectionMosaic, DetectionPaddedRescale, DetectionTargetsFormatTransform
 from super_gradients.training.utils.detection_utils import DetectionTargetsFormat
 from super_gradients.training.exceptions.dataset_exceptions import EmptyDatasetException
@@ -11,29 +11,62 @@ class DatasetIntegrationTest(unittest.TestCase):
     def setUp(self) -> None:
         super_gradients.init_trainer()
         self.batch_size = 64
+        self.max_samples_per_plot = 16
+        self.n_plot = 1
+        transforms = [DetectionMosaic(input_dim=(640, 640), prob=0.8),
+                      DetectionPaddedRescale(input_dim=(640, 640), max_targets=120),
+                      DetectionTargetsFormatTransform(output_format=DetectionTargetsFormat.XYXY_LABEL)]
+
         self.pascal_class_inclusion_lists = [['aeroplane', 'bicycle'],
                                              ['bird', 'boat', 'bottle', 'bus'],
                                              ['pottedplant'],
                                              ['person']]
-        transforms = [DetectionMosaic(input_dim=(640, 640), prob=0.8),
-                      DetectionPaddedRescale(input_dim=(640, 640), max_targets=120),
-                      DetectionTargetsFormatTransform(output_format=DetectionTargetsFormat.XYXY_LABEL)]
         self.pascal_base_config = dict(data_dir='/home/louis.dupont/data/pascal_unified_coco_format/',
                                        images_sub_directory='images/train2012/',
                                        input_dim=(640, 640),
                                        transforms=transforms)
 
+        self.coco_class_inclusion_lists = [['airplane', 'bicycle'],
+                                           ['bird', 'boat', 'bottle', 'bus'],
+                                           ['potted plant'],
+                                           ['person']]
+        self.dataset_coco_base_config = dict(data_dir="/data/coco",
+                                             subdir="images/val2017",
+                                             json_file="instances_val2017.json",
+                                             input_dim=(640, 640),
+                                             transforms=transforms,)
+
     def test_multiple_pascal_dataset_subclass_before_transforms(self):
         """Run test_pascal_dataset_subclass on multiple inclusion lists"""
         for class_inclusion_list in self.pascal_class_inclusion_lists:
-            dataset = PascalVOCDetectionDataset(class_inclusion_list=class_inclusion_list, **self.pascal_base_config)
-            dataset.plot(max_samples_per_plot=16, n_plots=1, plot_transformed_data=False)
+            dataset = PascalVOCDetectionDataset(class_inclusion_list=class_inclusion_list,
+                                                max_num_samples=self.max_samples_per_plot * self.n_plot,
+                                                **self.pascal_base_config)
+            dataset.plot(max_samples_per_plot=self.max_samples_per_plot, n_plots=self.n_plot, plot_transformed_data=False)
 
     def test_multiple_pascal_dataset_subclass_after_transforms(self):
         """Run test_pascal_dataset_subclass on multiple inclusion lists"""
         for class_inclusion_list in self.pascal_class_inclusion_lists:
-            dataset = PascalVOCDetectionDataset(class_inclusion_list=class_inclusion_list, **self.pascal_base_config)
-            dataset.plot(max_samples_per_plot=16, n_plots=1, plot_transformed_data=True)
+            dataset = PascalVOCDetectionDataset(class_inclusion_list=class_inclusion_list,
+                                                max_num_samples=self.max_samples_per_plot * self.n_plot,
+                                                **self.pascal_base_config)
+            dataset.plot(max_samples_per_plot=self.max_samples_per_plot, n_plots=self.n_plot, plot_transformed_data=True)
+
+    def test_multiple_coco_dataset_subclass_before_transforms(self):
+        """Check subclass on multiple inclusions before transform"""
+        for class_inclusion_list in self.coco_class_inclusion_lists:
+            dataset = COCODetectionDataset(class_inclusion_list=class_inclusion_list,
+                                           max_num_samples=self.max_samples_per_plot * self.n_plot,
+                                           **self.dataset_coco_base_config)
+            dataset.plot(max_samples_per_plot=self.max_samples_per_plot, n_plots=self.n_plot, plot_transformed_data=False)
+
+    def test_multiple_coco_dataset_subclass_after_transforms(self):
+        """Check subclass on multiple inclusions after transform"""
+        for class_inclusion_list in self.coco_class_inclusion_lists:
+            dataset = COCODetectionDataset(class_inclusion_list=class_inclusion_list,
+                                           max_num_samples=self.max_samples_per_plot * self.n_plot,
+                                           **self.dataset_coco_base_config)
+            dataset.plot(max_samples_per_plot=self.max_samples_per_plot, n_plots=self.n_plot, plot_transformed_data=True)
 
     def test_subclass_non_existing_class(self):
         """Check that EmptyDatasetException is raised when unknown label."""

@@ -1,6 +1,7 @@
 import unittest
 
-from super_gradients.training.datasets.dataset_interfaces.dataset_interface import PascalVOCUnifiedDetectionDatasetInterface
+from super_gradients.training.datasets.dataset_interfaces.dataset_interface import PascalVOCUnifiedDetectionDatasetInterface,\
+    CoCoDetectionDatasetInterface
 from super_gradients.training.transforms.transforms import DetectionPaddedRescale, DetectionTargetsFormatTransform, DetectionMosaic, DetectionRandomAffine,\
     DetectionHSV
 from super_gradients.training.utils.detection_utils import DetectionTargetsFormat
@@ -11,7 +12,7 @@ from super_gradients.training import utils as core_utils
 
 class TestDatasetInterface(unittest.TestCase):
     def setUp(self) -> None:
-        self.root_dir = "/home/data/"
+        self.root_dir = "/home/louis.dupont/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)
@@ -20,16 +21,18 @@ class TestDatasetInterface(unittest.TestCase):
         self.val_max_num_samples = 90
 
     def setup_pascal_voc_interface(self):
-        """setup PascalVOCUnifiedDetectionDataSetInterfaceV2 and return dataloaders"""
+        """setup PascalVOCUnifiedDetectionDatasetInterface 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),
@@ -52,6 +55,70 @@ class TestDatasetInterface(unittest.TestCase):
         train_loader, valid_loader, _test_loader, _classes = dataset_interface.get_data_loaders()
         return train_loader, valid_loader
 
+    def setup_coco_detection_interface(self):
+        """setup CoCoDetectionDatasetInterface and return dataloaders"""
+        dataset_params = {
+            "data_dir": "/data/coco",
+            "train_subdir": "images/train2017",  # sub directory path of data_dir containing the train data.
+            "val_subdir": "images/val2017",  # sub directory path of data_dir containing the validation data.
+            "train_json_file": "instances_train2017.json",  # path to coco train json file, data_dir/annotations/train_json_file.
+            "val_json_file": "instances_val2017.json",  # path to coco validation json file, data_dir/annotations/val_json_file.
+
+            "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,
+
+            "mixup_prob": 1.0,  # probability to apply per-sample mixup
+            "degrees": 10.,  # rotation degrees, randomly sampled from [-degrees, degrees]
+            "shear": 2.0,  # shear degrees, randomly sampled from [-degrees, degrees]
+            "flip_prob": 0.5,  # probability to apply horizontal flip
+            "hsv_prob": 1.0,  # probability to apply HSV transform
+            "hgain": 5,  # HSV transform hue gain (randomly sampled from [-hgain, hgain])
+            "sgain": 30,  # HSV transform saturation gain (randomly sampled from [-sgain, sgain])
+            "vgain": 30,  # HSV transform value gain (randomly sampled from [-vgain, vgain])
+            "mosaic_scale": [0.1, 2],  # random rescale range (keeps size by padding/cropping) after mosaic transform.
+            "mixup_scale": [0.5, 1.5],  # random rescale range for the additional sample in mixup
+            "mosaic_prob": 1.,  # probability to apply mosaic
+            "translate": 0.1,  # image translation fraction
+            "filter_box_candidates": False,  # whether to filter out transformed bboxes by edge size, area ratio, and aspect ratio.
+            "wh_thr": 2,  # edge size threshold when filter_box_candidates = True (pixels)
+            "ar_thr": 20,  # aspect ratio threshold when filter_box_candidates = True
+            "area_thr": 0.1,  # threshold for area ratio between original image and the transformed one, when when filter_box_candidates = True
+            "tight_box_rotation": False,
+
+            "download": False,
+            "train_collate_fn": DetectionCollateFN(),
+            "val_collate_fn": DetectionCollateFN(),
+            "cache_train_images": False,
+            "cache_val_images": False,
+            "cache_dir": "/home/data/cache",  # Depends on the user
+            "class_inclusion_list": None
+            # "with_crowd": True
+        }
+
+        dataset_interface = CoCoDetectionDatasetInterface(dataset_params=dataset_params)
+        train_loader, valid_loader, _test_loader, _classes = dataset_interface.get_data_loaders()
+        return train_loader, valid_loader
+
+    def test_coco_detection(self):
+        """Check that the dataset interface is correctly instantiated, and that the batch items are of expected size"""
+        train_loader, valid_loader = self.setup_coco_detection_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))
+
     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()

@@ -55,7 +55,7 @@ class TestDetectionDatasetCaching(unittest.TestCase):
 
         datasets = [
             DummyDetectionDataset(input_dim=(640, 512), ignore_empty_annotations=False, class_inclusion_list=class_inclusion_list,
-                                  cache=True, cache_path=self.temp_cache_dir, data_dir='/home/')
+                                  cache=True, cache_dir=self.temp_cache_dir, data_dir='/home/')
             for class_inclusion_list in [["class_0", "class_1", "class_2"], ["class_0"], ["class_1"], ["class_2"], ["class_1", "class_2"]]
         ]
 
@@ -70,7 +70,7 @@ class TestDetectionDatasetCaching(unittest.TestCase):
 
         datasets = [
             DummyDetectionDataset(input_dim=(640, 512), ignore_empty_annotations=True, class_inclusion_list=class_inclusion_list,
-                                  cache=True, cache_path=self.temp_cache_dir, data_dir='/home/')
+                                  cache=True, cache_dir=self.temp_cache_dir, data_dir='/home/')
             for class_inclusion_list in [["class_0", "class_1", "class_2"], ["class_0"], ["class_1"], ["class_2"], ["class_1", "class_2"]]
         ]
 
@@ -87,12 +87,12 @@ class TestDetectionDatasetCaching(unittest.TestCase):
         self.assertEqual(0, self._count_cached_array())
 
         _ = DummyDetectionDataset(input_dim=(640, 512), ignore_empty_annotations=True,
-                                  cache=True, cache_path=self.temp_cache_dir, data_dir='/home/')
+                                  cache=True, cache_dir=self.temp_cache_dir, data_dir='/home/')
         self.assertEqual(1, self._count_cached_array())
 
         for _ in range(5):
             _ = DummyDetectionDataset(input_dim=(640, 512), ignore_empty_annotations=True,
-                                      cache=True, cache_path=self.temp_cache_dir, data_dir='/home/')
+                                      cache=True, cache_dir=self.temp_cache_dir, data_dir='/home/')
             self.assertEqual(1, self._count_cached_array())
 
         self._empty_cache()