[Feature] Support NYU depth estimation dataset

docs/en/user_guides/2_dataset_prepare.md

@@ -198,6 +198,13 @@ mmsegmentation
 | │   │   │ └── rles
 | │ │ │ │ ├──sem_seg_train.json
 | │ │ │ │ └──sem_seg_val.json
+│ ├── nyu
+│ │ ├── images
+│ │ │ ├── train
+│ │ │ ├── test
+│ │ ├── annotations
+│ │ │ ├── train
+│ │ │ ├── test
 ```
 
 ## Cityscapes
@@ -716,3 +723,13 @@ mmsegmentation
 | │ │ │ │ ├──sem_seg_train.json
 | │ │ │ │ └──sem_seg_val.json
 ```
+
+## NYU
+
+- To access the NYU dataset, you can download it from [this link](https://drive.google.com/file/d/1wC-io-14RCIL4XTUrQLk6lBqU2AexLVp/view?usp=share_link)
+
+- Once the download is complete, you can utilize the [tools/dataset_converters/nyu.py](/tools/dataset_converters/nyu.py) script to extract and organize the data into the required format. Run the following command in your terminal:
+
+ ```bash
+ python tools/dataset_converters/nyu.py nyu.zip
+ ```

docs/zh_cn/user_guides/2_dataset_prepare.md

@@ -198,6 +198,13 @@ mmsegmentation
 | │   │   │ └── rles
 | │ │ │ │ ├──sem_seg_train.json
 | │ │ │ │ └──sem_seg_val.json
+│ ├── nyu
+│ │ ├── images
+│ │ │ ├── train
+│ │ │ ├── test
+│ │ ├── annotations
+│ │ │ ├── train
+│ │ │ ├── test
 ```
 
 ## Cityscapes
@@ -712,3 +719,13 @@ mmsegmentation
 | │ │ │ │ ├──sem_seg_train.json
 | │ │ │ │ └──sem_seg_val.json
 ```
+
+## NYU
+
+- 您可以从 [这个链接](https://drive.google.com/file/d/1wC-io-14RCIL4XTUrQLk6lBqU2AexLVp/view?usp=share_link) 下载 NYU 数据集
+
+- 下载完成后，您可以使用 [tools/dataset_converters/nyu.py](/tools/dataset_converters/nyu.py) 脚本来解压和组织数据到所需的格式
+
+ ```bash
+ python tools/dataset_converters/nyu.py nyu.zip
+ ```

mmseg/datasets/__init__.py

@@ -19,6 +19,7 @@
 from .loveda import LoveDADataset
 from .mapillary import MapillaryDataset_v1, MapillaryDataset_v2
 from .night_driving import NightDrivingDataset
+from .nyu import NYUDataset
 from .pascal_context import PascalContextDataset, PascalContextDataset59
 from .potsdam import PotsdamDataset
 from .refuge import REFUGEDataset
@@ -58,5 +59,6 @@
  'SynapseDataset', 'REFUGEDataset', 'MapillaryDataset_v1',
  'MapillaryDataset_v2', 'Albu', 'LEVIRCDDataset',
  'LoadMultipleRSImageFromFile', 'LoadSingleRSImageFromFile',
- 'ConcatCDInput', 'BaseCDDataset', 'DSDLSegDataset', 'BDD100KDataset'
+ 'ConcatCDInput', 'BaseCDDataset', 'DSDLSegDataset', 'BDD100KDataset',
+ 'NYUDataset'
 ]

mmseg/datasets/nyu.py

@@ -0,0 +1,123 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import os.path as osp
+from typing import List
+
+import mmengine.fileio as fileio
+
+from mmseg.registry import DATASETS
+from .basesegdataset import BaseSegDataset
+
+
+@DATASETS.register_module()
+class NYUDataset(BaseSegDataset):
+ """NYU depth estimation dataset. The file structure should be.
+
+ .. code-block:: none
+
+ ├── data
+ │ ├── nyu
+ │ │ ├── images
+ │ │ │ ├── train
+ │ │ │ │ ├── scene_xxx.jpg
+ │ │ │ │ ├── ...
+ │ │ │ ├── test
+ │ │ ├── annotations
+ │ │ │ ├── train
+ │ │ │ │ ├── scene_xxx.png
+ │ │ │ │ ├── ...
+ │ │ │ ├── test
+
+ Args:
+ ann_file (str): Annotation file path. Defaults to ''.
+ metainfo (dict, optional): Meta information for dataset, such as
+ specify classes to load. Defaults to None.
+ data_root (str, optional): The root directory for ``data_prefix`` and
+ ``ann_file``. Defaults to None.
+ data_prefix (dict, optional): Prefix for training data. Defaults to
+ dict(img_path='images', depth_map_path='annotations').
+ img_suffix (str): Suffix of images. Default: '.jpg'
+ seg_map_suffix (str): Suffix of segmentation maps. Default: '.png'
+ filter_cfg (dict, optional): Config for filter data. Defaults to None.
+ indices (int or Sequence[int], optional): Support using first few
+ data in annotation file to facilitate training/testing on a smaller
+ dataset. Defaults to None which means using all ``data_infos``.
+ serialize_data (bool, optional): Whether to hold memory using
+ serialized objects, when enabled, data loader workers can use
+ shared RAM from master process instead of making a copy. Defaults
+ to True.
+ pipeline (list, optional): Processing pipeline. Defaults to [].
+ test_mode (bool, optional): ``test_mode=True`` means in test phase.
+ Defaults to False.
+ lazy_init (bool, optional): Whether to load annotation during
+ instantiation. In some cases, such as visualization, only the meta
+ information of the dataset is needed, which is not necessary to
+ load annotation file. ``Basedataset`` can skip load annotations to
+ save time by set ``lazy_init=True``. Defaults to False.
+ max_refetch (int, optional): If ``Basedataset.prepare_data`` get a
+ None img. The maximum extra number of cycles to get a valid
+ image. Defaults to 1000.
+ ignore_index (int): The label index to be ignored. Default: 255
+ reduce_zero_label (bool): Whether to mark label zero as ignored.
+ Default to False.
+ backend_args (dict, Optional): Arguments to instantiate a file backend.
+ See https://mmengine.readthedocs.io/en/latest/api/fileio.htm
+ for details. Defaults to None.
+ Notes: mmcv>=2.0.0rc4, mmengine>=0.2.0 required.
+ """
+ METAINFO = dict(
+ classes=('printer_room', 'bathroom', 'living_room', 'study',
+ 'conference_room', 'study_room', 'kitchen', 'home_office',
+ 'bedroom', 'dinette', 'playroom', 'indoor_balcony',
+ 'laundry_room', 'basement', 'excercise_room', 'foyer',
+ 'home_storage', 'cafe', 'furniture_store', 'office_kitchen',
+ 'student_lounge', 'dining_room', 'reception_room',
+ 'computer_lab', 'classroom', 'office', 'bookstore'))
+
+ def __init__(self,
+ data_prefix=dict(
+ img_path='images', depth_map_path='annotations'),
+ img_suffix='.jpg',
+ depth_map_suffix='.png',
+ **kwargs) -> None:
+ super().__init__(
+ data_prefix=data_prefix,
+ img_suffix=img_suffix,
+ seg_map_suffix=depth_map_suffix,
+ **kwargs)
+
+ def _get_category_id_from_filename(self, image_fname: str) -> int:
+ """Retrieve the category ID from the given image filename."""
+ image_fname = osp.basename(image_fname)
+ position = image_fname.find(next(filter(str.isdigit, image_fname)), 0)
+ categoty_name = image_fname[:position - 1]
+ if categoty_name not in self._metainfo['classes']:
+ return -1
+ else:
+ return self._metainfo['classes'].index(categoty_name)
+
+ def load_data_list(self) -> List[dict]:
+ """Load annotation from directory or annotation file.
+
+ Returns:
+ list[dict]: All data info of dataset.
+ """
+ data_list = []
+ img_dir = self.data_prefix.get('img_path', None)
+ ann_dir = self.data_prefix.get('depth_map_path', None)
+
+ _suffix_len = len(self.img_suffix)
+ for img in fileio.list_dir_or_file(
+ dir_path=img_dir,
+ list_dir=False,
+ suffix=self.img_suffix,
+ recursive=True,
+ backend_args=self.backend_args):
+ data_info = dict(img_path=osp.join(img_dir, img))
+ if ann_dir is not None:
+ depth_map = img[:-_suffix_len] + self.seg_map_suffix
+ data_info['depth_map_path'] = osp.join(ann_dir, depth_map)
+ data_info['seg_fields'] = []
+ data_info['category_id'] = self._get_category_id_from_filename(img)
+ data_list.append(data_info)
+ data_list = sorted(data_list, key=lambda x: x['img_path'])
+ return data_list

mmseg/datasets/transforms/__init__.py

@@ -2,8 +2,8 @@
 from .formatting import PackSegInputs
 from .loading import (LoadAnnotations, LoadBiomedicalAnnotation,
  LoadBiomedicalData, LoadBiomedicalImageFromFile,
- LoadImageFromNDArray, LoadMultipleRSImageFromFile,
- LoadSingleRSImageFromFile)
+ LoadDepthAnnotation, LoadImageFromNDArray,
+ LoadMultipleRSImageFromFile, LoadSingleRSImageFromFile)
 # yapf: disable
 from .transforms import (CLAHE, AdjustGamma, Albu, BioMedical3DPad,
  BioMedical3DRandomCrop, BioMedical3DRandomFlip,
@@ -24,5 +24,5 @@
  'ResizeShortestEdge', 'BioMedicalGaussianNoise', 'BioMedicalGaussianBlur',
  'BioMedical3DRandomFlip', 'BioMedicalRandomGamma', 'BioMedical3DPad',
  'RandomRotFlip', 'Albu', 'LoadSingleRSImageFromFile', 'ConcatCDInput',
- 'LoadMultipleRSImageFromFile'
+ 'LoadMultipleRSImageFromFile', 'LoadDepthAnnotation'
 ]

mmseg/datasets/transforms/formatting.py

@@ -92,6 +92,11 @@ def transform(self, results: dict) -> dict:
  ...].astype(np.int64)))
  data_sample.set_data(dict(gt_edge_map=PixelData(**gt_edge_data)))
 
+ if 'gt_depth_map' in results:
+ gt_depth_data = dict(
+ data=to_tensor(results['gt_depth_map'][None, ...]))
+ data_sample.set_data(dict(gt_depth_map=PixelData(**gt_depth_data)))
+
  img_meta = {}
  for key in self.meta_keys:
  if key in results:

mmseg/datasets/transforms/loading.py

@@ -625,3 +625,77 @@ def __repr__(self):
  repr_str = (f'{self.__class__.__name__}('
  f'to_float32={self.to_float32})')
  return repr_str
+
+
+@TRANSFORMS.register_module()
+class LoadDepthAnnotation(BaseTransform):
+ """Load ``depth_map`` annotation provided by depth estimation dataset.
+
+ The annotation format is as the following:
+
+ .. code-block:: python
+
+ {
+ 'gt_depth_map': np.ndarray [Y, X]
+ }
+
+ Required Keys:
+
+ - seg_depth_path
+
+ Added Keys:
+
+ - gt_depth_map (np.ndarray): Depth map with shape (Y, X) by
+ default, and data type is float32 if set to_float32 = True.
+
+ Args:
+ decode_backend (str): The data decoding backend type. Options are
+ 'numpy', 'nifti', and 'cv2'. Defaults to 'cv2'.
+ to_float32 (bool): Whether to convert the loaded depth map to a float32
+ numpy array. If set to False, the loaded image is an uint16 array.
+ Defaults to True.
+ depth_rescale_factor (float): Factor to rescale the depth value to
+ limit the range. Defaults to 1.0.
+ backend_args (dict, Optional): Arguments to instantiate a file backend.
+ See :class:`mmengine.fileio` for details.
+ Defaults to None.
+ Notes: mmcv>=2.0.0rc4, mmengine>=0.2.0 required.
+ """
+
+ def __init__(self,
+ decode_backend: str = 'cv2',
+ to_float32: bool = True,
+ depth_rescale_factor: float = 1.0,
+ backend_args: Optional[dict] = None) -> None:
+ super().__init__()
+ self.decode_backend = decode_backend
+ self.to_float32 = to_float32
+ self.depth_rescale_factor = depth_rescale_factor
+ self.backend_args = backend_args.copy() if backend_args else None
+
+ def transform(self, results: Dict) -> Dict:
+ """Functions to load depth map.
+
+ Args:
+ results (dict): Result dict from :obj:``mmcv.BaseDataset``.
+
+ Returns:
+ dict: The dict contains loaded depth map.
+ """
+ data_bytes = fileio.get(results['depth_map_path'], self.backend_args)
+ gt_depth_map = datafrombytes(data_bytes, backend=self.decode_backend)
+
+ if self.to_float32:
+ gt_depth_map = gt_depth_map.astype(np.float32)
+
+ gt_depth_map *= self.depth_rescale_factor
+ results['gt_depth_map'] = gt_depth_map
+ results['seg_fields'].append('gt_depth_map')
+ return results
+
+ def __repr__(self):
+ repr_str = (f'{self.__class__.__name__}('
+ f"decode_backend='{self.decode_backend}', "
+ f'to_float32={self.to_float32}, '
+ f'backend_args={self.backend_args})')
+ return repr_str

mmseg/utils/io.py

@@ -3,6 +3,7 @@
 import io
 import pickle
 
+import cv2
 import numpy as np
 
 
@@ -12,7 +13,7 @@ def datafrombytes(content: bytes, backend: str = 'numpy') -> np.ndarray:
  Args:
  content (bytes): The data bytes got from files or other streams.
  backend (str): The data decoding backend type. Options are 'numpy',
- 'nifti' and 'pickle'. Defaults to 'numpy'.
+ 'nifti', 'cv2' and 'pickle'. Defaults to 'numpy'.
 
  Returns:
  numpy.ndarray: Loaded data array.
@@ -33,6 +34,9 @@ def datafrombytes(content: bytes, backend: str = 'numpy') -> np.ndarray:
  data = Nifti1Image.from_bytes(data.to_bytes()).get_fdata()
  elif backend == 'numpy':
  data = np.load(f)
+ elif backend == 'cv2':
+ data = np.frombuffer(f.read(), dtype=np.uint16)
+ data = cv2.imdecode(data, 2)
  else:
  raise ValueError
  return data

tests/test_datasets/test_dataset.py

@@ -9,7 +9,7 @@
  CityscapesDataset, COCOStuffDataset,
  DecathlonDataset, DSDLSegDataset, ISPRSDataset,
  LIPDataset, LoveDADataset, MapillaryDataset_v1,
- MapillaryDataset_v2, PascalVOCDataset,
+ MapillaryDataset_v2, NYUDataset, PascalVOCDataset,
  PotsdamDataset, REFUGEDataset, SynapseDataset,
  iSAIDDataset)
 from mmseg.registry import DATASETS
@@ -462,3 +462,14 @@ def test_dsdlseg_dataset():
  assert len(dataset.metainfo['classes']) == 21
  else:
  ImportWarning('Package `dsdl` is not installed.')
+
+
+def test_nyu_dataset():
+ dataset = NYUDataset(
+ data_root='tests/data/pseudo_nyu_dataset',
+ data_prefix=dict(img_path='images', depth_map_path='annotations'),
+ )
+ assert len(dataset) == 1
+ data = dataset[0]
+ assert data.get('depth_map_path', None) is not None
+ assert data.get('category_id', -1) == 26

tests/test_datasets/test_loading.py

@@ -7,10 +7,11 @@
 import numpy as np
 from mmcv.transforms import LoadImageFromFile
 
-from mmseg.datasets.transforms import (LoadAnnotations,
- LoadBiomedicalAnnotation,
+from mmseg.datasets.transforms import LoadAnnotations # noqa
+from mmseg.datasets.transforms import (LoadBiomedicalAnnotation,
  LoadBiomedicalData,
  LoadBiomedicalImageFromFile,
+ LoadDepthAnnotation,
  LoadImageFromNDArray)
 
 
@@ -276,3 +277,19 @@ def test_load_biomedical_data(self):
  "decode_backend='numpy', "
  'to_xyz=False, '
  'backend_args=None)')
+
+ def test_load_depth_annotation(self):
+ input_results = dict(
+ img_path='tests/data/pseudo_nyu_dataset/images/'
+ 'bookstore_0001d_00001.jpg',
+ depth_map_path='tests/data/pseudo_nyu_dataset/'
+ 'annotations/bookstore_0001d_00001.png',
+ category_id=-1,
+ seg_fields=[])
+ transform = LoadDepthAnnotation(depth_rescale_factor=0.001)
+ results = transform(input_results)
+ assert 'gt_depth_map' in results
+ assert results['gt_depth_map'].shape[:2] == mmcv.imread(
+ input_results['depth_map_path']).shape[:2]
+ assert results['gt_depth_map'].dtype == np.float32
+ assert 'gt_depth_map' in results['seg_fields']