scannet_dataset.py 11.2 KB
Newer Older
dingchang's avatar
dingchang committed
1
# Copyright (c) OpenMMLab. All rights reserved.
2
import warnings
zhangwenwei's avatar
zhangwenwei committed
3
from os import path as osp
ZCMax's avatar
ZCMax committed
4
from typing import Callable, List, Optional, Union
5

6
7
import numpy as np

8
from mmdet3d.registry import DATASETS
zhangshilong's avatar
zhangshilong committed
9
from mmdet3d.structures import DepthInstance3DBoxes
jshilong's avatar
jshilong committed
10
from .det3d_dataset import Det3DDataset
ZCMax's avatar
ZCMax committed
11
from .seg3d_dataset import Seg3DDataset
12
13
14


@DATASETS.register_module()
jshilong's avatar
jshilong committed
15
class ScanNetDataset(Det3DDataset):
16
    r"""ScanNet Dataset for Detection Task.
17

wangtai's avatar
wangtai committed
18
19
    This class serves as the API for experiments on the ScanNet Dataset.

zhangwenwei's avatar
zhangwenwei committed
20
21
    Please refer to the `github repo <https://github.com/ScanNet/ScanNet>`_
    for data downloading.
wangtai's avatar
wangtai committed
22
23
24
25

    Args:
        data_root (str): Path of dataset root.
        ann_file (str): Path of annotation file.
jshilong's avatar
jshilong committed
26
27
28
29
30
31
32
        metainfo (dict, optional): Meta information for dataset, such as class
            information. Defaults to None.
        data_prefix (dict): Prefix for data. Defaults to
            `dict(pts='points',
                pts_isntance_mask='instance_mask',
                pts_semantic_mask='semantic_mask')`.
        pipeline (list[dict]): Pipeline used for data processing.
wangtai's avatar
wangtai committed
33
            Defaults to None.
jshilong's avatar
jshilong committed
34
        modality (dict): Modality to specify the sensor data used
wangtai's avatar
wangtai committed
35
            as input. Defaults to None.
jshilong's avatar
jshilong committed
36
        box_type_3d (str): Type of 3D box of this dataset.
wangtai's avatar
wangtai committed
37
38
39
40
            Based on the `box_type_3d`, the dataset will encapsulate the box
            to its original format then converted them to `box_type_3d`.
            Defaults to 'Depth' in this dataset. Available options includes

wangtai's avatar
wangtai committed
41
42
43
            - 'LiDAR': Box in LiDAR coordinates.
            - 'Depth': Box in depth coordinates, usually for indoor dataset.
            - 'Camera': Box in camera coordinates.
jshilong's avatar
jshilong committed
44
        filter_empty_gt (bool): Whether to filter empty GT.
wangtai's avatar
wangtai committed
45
            Defaults to True.
jshilong's avatar
jshilong committed
46
        test_mode (bool): Whether the dataset is in test mode.
wangtai's avatar
wangtai committed
47
48
            Defaults to False.
    """
jshilong's avatar
jshilong committed
49
50
51
52
53
54
    METAINFO = {
        'CLASSES':
        ('cabinet', 'bed', 'chair', 'sofa', 'table', 'door', 'window',
         'bookshelf', 'picture', 'counter', 'desk', 'curtain', 'refrigerator',
         'showercurtrain', 'toilet', 'sink', 'bathtub', 'garbagebin')
    }
55
56

    def __init__(self,
jshilong's avatar
jshilong committed
57
58
59
60
61
                 data_root: str,
                 ann_file: str,
                 metainfo: dict = None,
                 data_prefix: dict = dict(
                     pts='points',
62
                     pts_instance_mask='instance_mask',
jshilong's avatar
jshilong committed
63
64
65
66
67
68
                     pts_semantic_mask='semantic_mask'),
                 pipeline: List[Union[dict, Callable]] = [],
                 modality=dict(use_camera=False, use_lidar=True),
                 box_type_3d: str = 'Depth',
                 filter_empty_gt: bool = True,
                 test_mode: bool = False,
69
                 **kwargs):
70
71
72
        super().__init__(
            data_root=data_root,
            ann_file=ann_file,
jshilong's avatar
jshilong committed
73
74
            metainfo=metainfo,
            data_prefix=data_prefix,
75
76
77
78
            pipeline=pipeline,
            modality=modality,
            box_type_3d=box_type_3d,
            filter_empty_gt=filter_empty_gt,
79
80
            test_mode=test_mode,
            **kwargs)
81
        assert 'use_camera' in self.modality and \
jshilong's avatar
jshilong committed
82
83
               'use_lidar' in self.modality
        assert self.modality['use_camera'] or self.modality['use_lidar']
84

jshilong's avatar
jshilong committed
85
86
87
    @staticmethod
    def _get_axis_align_matrix(info: dict) -> dict:
        """Get axis_align_matrix from info. If not exist, return identity mat.
88
89

        Args:
jshilong's avatar
jshilong committed
90
            info (dict): Info of a single sample data.
91
92

        Returns:
jshilong's avatar
jshilong committed
93
            np.ndarray: 4x4 transformation matrix.
94
        """
jshilong's avatar
jshilong committed
95
96
        if 'axis_align_matrix' in info:
            return np.array(info['axis_align_matrix'])
97
        else:
jshilong's avatar
jshilong committed
98
99
100
101
            warnings.warn(
                'axis_align_matrix is not found in ScanNet data info, please '
                'use new pre-process scripts to re-generate ScanNet data')
            return np.eye(4).astype(np.float32)
liyinhao's avatar
liyinhao committed
102

jshilong's avatar
jshilong committed
103
104
    def parse_data_info(self, info: dict) -> dict:
        """Process the raw data info.
105

jshilong's avatar
jshilong committed
106
107
        The only difference with it in `Det3DDataset`
        is the specific process for `axis_align_matrix'.
108
109

        Args:
jshilong's avatar
jshilong committed
110
            info (dict): Raw info dict.
111
112

        Returns:
jshilong's avatar
jshilong committed
113
            dict: Data information that will be passed to the data
zhangshilong's avatar
zhangshilong committed
114
            preprocessing transforms. It includes the following keys:
115
        """
jshilong's avatar
jshilong committed
116
117
118
119
120
121
122
123
124
125
126
127
128
        info['axis_align_matrix'] = self._get_axis_align_matrix(info)
        info['pts_instance_mask_path'] = osp.join(
            self.data_prefix.get('pts_instance_mask', ''),
            info['pts_instance_mask_path'])
        info['pts_semantic_mask_path'] = osp.join(
            self.data_prefix.get('pts_semantic_mask', ''),
            info['pts_semantic_mask_path'])

        info = super().parse_data_info(info)
        return info

    def parse_ann_info(self, info: dict) -> dict:
        """Process the `instances` in data info to `ann_info`
129
130

        Args:
jshilong's avatar
jshilong committed
131
            info (dict): Info dict.
132
133

        Returns:
jshilong's avatar
jshilong committed
134
            dict: Processed `ann_info`
135
        """
jshilong's avatar
jshilong committed
136
        ann_info = super().parse_ann_info(info)
137
138
        # empty gt
        if ann_info is None:
jshilong's avatar
jshilong committed
139
            ann_info = dict()
140
141
            ann_info['gt_bboxes_3d'] = np.zeros((0, 6), dtype=np.float32)
            ann_info['gt_labels_3d'] = np.zeros((0, ), dtype=np.int64)
jshilong's avatar
jshilong committed
142
        # to target box structure
143

jshilong's avatar
jshilong committed
144
145
146
147
148
149
150
        ann_info['gt_bboxes_3d'] = DepthInstance3DBoxes(
            ann_info['gt_bboxes_3d'],
            box_dim=ann_info['gt_bboxes_3d'].shape[-1],
            with_yaw=False,
            origin=(0.5, 0.5, 0.5)).convert_to(self.box_mode_3d)

        return ann_info
151

152
153

@DATASETS.register_module()
ZCMax's avatar
ZCMax committed
154
class ScanNetSegDataset(Seg3DDataset):
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
    r"""ScanNet Dataset for Semantic Segmentation Task.

    This class serves as the API for experiments on the ScanNet Dataset.

    Please refer to the `github repo <https://github.com/ScanNet/ScanNet>`_
    for data downloading.

    Args:
        data_root (str): Path of dataset root.
        ann_file (str): Path of annotation file.
        pipeline (list[dict], optional): Pipeline used for data processing.
            Defaults to None.
        classes (tuple[str], optional): Classes used in the dataset.
            Defaults to None.
        palette (list[list[int]], optional): The palette of segmentation map.
            Defaults to None.
        modality (dict, optional): Modality to specify the sensor data used
            as input. Defaults to None.
        test_mode (bool, optional): Whether the dataset is in test mode.
            Defaults to False.
175
        ignore_index (int, optional): The label index to be ignored, e.g.
176
177
178
179
180
181
            unannotated points. If None is given, set to len(self.CLASSES).
            Defaults to None.
        scene_idxs (np.ndarray | str, optional): Precomputed index to load
            data. For scenes with many points, we may sample it several times.
            Defaults to None.
    """
ZCMax's avatar
ZCMax committed
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
    METAINFO = {
        'CLASSES':
        ('wall', 'floor', 'cabinet', 'bed', 'chair', 'sofa', 'table', 'door',
         'window', 'bookshelf', 'picture', 'counter', 'desk', 'curtain',
         'refrigerator', 'showercurtrain', 'toilet', 'sink', 'bathtub',
         'otherfurniture'),
        'PALETTE': [
            [174, 199, 232],
            [152, 223, 138],
            [31, 119, 180],
            [255, 187, 120],
            [188, 189, 34],
            [140, 86, 75],
            [255, 152, 150],
            [214, 39, 40],
            [197, 176, 213],
            [148, 103, 189],
            [196, 156, 148],
            [23, 190, 207],
            [247, 182, 210],
            [219, 219, 141],
            [255, 127, 14],
            [158, 218, 229],
            [44, 160, 44],
            [112, 128, 144],
            [227, 119, 194],
            [82, 84, 163],
        ],
        'valid_class_ids': (1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 16, 24,
                            28, 33, 34, 36, 39),
        'all_class_ids':
        tuple(range(41)),
    }
215
216

    def __init__(self,
ZCMax's avatar
ZCMax committed
217
218
219
220
221
222
223
                 data_root: Optional[str] = None,
                 ann_file: str = '',
                 metainfo: Optional[dict] = None,
                 data_prefix: dict = dict(
                     pts='points', img='', instance_mask='', semantic_mask=''),
                 pipeline: List[Union[dict, Callable]] = [],
                 modality: dict = dict(use_lidar=True, use_camera=False),
224
                 ignore_index=None,
225
                 scene_idxs=None,
ZCMax's avatar
ZCMax committed
226
227
                 test_mode=False,
                 **kwargs) -> None:
228
229
230
        super().__init__(
            data_root=data_root,
            ann_file=ann_file,
ZCMax's avatar
ZCMax committed
231
232
            metainfo=metainfo,
            data_prefix=data_prefix,
233
234
235
            pipeline=pipeline,
            modality=modality,
            ignore_index=ignore_index,
236
            scene_idxs=scene_idxs,
ZCMax's avatar
ZCMax committed
237
            test_mode=test_mode,
238
            **kwargs)
239

240
241
    def get_scene_idxs(self, scene_idxs):
        """Compute scene_idxs for data sampling.
242

243
        We sample more times for scenes with more points.
244
245
246
247
248
249
        """
        # when testing, we load one whole scene every time
        if not self.test_mode and scene_idxs is None:
            raise NotImplementedError(
                'please provide re-sampled scene indexes for training')

250
        return super().get_scene_idxs(scene_idxs)
251

252
253

@DATASETS.register_module()
ZCMax's avatar
ZCMax committed
254
class ScanNetInstanceSegDataset(Seg3DDataset):
255

ZCMax's avatar
ZCMax committed
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
    METAINFO = {
        'CLASSES':
        ('cabinet', 'bed', 'chair', 'sofa', 'table', 'door', 'window',
         'bookshelf', 'picture', 'counter', 'desk', 'curtain', 'refrigerator',
         'showercurtrain', 'toilet', 'sink', 'bathtub', 'garbagebin'),
        'PLATTE': [
            [174, 199, 232],
            [152, 223, 138],
            [31, 119, 180],
            [255, 187, 120],
            [188, 189, 34],
            [140, 86, 75],
            [255, 152, 150],
            [214, 39, 40],
            [197, 176, 213],
            [148, 103, 189],
            [196, 156, 148],
            [23, 190, 207],
            [247, 182, 210],
            [219, 219, 141],
            [255, 127, 14],
            [158, 218, 229],
            [44, 160, 44],
            [112, 128, 144],
            [227, 119, 194],
            [82, 84, 163],
        ],
        'valid_class_ids':
        (3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 16, 24, 28, 33, 34, 36, 39),
        'all_class_ids':
        tuple(range(41))
    }
288

ZCMax's avatar
ZCMax committed
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
    def __init__(self,
                 data_root: Optional[str] = None,
                 ann_file: str = '',
                 metainfo: Optional[dict] = None,
                 data_prefix: dict = dict(
                     pts='points', img='', instance_mask='', semantic_mask=''),
                 pipeline: List[Union[dict, Callable]] = [],
                 modality: dict = dict(use_lidar=True, use_camera=False),
                 test_mode=False,
                 ignore_index=None,
                 scene_idxs=None,
                 file_client_args=dict(backend='disk'),
                 **kwargs) -> None:
        super().__init__(
            data_root=data_root,
            ann_file=ann_file,
            metainfo=metainfo,
            pipeline=pipeline,
            data_prefix=data_prefix,
            modality=modality,
            test_mode=test_mode,
            ignore_index=ignore_index,
            scene_idxs=scene_idxs,
            file_client_args=file_client_args,
            **kwargs)