Cityscapes AutoLabelling dataset

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/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
@@ -148,3 +148,32 @@ 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):
+        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
+            ]
+        )