|
| 1 | +# Copyright (c) OpenMMLab. All rights reserved. |
| 2 | +import os.path as osp |
| 3 | +from typing import List |
| 4 | + |
| 5 | +import mmengine.fileio as fileio |
| 6 | + |
| 7 | +from mmseg.registry import DATASETS |
| 8 | +from .basesegdataset import BaseSegDataset |
| 9 | + |
| 10 | + |
| 11 | +@DATASETS.register_module() |
| 12 | +class NYUDataset(BaseSegDataset): |
| 13 | + """NYU depth estimation dataset. The file structure should be. |
| 14 | +
|
| 15 | + .. code-block:: none |
| 16 | +
|
| 17 | + ├── data |
| 18 | + │ ├── nyu |
| 19 | + │ │ ├── images |
| 20 | + │ │ │ ├── train |
| 21 | + │ │ │ │ ├── scene_xxx.jpg |
| 22 | + │ │ │ │ ├── ... |
| 23 | + │ │ │ ├── test |
| 24 | + │ │ ├── annotations |
| 25 | + │ │ │ ├── train |
| 26 | + │ │ │ │ ├── scene_xxx.png |
| 27 | + │ │ │ │ ├── ... |
| 28 | + │ │ │ ├── test |
| 29 | +
|
| 30 | + Args: |
| 31 | + ann_file (str): Annotation file path. Defaults to ''. |
| 32 | + metainfo (dict, optional): Meta information for dataset, such as |
| 33 | + specify classes to load. Defaults to None. |
| 34 | + data_root (str, optional): The root directory for ``data_prefix`` and |
| 35 | + ``ann_file``. Defaults to None. |
| 36 | + data_prefix (dict, optional): Prefix for training data. Defaults to |
| 37 | + dict(img_path='images', depth_map_path='annotations'). |
| 38 | + img_suffix (str): Suffix of images. Default: '.jpg' |
| 39 | + seg_map_suffix (str): Suffix of segmentation maps. Default: '.png' |
| 40 | + filter_cfg (dict, optional): Config for filter data. Defaults to None. |
| 41 | + indices (int or Sequence[int], optional): Support using first few |
| 42 | + data in annotation file to facilitate training/testing on a smaller |
| 43 | + dataset. Defaults to None which means using all ``data_infos``. |
| 44 | + serialize_data (bool, optional): Whether to hold memory using |
| 45 | + serialized objects, when enabled, data loader workers can use |
| 46 | + shared RAM from master process instead of making a copy. Defaults |
| 47 | + to True. |
| 48 | + pipeline (list, optional): Processing pipeline. Defaults to []. |
| 49 | + test_mode (bool, optional): ``test_mode=True`` means in test phase. |
| 50 | + Defaults to False. |
| 51 | + lazy_init (bool, optional): Whether to load annotation during |
| 52 | + instantiation. In some cases, such as visualization, only the meta |
| 53 | + information of the dataset is needed, which is not necessary to |
| 54 | + load annotation file. ``Basedataset`` can skip load annotations to |
| 55 | + save time by set ``lazy_init=True``. Defaults to False. |
| 56 | + max_refetch (int, optional): If ``Basedataset.prepare_data`` get a |
| 57 | + None img. The maximum extra number of cycles to get a valid |
| 58 | + image. Defaults to 1000. |
| 59 | + ignore_index (int): The label index to be ignored. Default: 255 |
| 60 | + reduce_zero_label (bool): Whether to mark label zero as ignored. |
| 61 | + Default to False. |
| 62 | + backend_args (dict, Optional): Arguments to instantiate a file backend. |
| 63 | + See https://mmengine.readthedocs.io/en/latest/api/fileio.htm |
| 64 | + for details. Defaults to None. |
| 65 | + Notes: mmcv>=2.0.0rc4, mmengine>=0.2.0 required. |
| 66 | + """ |
| 67 | + METAINFO = dict( |
| 68 | + classes=('printer_room', 'bathroom', 'living_room', 'study', |
| 69 | + 'conference_room', 'study_room', 'kitchen', 'home_office', |
| 70 | + 'bedroom', 'dinette', 'playroom', 'indoor_balcony', |
| 71 | + 'laundry_room', 'basement', 'excercise_room', 'foyer', |
| 72 | + 'home_storage', 'cafe', 'furniture_store', 'office_kitchen', |
| 73 | + 'student_lounge', 'dining_room', 'reception_room', |
| 74 | + 'computer_lab', 'classroom', 'office', 'bookstore')) |
| 75 | + |
| 76 | + def __init__(self, |
| 77 | + data_prefix=dict( |
| 78 | + img_path='images', depth_map_path='annotations'), |
| 79 | + img_suffix='.jpg', |
| 80 | + depth_map_suffix='.png', |
| 81 | + **kwargs) -> None: |
| 82 | + super().__init__( |
| 83 | + data_prefix=data_prefix, |
| 84 | + img_suffix=img_suffix, |
| 85 | + seg_map_suffix=depth_map_suffix, |
| 86 | + **kwargs) |
| 87 | + |
| 88 | + def _get_category_id_from_filename(self, image_fname: str) -> int: |
| 89 | + """Retrieve the category ID from the given image filename.""" |
| 90 | + image_fname = osp.basename(image_fname) |
| 91 | + position = image_fname.find(next(filter(str.isdigit, image_fname)), 0) |
| 92 | + categoty_name = image_fname[:position - 1] |
| 93 | + if categoty_name not in self._metainfo['classes']: |
| 94 | + return -1 |
| 95 | + else: |
| 96 | + return self._metainfo['classes'].index(categoty_name) |
| 97 | + |
| 98 | + def load_data_list(self) -> List[dict]: |
| 99 | + """Load annotation from directory or annotation file. |
| 100 | +
|
| 101 | + Returns: |
| 102 | + list[dict]: All data info of dataset. |
| 103 | + """ |
| 104 | + data_list = [] |
| 105 | + img_dir = self.data_prefix.get('img_path', None) |
| 106 | + ann_dir = self.data_prefix.get('depth_map_path', None) |
| 107 | + |
| 108 | + _suffix_len = len(self.img_suffix) |
| 109 | + for img in fileio.list_dir_or_file( |
| 110 | + dir_path=img_dir, |
| 111 | + list_dir=False, |
| 112 | + suffix=self.img_suffix, |
| 113 | + recursive=True, |
| 114 | + backend_args=self.backend_args): |
| 115 | + data_info = dict(img_path=osp.join(img_dir, img)) |
| 116 | + if ann_dir is not None: |
| 117 | + depth_map = img[:-_suffix_len] + self.seg_map_suffix |
| 118 | + data_info['depth_map_path'] = osp.join(ann_dir, depth_map) |
| 119 | + data_info['seg_fields'] = [] |
| 120 | + data_info['category_id'] = self._get_category_id_from_filename(img) |
| 121 | + data_list.append(data_info) |
| 122 | + data_list = sorted(data_list, key=lambda x: x['img_path']) |
| 123 | + return data_list |
0 commit comments