Spaces:
Runtime error
Runtime error
File size: 5,707 Bytes
412c852 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 |
# 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
|