@@ -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"

# 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

# 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

@@ -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>
 

@@ -16,7 +16,7 @@ from super_gradients.training.datasets.segmentation_datasets.pascal_voc_segmenta
     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 @@ __all__ = [
     "DirectoryDataSet",
     "SegmentationDataSet",
     "CityscapesDataset",
+    "CityscapesConcatDataset",
     "PascalVOC2012SegmentationDataSet",
     "PascalAUG2012SegmentationDataSet",
     "PascalVOCAndAUGUnifiedDataset",

@@ -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 @@ __all__ = [
     "PascalAUG2012SegmentationDataSet",
     "PascalVOC2012SegmentationDataSet",
     "CityscapesDataset",
+    "CityscapesConcatDataset",
     "SuperviselyPersonsDataset",
     "PascalVOCAndAUGUnifiedDataset",
     "MapillaryDataset",

@@ -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 @@ class CityscapesDataset(SegmentationDataSet):
         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
+            ]
+        )

@@ -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,30 +16,46 @@ from super_gradients.training.dataloaders.dataloaders import (
     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]
 
@@ -46,6 +63,16 @@ class CityscapesDatasetTest(unittest.TestCase):
         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)