[Feature] Support TSP6K dataset(CVPR'2024)

README.md

@@ -340,6 +340,7 @@ Results and models are available in the [model zoo](docs/en/model_zoo.md).
           <li><a href="https://github.com/open-mmlab/mmsegmentation/blob/main/docs/en/user_guides/2_dataset_prepare.md#bdd100K">BDD100K</a></li>
           <li><a href="https://github.com/open-mmlab/mmsegmentation/blob/main/docs/en/user_guides/2_dataset_prepare.md#nyu">NYU</a></li>
           <li><a href="https://github.com/open-mmlab/mmsegmentation/blob/main/docs/en/user_guides/2_dataset_prepare.md#hsi-drive-2.0">HSIDrive20</a></li>
+          <li><a href="https://github.com/open-mmlab/mmsegmentation/blob/main/docs/en/user_guides/2_dataset_prepare.md#TSP6K">TSP6K</a></li>
         </ul>
       </td>
       <td>

README_zh-CN.md

@@ -329,6 +329,7 @@ MMSegmentation v1.x 在 0.x 版本的基础上有了显著的提升，提供了
           <li><a href="https://github.com/open-mmlab/mmsegmentation/blob/main/docs/zh_cn/user_guides/2_dataset_prepare.md#bdd100K">BDD100K</a></li>
           <li><a href="https://github.com/open-mmlab/mmsegmentation/blob/main/docs/en/user_guides/2_dataset_prepare.md#nyu">NYU</a></li>
           <li><a href="https://github.com/open-mmlab/mmsegmentation/blob/main/docs/en/user_guides/2_dataset_prepare.md#hsi-drive-2.0">HSIDrive20</a></li>
+          <li><a href="https://github.com/open-mmlab/mmsegmentation/blob/main/docs/en/user_guides/2_dataset_prepare.md#TSP6K">TSP6K</a></li>
         </ul>
       </td>
       <td>

configs/_base_/datasets/tsp6k.py

@@ -0,0 +1,65 @@
+# dataset settings
+dataset_type = 'TSP6KDataset'
+data_root = 'data/TSP6K/'
+crop_size = (512, 1024)
+train_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(type='LoadAnnotations'),
+    dict(
+        type='RandomResize',
+        scale=(2048, 1024),
+        ratio_range=(0.5, 2.0),
+        keep_ratio=True),
+    dict(type='RandomCrop', crop_size=crop_size, cat_max_ratio=0.75),
+    dict(type='RandomFlip', prob=0.5),
+    dict(type='PhotoMetricDistortion'),
+    dict(type='PackSegInputs')
+]
+test_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(type='Resize', scale=(2048, 1024), keep_ratio=True),
+    # add loading annotation after ``Resize`` because ground truth
+    # does not need to do resize data transform
+    dict(type='LoadAnnotations'),
+    dict(type='PackSegInputs')
+]
+img_ratios = [0.5, 0.75, 1.0, 1.25, 1.5, 1.75]
+tta_pipeline = [
+    dict(type='LoadImageFromFile', backend_args=None),
+    dict(
+        type='TestTimeAug',
+        transforms=[
+            [
+                dict(type='Resize', scale_factor=r, keep_ratio=True)
+                for r in img_ratios
+            ],
+            [
+                dict(type='RandomFlip', prob=0., direction='horizontal'),
+                dict(type='RandomFlip', prob=1., direction='horizontal')
+            ], [dict(type='LoadAnnotations')], [dict(type='PackSegInputs')]
+        ])
+]
+train_dataloader = dict(
+    batch_size=2,
+    num_workers=2,
+    persistent_workers=True,
+    sampler=dict(type='InfiniteSampler', shuffle=True),
+    dataset=dict(
+        type=dataset_type,
+        data_root=data_root,
+        data_prefix=dict(img_path='image/train', seg_map_path='label/train'),
+        pipeline=train_pipeline))
+val_dataloader = dict(
+    batch_size=1,
+    num_workers=4,
+    persistent_workers=True,
+    sampler=dict(type='DefaultSampler', shuffle=False),
+    dataset=dict(
+        type=dataset_type,
+        data_root=data_root,
+        data_prefix=dict(img_path='image/val', seg_map_path='label/val'),
+        pipeline=test_pipeline))
+test_dataloader = val_dataloader
+
+val_evaluator = dict(type='IoUMetric', iou_metrics=['mIoU'])
+test_evaluator = val_evaluator

configs/_base_/datasets/tsp6k_1024x1024.py

@@ -0,0 +1,29 @@
+_base_ = './tsp6k.py'
+crop_size = (1024, 1024)
+train_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(type='LoadAnnotations'),
+    dict(
+        type='RandomResize',
+        scale=(2048, 1024),
+        ratio_range=(0.5, 2.0),
+        keep_ratio=True),
+    dict(type='RandomCrop', crop_size=crop_size, cat_max_ratio=0.75),
+    dict(type='RandomFlip', prob=0.5),
+    dict(type='PhotoMetricDistortion'),
+    dict(type='PackSegInputs')
+]
+test_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(type='Resize', scale=(2048, 1024), keep_ratio=True),
+    # add loading annotation after ``Resize`` because ground truth
+    # does not need to do resize data transform
+    dict(type='LoadAnnotations'),
+    dict(type='PackSegInputs')
+]
+train_dataloader = dict(dataset=dict(pipeline=train_pipeline))
+val_dataloader = dict(dataset=dict(pipeline=test_pipeline))
+test_dataloader = val_dataloader
+
+val_evaluator = dict(type='IoUMetric', iou_metrics=['mIoU'])
+test_evaluator = val_evaluator

configs/_base_/datasets/tsp6k_769x769.py

@@ -0,0 +1,29 @@
+_base_ = './tsp6k.py'
+crop_size = (769, 769)
+train_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(type='LoadAnnotations'),
+    dict(
+        type='RandomResize',
+        scale=(2049, 1025),
+        ratio_range=(0.5, 2.0),
+        keep_ratio=True),
+    dict(type='RandomCrop', crop_size=crop_size, cat_max_ratio=0.75),
+    dict(type='RandomFlip', prob=0.5),
+    dict(type='PhotoMetricDistortion'),
+    dict(type='PackSegInputs')
+]
+test_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(type='Resize', scale=(2049, 1025), keep_ratio=True),
+    # add loading annotation after ``Resize`` because ground truth
+    # does not need to do resize data transform
+    dict(type='LoadAnnotations'),
+    dict(type='PackSegInputs')
+]
+train_dataloader = dict(dataset=dict(pipeline=train_pipeline))
+val_dataloader = dict(dataset=dict(pipeline=test_pipeline))
+test_dataloader = val_dataloader
+
+val_evaluator = dict(type='IoUMetric', iou_metrics=['mIoU'])
+test_evaluator = val_evaluator

configs/deeplabv3plus/deeplabv3plus_r50b-d8_4xb2_80k_tsp6k_769x769.py

@@ -0,0 +1,12 @@
+_base_ = [
+    '../_base_/models/deeplabv3_r50-d8.py',
+    '../_base_/datasets/tsp6k_769x769.py', '../_base_/default_runtime.py',
+    '../_base_/schedules/schedule_80k.py'
+]
+crop_size = (769, 769)
+data_preprocessor = dict(size=crop_size)
+model = dict(
+    data_preprocessor=data_preprocessor,
+    decode_head=dict(align_corners=True, num_classes=21),
+    auxiliary_head=dict(align_corners=True, num_classes=21),
+    test_cfg=dict(mode='slide', crop_size=(769, 769), stride=(513, 513)))

docs/en/user_guides/2_dataset_prepare.md

@@ -214,6 +214,13 @@ mmsegmentation
 │   │   │   ├── train
 │   │   │   ├── validation
 │   │   │   ├── test
+│   ├── TSP6K
+│   │   ├── image
+│   │   |   ├── train
+│   │   │   ├── val
+│   │   ├── label
+│   │   |   ├── train
+│   │   │   ├── val
 ```
 
 ## Download dataset via MIM
@@ -804,3 +811,23 @@ mmsegmentation
 │   ├── HSI_Drive_v2_0_release_notes_Python_version.md
 │   ├── image_numbering.pdf
 ```
+
+## TSP6K
+
+- The TSP6K dataset can be downloaded from [here](https://drive.google.com/file/d/1yDCL8vLOYcRuInCPUij3_pmq9yPk7KQj/view?usp=sharing).
+
+- Then, unzip `TSP6K_release.zip` and the contents of original datasets include:
+
+```none
+├── TSP6K_release
+│   ├── release
+│   │   ├── image
+│   │   ├── label_trainval
+│   │   ├── split
+```
+
+- To organize the data into the required format, run the following command in your terminal:
+
+  ```bash
+  python tools/dataset_converters/tsp6k.py /path/to/TSP6K_release/release
+  ```

docs/zh_cn/user_guides/2_dataset_prepare.md

@@ -214,6 +214,14 @@ mmsegmentation
 │   │   │   ├── train
 │   │   │   ├── validation
 │   │   │   ├── test
+│   │   │   ├── test
+│   ├── TSP6K
+│   │   ├── image
+│   │   |   ├── train
+│   │   │   ├── val
+│   │   ├── label
+│   │   |   ├── train
+│   │   │   ├── val
 ```
 
 ## 用 MIM 下载数据集
@@ -800,3 +808,23 @@ mmsegmentation
 │   ├── HSI_Drive_v2_0_release_notes_Python_version.md
 │   ├── image_numbering.pdf
 ```
+
+## TSP6K
+
+- 您可以从以下位置下载TSP6K数据集 [here](https://drive.google.com/file/d/1yDCL8vLOYcRuInCPUij3_pmq9yPk7KQj/view?usp=sharing)。
+
+- 下载后，解压文件 `TSP6K_release.zip` 然后会得到如下内容：
+
+```none
+├── TSP6K_release
+│   ├── release
+│   │   ├── image
+│   │   ├── label_trainval
+│   │   ├── split
+```
+
+- 请在命令行中运行以下指令来把数据结构组织成需要的格式：
+
+  ```bash
+  python tools/dataset_converters/tsp6k.py /path/to/TSP6K_release/release
+  ```

mmseg/datasets/__init__.py

@@ -39,6 +39,7 @@
                          RandomMosaic, RandomRotate, RandomRotFlip, Rerange,
                          ResizeShortestEdge, ResizeToMultiple, RGB2Gray,
                          SegRescale)
+from .tsp6k import TSP6KDataset
 from .voc import PascalVOCDataset
 
 # yapf: enable
@@ -61,5 +62,5 @@
     'MapillaryDataset_v2', 'Albu', 'LEVIRCDDataset',
     'LoadMultipleRSImageFromFile', 'LoadSingleRSImageFromFile',
     'ConcatCDInput', 'BaseCDDataset', 'DSDLSegDataset', 'BDD100KDataset',
-    'NYUDataset', 'HSIDrive20Dataset'
+    'NYUDataset', 'HSIDrive20Dataset', 'TSP6KDataset'
 ]

mmseg/datasets/tsp6k.py

@@ -0,0 +1,27 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+from mmseg.registry import DATASETS
+from .basesegdataset import BaseSegDataset
+
+
+@DATASETS.register_module()
+class TSP6KDataset(BaseSegDataset):
+    """TSP6K dataset."""
+    METAINFO = dict(
+        classes=('road', 'sidewalk', 'building', 'wall', 'railing',
+                 'vegetation', 'terrain', 'sky', 'person', 'rider', 'car',
+                 'truck', 'bus', 'motorcycle', 'bicycle', 'indication line',
+                 'lane line', 'crosswalk', 'pole', 'traffic light',
+                 'traffic sign'),
+        palette=[[[128, 64, 128], [244, 35, 232], [70, 70, 70], [80, 90, 40],
+                  [180, 165, 180], [107, 142, 35], [152, 251, 152],
+                  [70, 130, 180], [255, 0, 0], [255, 100, 0], [0, 0, 142],
+                  [0, 0, 70], [0, 60, 100], [0, 0, 230], [119, 11, 32],
+                  [250, 170, 160], [250, 200, 160], [250, 240, 180],
+                  [153, 153, 153], [250, 170, 30], [220, 220, 0]]])
+
+    def __init__(self,
+                 img_suffix='.jpg',
+                 seg_map_suffix='_sem.png',
+                 **kwargs) -> None:
+        super().__init__(
+            img_suffix=img_suffix, seg_map_suffix=seg_map_suffix, **kwargs)

tools/dataset_converters/tsp6k.py

@@ -0,0 +1,48 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import argparse
+import os.path as osp
+import shutil
+
+from mmengine.utils import mkdir_or_exist
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description='Convert TSP6K dataset to mmsegmentation format')
+    parser.add_argument('raw_data', help='the path of raw data')
+    parser.add_argument(
+        '-o', '--out_dir', help='output path', default='./data/TSP6K')
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = parse_args()
+
+    mkdir_or_exist(args.out_dir)
+    for subdir in ['label/val', 'image/val', 'label/train', 'image/train']:
+        mkdir_or_exist(osp.join(args.out_dir, subdir))
+
+    splits = ['train', 'val']
+
+    for split in splits:
+        split_path = osp.join(args.raw_data, f'split/{split}.txt')
+        with open(split_path) as split_read:
+            image_names = [line[:-1] for line in split_read.readlines()]
+        for image_name in image_names:
+            ori_label_path = osp.join(args.raw_data,
+                                      f'label_trainval/{image_name}_sem.png')
+            tar_label_path = osp.join(args.out_dir,
+                                      f'label/{split}/{image_name}_sem.png')
+            shutil.copy(ori_label_path, tar_label_path)
+
+            ori_img_path = osp.join(args.raw_data, f'image/{image_name}.jpg')
+            tar_img_path = osp.join(args.out_dir,
+                                    f'image/{split}/{image_name}.jpg')
+            shutil.copy(ori_img_path, tar_img_path)
+
+    print('Done!')
+
+
+if __name__ == '__main__':
+    main()