scannet_dataset.py 12.8 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

ZCMax's avatar
ZCMax committed
8
from mmdet3d.core import show_result
wuyuefeng's avatar
wuyuefeng committed
9
from mmdet3d.core.bbox import DepthInstance3DBoxes
10
from mmdet3d.registry import DATASETS
jshilong's avatar
jshilong committed
11
from .det3d_dataset import Det3DDataset
12
from .pipelines import Compose
ZCMax's avatar
ZCMax committed
13
from .seg3d_dataset import Seg3DDataset
14
15
16


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

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

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

    Args:
        data_root (str): Path of dataset root.
        ann_file (str): Path of annotation file.
jshilong's avatar
jshilong committed
28
29
30
31
32
33
34
        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
35
            Defaults to None.
jshilong's avatar
jshilong committed
36
        modality (dict): Modality to specify the sensor data used
wangtai's avatar
wangtai committed
37
            as input. Defaults to None.
jshilong's avatar
jshilong committed
38
        box_type_3d (str): Type of 3D box of this dataset.
wangtai's avatar
wangtai committed
39
40
41
42
            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
43
44
45
            - 'LiDAR': Box in LiDAR coordinates.
            - 'Depth': Box in depth coordinates, usually for indoor dataset.
            - 'Camera': Box in camera coordinates.
jshilong's avatar
jshilong committed
46
        filter_empty_gt (bool): Whether to filter empty GT.
wangtai's avatar
wangtai committed
47
            Defaults to True.
jshilong's avatar
jshilong committed
48
        test_mode (bool): Whether the dataset is in test mode.
wangtai's avatar
wangtai committed
49
50
            Defaults to False.
    """
jshilong's avatar
jshilong committed
51
52
53
54
55
56
    METAINFO = {
        'CLASSES':
        ('cabinet', 'bed', 'chair', 'sofa', 'table', 'door', 'window',
         'bookshelf', 'picture', 'counter', 'desk', 'curtain', 'refrigerator',
         'showercurtrain', 'toilet', 'sink', 'bathtub', 'garbagebin')
    }
57
58

    def __init__(self,
jshilong's avatar
jshilong committed
59
60
61
62
63
                 data_root: str,
                 ann_file: str,
                 metainfo: dict = None,
                 data_prefix: dict = dict(
                     pts='points',
64
                     pts_instance_mask='instance_mask',
jshilong's avatar
jshilong committed
65
66
67
68
69
70
                     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,
71
                 **kwargs):
72
73
74
        super().__init__(
            data_root=data_root,
            ann_file=ann_file,
jshilong's avatar
jshilong committed
75
76
            metainfo=metainfo,
            data_prefix=data_prefix,
77
78
79
80
            pipeline=pipeline,
            modality=modality,
            box_type_3d=box_type_3d,
            filter_empty_gt=filter_empty_gt,
81
82
            test_mode=test_mode,
            **kwargs)
83
        assert 'use_camera' in self.modality and \
jshilong's avatar
jshilong committed
84
85
               'use_lidar' in self.modality
        assert self.modality['use_camera'] or self.modality['use_lidar']
86

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

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

        Returns:
jshilong's avatar
jshilong committed
95
            np.ndarray: 4x4 transformation matrix.
96
        """
jshilong's avatar
jshilong committed
97
98
        if 'axis_align_matrix' in info:
            return np.array(info['axis_align_matrix'])
99
        else:
jshilong's avatar
jshilong committed
100
101
102
103
            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
104

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

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

        Args:
jshilong's avatar
jshilong committed
112
            info (dict): Raw info dict.
113
114

        Returns:
jshilong's avatar
jshilong committed
115
116
            dict: Data information that will be passed to the data
            preprocessing pipelines. It includes the following keys:
117
        """
jshilong's avatar
jshilong committed
118
119
120
121
122
123
124
125
126
127
128
129
130
        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`
131
132

        Args:
jshilong's avatar
jshilong committed
133
            info (dict): Info dict.
134
135

        Returns:
jshilong's avatar
jshilong committed
136
            dict: Processed `ann_info`
137
        """
jshilong's avatar
jshilong committed
138
        ann_info = super().parse_ann_info(info)
139
140
141
142
        # empty gt
        if ann_info is None:
            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
143
        # to target box structure
144

jshilong's avatar
jshilong committed
145
146
147
148
149
150
151
        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
152

153
154
155
156
157
158
159
160
161
    def _build_default_pipeline(self):
        """Build the default pipeline for this dataset."""
        pipeline = [
            dict(
                type='LoadPointsFromFile',
                coord_type='DEPTH',
                shift_height=False,
                load_dim=6,
                use_dim=[0, 1, 2]),
162
            dict(type='GlobalAlignment', rotation_axis=2),
163
164
165
166
167
168
169
170
171
            dict(
                type='DefaultFormatBundle3D',
                class_names=self.CLASSES,
                with_label=False),
            dict(type='Collect3D', keys=['points'])
        ]
        return Compose(pipeline)

    def show(self, results, out_dir, show=True, pipeline=None):
172
173
174
175
176
        """Results visualization.

        Args:
            results (list[dict]): List of bounding boxes results.
            out_dir (str): Output directory of visualization result.
177
            show (bool): Visualize the results online.
178
179
            pipeline (list[dict], optional): raw data loading for showing.
                Default: None.
180
        """
liyinhao's avatar
liyinhao committed
181
        assert out_dir is not None, 'Expect out_dir, got none.'
182
        pipeline = self._get_pipeline(pipeline)
liyinhao's avatar
liyinhao committed
183
        for i, result in enumerate(results):
jshilong's avatar
jshilong committed
184
185
            data_info = self.get_data_info[i]
            pts_path = data_info['lidar_points']['lidar_path']
liyinhao's avatar
liyinhao committed
186
            file_name = osp.split(pts_path)[-1].split('.')[0]
187
            points = self._extract_data(i, pipeline, 'points').numpy()
188
            gt_bboxes = self.get_ann_info(i)['gt_bboxes_3d'].tensor.numpy()
liyinhao's avatar
liyinhao committed
189
            pred_bboxes = result['boxes_3d'].tensor.numpy()
190
191
            show_result(points, gt_bboxes, pred_bboxes, out_dir, file_name,
                        show)
192
193
194


@DATASETS.register_module()
ZCMax's avatar
ZCMax committed
195
class ScanNetSegDataset(Seg3DDataset):
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
    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.
216
        ignore_index (int, optional): The label index to be ignored, e.g.
217
218
219
220
221
222
            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
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
    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)),
    }
256
257

    def __init__(self,
ZCMax's avatar
ZCMax committed
258
259
260
261
262
263
264
                 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),
265
                 ignore_index=None,
266
                 scene_idxs=None,
ZCMax's avatar
ZCMax committed
267
268
                 test_mode=False,
                 **kwargs) -> None:
269
270
271
        super().__init__(
            data_root=data_root,
            ann_file=ann_file,
ZCMax's avatar
ZCMax committed
272
273
            metainfo=metainfo,
            data_prefix=data_prefix,
274
275
276
            pipeline=pipeline,
            modality=modality,
            ignore_index=ignore_index,
277
            scene_idxs=scene_idxs,
ZCMax's avatar
ZCMax committed
278
            test_mode=test_mode,
279
            **kwargs)
280

281
282
    def get_scene_idxs(self, scene_idxs):
        """Compute scene_idxs for data sampling.
283

284
        We sample more times for scenes with more points.
285
286
287
288
289
290
        """
        # 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')

291
        return super().get_scene_idxs(scene_idxs)
292

293
294

@DATASETS.register_module()
ZCMax's avatar
ZCMax committed
295
class ScanNetInstanceSegDataset(Seg3DDataset):
296

ZCMax's avatar
ZCMax committed
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
    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))
    }
329

ZCMax's avatar
ZCMax committed
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
    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)