Fix Data Adapter for Segmentation (#1654)

requirements.txt

@@ -36,4 +36,4 @@ numpy<=1.23
 rapidfuzz
 json-tricks==3.16.1
 onnx-simplifier>=0.4.3,<1.0
-data-gradients>=0.3.0
+data-gradients~=0.3.1

src/super_gradients/training/utils/collate_fn/adapters/segmentation_adapter_collate_fn.py

@@ -48,5 +48,4 @@ def __call__(self, samples: Iterable[SupportedDataType]) -> Tuple[torch.Tensor,
         images, targets = super().__call__(samples=samples)  # This already returns a batch of (images, targets)
         transform = SegmentationDataSet.get_normalize_transform()
         images = transform(images / 255)  # images are [0-255] after the data adapter
-        targets = targets.argmax(1)
         return images, targets

tests/deci_core_unit_test_suite_runner.py

@@ -27,6 +27,9 @@
     DynamicModelTests,
     TestExportRecipe,
     TestMixedPrecisionDisabled,
+    TestClassificationAdapter,
+    TestDetectionAdapter,
+    TestSegmentationAdapter,
 )
 from tests.end_to_end_tests import TestTrainer
 from tests.unit_tests.test_convert_recipe_to_code import TestConvertRecipeToCode
@@ -172,6 +175,9 @@ def _add_modules_to_unit_tests_suite(self):
         self.unit_tests_suite.addTest(self.test_loader.loadTestsFromModule(TestConvertRecipeToCode))
         self.unit_tests_suite.addTest(self.test_loader.loadTestsFromModule(TestVersionCheck))
         self.unit_tests_suite.addTest(self.test_loader.loadTestsFromModule(TestModelWeightAveraging))
+        self.unit_tests_suite.addTest(self.test_loader.loadTestsFromModule(TestClassificationAdapter))
+        self.unit_tests_suite.addTest(self.test_loader.loadTestsFromModule(TestDetectionAdapter))
+        self.unit_tests_suite.addTest(self.test_loader.loadTestsFromModule(TestSegmentationAdapter))
 
     def _add_modules_to_end_to_end_tests_suite(self):
         """

tests/integration_tests/data_adapter/test_dataloader_adapter.py

@@ -1,3 +1,4 @@
+import os.path
 import unittest
 import tempfile
 import shutil
@@ -10,6 +11,7 @@
 from data_gradients.managers.segmentation_manager import SegmentationAnalysisManager
 from data_gradients.managers.classification_manager import ClassificationAnalysisManager
 from data_gradients.dataset_adapters.config.data_config import SegmentationDataConfig
+from data_gradients.utils.data_classes.image_channels import ImageChannels
 
 from super_gradients.training.dataloaders.adapters import (
     DetectionDataloaderAdapterFactory,
@@ -20,7 +22,10 @@
 
 class DataloaderAdapterTest(unittest.TestCase):
     def setUp(self) -> None:
-        self.tmp_dir = tempfile.mkdtemp()
+        if os.getenv("DEBUG_DIR"):  # This is useful when debugging locally, to avoid downloading the dataset everytime
+            self.tmp_dir = os.path.join(os.getenv("DEBUG_DIR"), "DataloaderAdapterNonRegressionTest")
+        else:
+            self.tmp_dir = tempfile.mkdtemp()
 
     def tearDown(self):
         shutil.rmtree(self.tmp_dir)
@@ -38,8 +43,10 @@ def __call__(self, pic):
             log_dir=self.tmp_dir,
             report_title="Caltech101",
             class_names=train_set.categories,
+            image_channels=ImageChannels.from_str("RGB"),
+            is_batch=False,
+            labels_extractor="[1]",  # dataset returns (image, label)
             batches_early_stop=4,
-            n_image_channels=3,
             use_cache=True,
         )
         analyzer.run()
@@ -121,7 +128,8 @@ def voc_format_to_bbox(sample: tuple) -> np.ndarray:
             train_data=train_set,
             val_data=val_set,
             labels_extractor=voc_format_to_bbox,
-            class_names=PASCAL_VOC_CLASS_NAMES,
+            class_names=list(PASCAL_VOC_CLASS_NAMES),
+            image_channels=ImageChannels.from_str("RGB"),
             # class_names=train_set,
             batches_early_stop=20,
             use_cache=True,  # With this we will be asked about the dataset information only once
@@ -181,7 +189,8 @@ def test_torchvision_segmentation(self):
             log_dir=self.tmp_dir,
             train_data=train_set,
             val_data=val_set,
-            class_names=list(range(256)),
+            class_names=[f"class_{i}" for i in range(256)],
+            image_channels=ImageChannels.from_str("RGB"),
             # class_names=train_set,
             batches_early_stop=20,
             use_cache=True,  # With this we will be asked about the dataset information only once

tests/integration_tests/data_adapter/test_dataloader_adapter_non_regression.py

@@ -7,6 +7,7 @@
 from data_gradients.managers.detection_manager import DetectionAnalysisManager
 from data_gradients.managers.segmentation_manager import SegmentationAnalysisManager
 from data_gradients.managers.classification_manager import ClassificationAnalysisManager
+from data_gradients.utils.data_classes.image_channels import ImageChannels
 
 from super_gradients.training.dataloaders.dataloaders import coco2017_val, cityscapes_stdc_seg50_val, cifar10_val
 from super_gradients.training.dataloaders.adapters import (
@@ -36,6 +37,7 @@ def test_adapter_on_coco2017_val(self):
             train_data=loader,
             val_data=loader,
             class_names=loader.dataset.classes,
+            image_channels=ImageChannels.from_str("RGB"),
             batches_early_stop=20,
             use_cache=True,  # With this we will be asked about the data information only once
             bbox_format="cxcywh",
@@ -60,6 +62,7 @@ def test_adapter_on_cityscapes_stdc_seg50_val(self):
             train_data=loader,
             val_data=loader,
             class_names=loader.dataset.classes + ["<unknown>"],
+            image_channels=ImageChannels.from_str("RGB"),
             batches_early_stop=1,
             use_cache=True,  # With this we will be asked about the data information only once
         )
@@ -69,6 +72,7 @@ def test_adapter_on_cityscapes_stdc_seg50_val(self):
 
         for (adapted_images, adapted_targets), (images, targets) in zip(adapted_loader, loader):
             assert np.isclose(adapted_targets, targets).all()
+            assert np.isclose(adapted_images, images).all()
         os.remove(analyzer.data_config.cache_path)
 
     def test_adapter_on_cifar10_val(self):
@@ -81,6 +85,7 @@ def test_adapter_on_cifar10_val(self):
             train_data=loader,
             val_data=loader,
             class_names=list(range(10)),
+            image_channels=ImageChannels.from_str("RGB"),
             batches_early_stop=20,
             use_cache=True,  # With this we will be asked about the data information only once
         )
@@ -93,88 +98,31 @@ def test_adapter_on_cifar10_val(self):
             assert np.isclose(adapted_images, images).all()
         os.remove(analyzer.data_config.cache_path)
 
-    def test_adpter_from_dict(self):
+    def test_ddp_python_based_adapter(self):
+        # setup_device(num_gpus=3)
+
         # We use Validation set because it does not include augmentation (which is random and makes it impossible to compare results)
-        loader = coco2017_val(
-            dataset_params={"max_num_samples": 500, "with_crowd": False},
-            dataloader_params={"collate_fn": "DetectionCollateFN"},
-        )  # `max_num_samples` To make it faster
+        loader = cifar10_val(dataset_params={"transforms": ["ToTensor"]})
 
-        analyzer = DetectionAnalysisManager(
-            report_title="coco2017_val_dict",
+        analyzer = ClassificationAnalysisManager(
+            report_title="test_python_classification",
             log_dir=self.tmp_dir,
             train_data=loader,
             val_data=loader,
-            class_names=loader.dataset.classes,
+            class_names=list(range(10)),
+            image_channels=ImageChannels.from_str("RGB"),
             batches_early_stop=20,
             use_cache=True,  # With this we will be asked about the data information only once
-            bbox_format="cxcywh",
-            is_label_first=True,
         )
         analyzer.run()
 
-        # Here we mimic how it works when loading from a recipe
-        adapted_loader = coco2017_val(
-            dataset_params={"max_num_samples": 500, "with_crowd": False},
-            dataloader_params={
-                "collate_fn": {
-                    "DetectionDatasetAdapterCollateFN": {
-                        "base_collate_fn": "DetectionCollateFN",
-                        "config_path": analyzer.data_config.cache_path,
-                    }
-                }
-            },
-        )
+        adapted_loader = ClassificationDataloaderAdapterFactory.from_dataloader(dataloader=loader, config_path=analyzer.data_config.cache_path)
 
         for (adapted_images, adapted_targets), (images, targets) in zip(adapted_loader, loader):
             assert np.isclose(adapted_targets, targets).all()
             assert np.isclose(adapted_images, images).all()
         os.remove(analyzer.data_config.cache_path)
 
-    def test_ddp_from_dict_based_adapter(self):
-        # setup_device(num_gpus=3)
-
-        # We use Validation set because it does not include augmentation (which is random and makes it impossible to compare results)
-        loader = coco2017_val(
-            dataset_params={"max_num_samples": 500, "with_crowd": False},
-            dataloader_params={"num_workers": 4, "collate_fn": "DetectionCollateFN"},
-        )
-
-        # We use Validation set because it does not include augmentation (which is random and makes it impossible to compare results)
-        adapted_loader = coco2017_val(
-            dataset_params={"max_num_samples": 500, "with_crowd": False},  # `max_num_samples` To make it faster
-            dataloader_params={
-                "num_workers": 4,
-                "collate_fn": {
-                    "DetectionDatasetAdapterCollateFN": {
-                        "base_collate_fn": "DetectionCollateFN",
-                        "config_path": os.path.join(self.tmp_dir, "test_ddp_from_dict_based_adapter.json"),
-                    }
-                },
-            },
-        )
-
-        for (adapted_images, adapted_targets), (images, targets) in zip(adapted_loader, loader):
-            assert np.isclose(adapted_targets, targets).all()
-            assert np.isclose(adapted_images, images).all()
-
-    def test_ddp_python_based_adapter(self):
-        # setup_device(num_gpus=3)
-
-        # We use Validation set because it does not include augmentation (which is random and makes it impossible to compare results)
-        loader = coco2017_val(
-            dataset_params={"max_num_samples": 500, "with_crowd": False},  # `max_num_samples` To make it faster
-            dataloader_params={"num_workers": 4, "collate_fn": "DetectionCollateFN"},
-        )
-        adapted_loader = DetectionDataloaderAdapterFactory.from_dataloader(
-            dataloader=loader,
-            config_path=os.path.join(self.tmp_dir, "test_ddp_python_based_adapter.json"),
-        )
-
-        for (adapted_images, adapted_targets), (images, targets) in zip(adapted_loader, loader):
-            assert np.isclose(adapted_targets, targets).all()
-            assert np.isclose(adapted_images, images).all()
-
 
 if __name__ == "__main__":
     DataloaderAdapterNonRegressionTest()

tests/unit_tests/__init__.py

@@ -29,6 +29,7 @@
 from tests.unit_tests.test_deprecate import TestDeprecationDecorator
 from tests.unit_tests.test_models_factory import DynamicModelTests
 from tests.unit_tests.test_mixed_precision_cpu import TestMixedPrecisionDisabled
+from tests.unit_tests.test_data_adapters import TestClassificationAdapter, TestDetectionAdapter, TestSegmentationAdapter
 
 __all__ = [
     "CrashTipTest",
@@ -61,4 +62,7 @@
     "TestMixedPrecisionDisabled",
     "DynamicModelTests",
     "TestExportRecipe",
+    "TestClassificationAdapter",
+    "TestDetectionAdapter",
+    "TestSegmentationAdapter",
 ]