Cityscapes AutoLabelling dataset (#1000)

* CityscapesConcatDataset

* documentation

* ddrnet recipe

* unit test

* docs

* add to init

src/super_gradients/common/object_names.py

@@ -403,6 +403,7 @@ class Datasets:
     PASCAL_AUG_2012_SEGMENTATION_DATASET = "PascalAUG2012SegmentationDataSet"
     PASCAL_VOC_2012_SEGMENTATION_DATASET = "PascalVOC2012SegmentationDataSet"
     CITYSCAPES_DATASET = "CityscapesDataset"
+    CITYSCAPES_CONCAT_DATASET = "CityscapesConcatDataset"
     MAPILLARY_DATASET = "MapillaryDataset"
     SUPERVISELY_PERSONS_DATASET = "SuperviselyPersonsDataset"
     PASCAL_VOC_AND_AUG_UNIFIED_DATASET = "PascalVOCAndAUGUnifiedDataset"

src/super_gradients/recipes/cityscapes_al_ddrnet.yaml

@@ -0,0 +1,89 @@
+# Instructions:
+#   0. Make sure that the data is stored in dataset_params.dataset_dir or add "dataset_params.data_dir=<PATH-TO-DATASET>" at the end of the command below (feel free to check ReadMe)
+#   1. Prepare the Cityscapes AutoLabelling dataset as described in `src/super_gradients/training/datasets/Dataset_Setup_Instructions.md`.
+#   2. Move to the project root (where you will find the ReadMe and src folder)
+#   3. Run the command:
+#      DDRNet39:        python -m super_gradients.train_from_recipe --config-name=cityscapes_ddrnet architecture=ddrnet_39
+# Note: add "checkpoint_params.checkpoint_path=<ddrnet39-backbone-pretrained-path>" to use pretrained backbone
+#
+#  Validation mIoU - Cityscapes, training time:
+#      DDRNet39:        input-size: [1024, 2048]     mIoU: 85.17    4 X RTX A5000, 38 H
+#
+#  Pretrained checkpoints:
+#      Backbones- downloaded from the author's official repo.
+#       https://deci-pretrained-models.s3.amazonaws.com/ddrnet/imagenet_pt_backbones/ddrnet39_bb_imagenet.pth
+#
+#      Network checkpoints:
+#       DDRNet39:       https://sghub.deci.ai/models/ddrnet_39_cityscapes.pth
+#
+#  Learning rate and batch size parameters, using 4 RTX A5000 with DDP:
+#      DDRNet39:        input-size: [1024, 1024]     initial_lr: 0.0075    batch-size: 6 * 4gpus = 24
+#
+#  Comments:
+#      * Pretrained backbones were used.
+
+defaults:
+  - training_hyperparams: cityscapes_default_train_params
+  - dataset_params: cityscapes_al_dataset_params
+  - checkpoint_params: default_checkpoint_params
+  - _self_
+  - variable_setup
+
+architecture: ddrnet_39
+
+dataset_params:
+  train_dataloader_params:
+    batch_size: 6
+  val_dataloader_params:
+    batch_size: 3
+  train_dataset_params:
+    transforms:
+      - SegColorJitter:
+          brightness: 0.5
+          contrast: 0.5
+          saturation: 0.5
+
+      - SegRandomFlip:
+          prob: 0.5
+
+      - SegRandomRescale:
+          scales: [ 0.5, 2. ]
+
+      - SegPadShortToCropSize:
+          crop_size: [ 1024, 1024 ]
+          fill_mask: 19
+
+      - SegCropImageAndMask:
+          crop_size: [ 1024, 1024 ]
+          mode: random
+
+training_hyperparams:
+  max_epochs: 200
+  initial_lr: 0.0075   # batch size 24
+  loss:
+    dice_ce_edge_loss:
+      num_classes: 19
+      ignore_index: 19
+      num_aux_heads: 1
+      num_detail_heads: 0
+      weights: [ 1., 0.4 ]
+      dice_ce_weights: [ 1., 1. ]
+      ce_edge_weights: [ .5, .5 ]
+      edge_kernel: 5
+  sync_bn: True
+
+arch_params:
+  num_classes: 19
+  use_aux_heads: True
+
+load_checkpoint: False
+checkpoint_params:
+  load_checkpoint: ${load_checkpoint}
+  checkpoint_path: ???
+  load_backbone: True
+  strict_load: no_key_matching
+
+experiment_name: ${architecture}_cityscapes_al
+
+multi_gpu: DDP
+num_gpus: 4

src/super_gradients/recipes/dataset_params/cityscapes_al_dataset_params.yaml

@@ -0,0 +1,42 @@
+# Cityscapes AutoLabelled dataset were introduced by NVIDIA research group.
+#   paper:
+#     Hierarchical Multi-Scale Attention for Semantic Segmentation", https://arxiv.org/abs/2005.10821
+#   Official repo:
+#     https://github.com/NVIDIA/semantic-segmentation
+#
+# AutoLabelled refer to the refinement of the Cityscapes coarse data and pseudo labels generation using their suggested
+# Hierarchical multi-scale attention model.
+#
+# For dataset preparation instruction please follow:
+#   https://github.com/Deci-AI/super-gradients/blob/master/src/super_gradients/training/datasets/Dataset_Setup_Instructions.md
+
+train_dataset_params:
+  root_dir: /data/cityscapes
+  labels_csv_path: lists/labels.csv
+  list_files:
+    - lists/train.lst
+    - lists/auto_labelling.lst
+  cache_labels: False
+  cache_images: False
+  transforms:
+
+val_dataset_params:
+  root_dir: /data/cityscapes
+  list_file: lists/val.lst
+  labels_csv_path: lists/labels.csv
+  cache_labels: False
+  cache_images: False
+  transforms:
+
+train_dataloader_params:
+  dataset: CityscapesConcatDataset
+  shuffle: True
+  batch_size: 8
+  num_workers: 8
+  drop_last: True                 # drop the last incomplete batch, if dataset size is not divisible by the batch size
+
+val_dataloader_params:
+  dataset: CityscapesDataset
+  batch_size: 8
+  num_workers: 8
+  drop_last: False

src/super_gradients/training/datasets/Dataset_Setup_Instructions.md

@@ -203,7 +203,8 @@ data_set = YoloDarknetFormatDetectionDataset(data_dir='<path-to>/data_dir', imag
       ├── test.lst
       ├── train.lst
       ├── trainval.lst
-      └── val.lst
+      ├── val.lst
+      └── auto_labelling.lst
 ```
 
 2. c. Move Metadata folder to the Cityscapes folder
@@ -224,6 +225,45 @@ from super_gradients.training.datasets import CityscapesDataset
 train_set = CityscapesDataset(root_dir='.../root_dir', list_file='lists/train.lst', labels_csv_path='lists/labels.csv', ...)
 ```
 
+4. AutoLabelling dataset [Optional]
+
+Cityscapes AutoLabelled dataset were introduced by NVIDIA research group
+in the [paper](https://arxiv.org/abs/2005.10821):
+"Hierarchical Multi-Scale Attention for Semantic Segmentation".
+
+AutoLabelled refer to the refinement of the Cityscapes coarse data and pseudo
+labels generation using their suggested Hierarchical multi-scale attention model.
+
+* To download the AutoLabelled labels please refer to the original 
+[repo](https://github.com/NVIDIA/semantic-segmentation#downloadprepare-data).
+Unzip and rename the folder to `AutoLabelling` as described bellow.
+
+* Download the coarse RGB images from cityscapes official site, 
+leftImg8bit_train_extra: https://www.cityscapes-dataset.com/file-handling/?packageID=4
+
+```
+  root_dir (in recipe default to /data/cityscapes)
+      ├─── gtFine
+      │       ├── test
+      │       │     └── ...
+      │       ├─── train
+      │       │     └── ...
+      │       └─── val
+      │             └── ...
+      ├─── leftImg8bit
+      │       ├── test
+      │       │     └── ...
+      │       ├─── train
+      │       │     └── ...
+      │       └─── val
+      │             └── ...
+      ├─── AutoLabelling
+      │       └─── train_extra
+      │             └── ...
+      └─── leftImg8bit
+              └─── train_extra
+                    └── ...
+```
 
  </details>
 

src/super_gradients/training/datasets/__init__.py

@@ -16,7 +16,7 @@
     PascalAUG2012SegmentationDataSet,
     PascalVOCAndAUGUnifiedDataset,
 )
-from super_gradients.training.datasets.segmentation_datasets.cityscape_segmentation import CityscapesDataset
+from super_gradients.training.datasets.segmentation_datasets.cityscape_segmentation import CityscapesDataset, CityscapesConcatDataset
 from super_gradients.training.datasets.segmentation_datasets.coco_segmentation import CoCoSegmentationDataSet
 from super_gradients.training.datasets.segmentation_datasets.supervisely_persons_segmentation import SuperviselyPersonsDataset
 from super_gradients.training.datasets.pose_estimation_datasets import COCOKeypointsDataset
@@ -30,6 +30,7 @@
     "DirectoryDataSet",
     "SegmentationDataSet",
     "CityscapesDataset",
+    "CityscapesConcatDataset",
     "PascalVOC2012SegmentationDataSet",
     "PascalAUG2012SegmentationDataSet",
     "PascalVOCAndAUGUnifiedDataset",

src/super_gradients/training/datasets/segmentation_datasets/__init__.py

@@ -1,4 +1,4 @@
-from super_gradients.training.datasets.segmentation_datasets.cityscape_segmentation import CityscapesDataset
+from super_gradients.training.datasets.segmentation_datasets.cityscape_segmentation import CityscapesDataset, CityscapesConcatDataset
 from super_gradients.training.datasets.segmentation_datasets.coco_segmentation import CoCoSegmentationDataSet
 from super_gradients.training.datasets.segmentation_datasets.mapillary_dataset import MapillaryDataset
 from super_gradients.training.datasets.segmentation_datasets.pascal_voc_segmentation import (
@@ -15,6 +15,7 @@
     "PascalAUG2012SegmentationDataSet",
     "PascalVOC2012SegmentationDataSet",
     "CityscapesDataset",
+    "CityscapesConcatDataset",
     "SuperviselyPersonsDataset",
     "PascalVOCAndAUGUnifiedDataset",
     "MapillaryDataset",

src/super_gradients/training/datasets/segmentation_datasets/cityscape_segmentation.py

@@ -1,14 +1,14 @@
+from typing import List
 import os
 import cv2
 import numpy as np
 from PIL import Image, ImageColor
+from torch.utils.data import ConcatDataset
 
 from super_gradients.common.object_names import Datasets
 from super_gradients.common.registry.registry import register_dataset
 from super_gradients.training.datasets.segmentation_datasets.segmentation_dataset import SegmentationDataSet
 
-# TODO - ADD COARSE DATA - right now cityscapes dataset includes fine annotations. It's optional to use extra coarse
-#  annotations.
 
 # label for background and labels to ignore during training and evaluation.
 CITYSCAPES_IGNORE_LABEL = 19
@@ -75,7 +75,7 @@ class CityscapesDataset(SegmentationDataSet):
 
     def __init__(self, root_dir: str, list_file: str, labels_csv_path: str, **kwargs):
         """
-        :param root:            Absolute path to root directory of the dataset.
+        :param root_dir:        Absolute path to root directory of the dataset.
         :param list_file:       List file that contains names of images to load, line format: <image_path> <label_path>. The path is relative to root.
         :param labels_csv_path: Path to csv file, with labels metadata and mapping. The path is relative to root.
         :param kwargs:          Any hyper params required for the dataset, i.e img_size, crop_size, cache_images
@@ -148,3 +148,39 @@ def target_transform(target):
         out = SegmentationDataSet.target_transform(target)
         out[out == 255] = CITYSCAPES_IGNORE_LABEL
         return out
+
+
+@register_dataset(Datasets.CITYSCAPES_CONCAT_DATASET)
+class CityscapesConcatDataset(ConcatDataset):
+    """
+    Support building a Cityscapes dataset which includes multiple group of samples from several list files.
+    i.e to initiate a trainval dataset:
+    >>> trainval_set = CityscapesConcatDataset(
+    >>>    root_dir='/data', list_files=['lists/train.lst', 'lists/val.lst'], labels_csv_path='lists/labels.csv', ...
+    >>> )
+
+    i.e to initiate a combination of the train-set with AutoLabelling-set:
+    >>> train_al_set = CityscapesConcatDataset(
+    >>>    root_dir='/data', list_files=['lists/train.lst', 'lists/auto_labelling.lst'], labels_csv_path='lists/labels.csv', ...
+    >>> )
+    """
+
+    def __init__(self, root_dir: str, list_files: List[str], labels_csv_path: str, **kwargs):
+        """
+        :param root_dir:        Absolute path to root directory of the dataset.
+        :param list_files:      List of list files that contains names of images to load,
+                                line format: <image_path> <label_path>. The path is relative to root.
+        :param labels_csv_path: Path to csv file, with labels metadata and mapping. The path is relative to root.
+        :param kwargs:          Any hyper params required for the dataset, i.e img_size, crop_size, cache_images
+        """
+        super().__init__(
+            datasets=[
+                CityscapesDataset(
+                    root_dir=root_dir,
+                    list_file=list_file,
+                    labels_csv_path=labels_csv_path,
+                    **kwargs,
+                )
+                for list_file in list_files
+            ]
+        )

tests/unit_tests/cityscapes_dataset_test.py

@@ -1,8 +1,9 @@
 import unittest
+from typing import Type
 
 import pkg_resources
 import yaml
-from torch.utils.data import DataLoader
+from torch.utils.data import DataLoader, Dataset
 
 from super_gradients.training.dataloaders.dataloaders import (
     cityscapes_train,
@@ -15,37 +16,63 @@
     cityscapes_regseg48_train,
     cityscapes_ddrnet_val,
     cityscapes_stdc_seg75_train,
+    get,
 )
-from super_gradients.training.datasets.segmentation_datasets.cityscape_segmentation import CityscapesDataset
+from super_gradients.training.datasets.segmentation_datasets.cityscape_segmentation import CityscapesDataset, CityscapesConcatDataset
 
 
 class CityscapesDatasetTest(unittest.TestCase):
-    def setUp(self) -> None:
+    def _cityscapes_dataset_params(self):
         default_config_path = pkg_resources.resource_filename("super_gradients.recipes", "dataset_params/cityscapes_dataset_params.yaml")
         with open(default_config_path, "r") as file:
-            self.recipe = yaml.safe_load(file)
+            dataset_params = yaml.safe_load(file)
+        return dataset_params
 
-    def dataloader_tester(self, dl: DataLoader):
+    def _cityscapes_al_dataset_params(self):
+        default_config_path = pkg_resources.resource_filename("super_gradients.recipes", "dataset_params/cityscapes_al_dataset_params.yaml")
+        with open(default_config_path, "r") as file:
+            dataset_params = yaml.safe_load(file)
+        return dataset_params
+
+    def dataloader_tester(self, dl: DataLoader, dataset_cls: Type[Dataset] = CityscapesDataset):
         self.assertTrue(isinstance(dl, DataLoader))
-        self.assertTrue(isinstance(dl.dataset, CityscapesDataset))
+        self.assertTrue(isinstance(dl.dataset, dataset_cls))
         it = iter(dl)
         for _ in range(10):
             next(it)
 
     def test_train_dataset_creation(self):
-        train_dataset = CityscapesDataset(**self.recipe["train_dataset_params"])
+        dataset_params = self._cityscapes_dataset_params()
+        train_dataset = CityscapesDataset(**dataset_params["train_dataset_params"])
+        for i in range(10):
+            image, mask = train_dataset[i]
+
+    def test_al_train_dataset_creation(self):
+        dataset_params = self._cityscapes_al_dataset_params()
+        train_dataset = CityscapesConcatDataset(**dataset_params["train_dataset_params"])
         for i in range(10):
             image, mask = train_dataset[i]
 
     def test_val_dataset_creation(self):
-        val_dataset = CityscapesDataset(**self.recipe["val_dataset_params"])
+        dataset_params = self._cityscapes_dataset_params()
+        val_dataset = CityscapesDataset(**dataset_params["val_dataset_params"])
         for i in range(10):
             image, mask = val_dataset[i]
 
     def test_cityscapes_train_dataloader(self):
         dl_train = cityscapes_train()
         self.dataloader_tester(dl_train)
 
+    def test_cityscapes_al_train_dataloader(self):
+        dataset_params = self._cityscapes_al_dataset_params()
+        # Same dataloader creation as in `train_from_recipe`
+        dl_train = get(
+            name=None,
+            dataset_params=dataset_params["train_dataset_params"],
+            dataloader_params=dataset_params["train_dataloader_params"],
+        )
+        self.dataloader_tester(dl_train, dataset_cls=CityscapesConcatDataset)
+
     def test_cityscapes_val_dataloader(self):
         dl_val = cityscapes_val()
         self.dataloader_tester(dl_val)