scannet_dataset.py 17.5 KB
Newer Older
dingchang's avatar
dingchang committed
1
# Copyright (c) OpenMMLab. All rights reserved.
2
import numpy as np
3
import tempfile
4
import warnings
zhangwenwei's avatar
zhangwenwei committed
5
from os import path as osp
6

7
from mmdet3d.core import show_result, show_seg_result
wuyuefeng's avatar
wuyuefeng committed
8
from mmdet3d.core.bbox import DepthInstance3DBoxes
9
from mmdet.datasets import DATASETS
10
from mmseg.datasets import DATASETS as SEG_DATASETS
zhangwenwei's avatar
zhangwenwei committed
11
from .custom_3d import Custom3DDataset
12
from .custom_3d_seg import Custom3DSegDataset
13
from .pipelines import Compose
14
15
16


@DATASETS.register_module()
zhangwenwei's avatar
zhangwenwei committed
17
class ScanNetDataset(Custom3DDataset):
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
28
29
30
31
32
33
34
35
36
37
38

    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.
        modality (dict, optional): Modality to specify the sensor data used
            as input. Defaults to None.
        box_type_3d (str, optional): Type of 3D box of this dataset.
            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
39
40
41
            - 'LiDAR': Box in LiDAR coordinates.
            - 'Depth': Box in depth coordinates, usually for indoor dataset.
            - 'Camera': Box in camera coordinates.
wangtai's avatar
wangtai committed
42
43
44
45
46
        filter_empty_gt (bool, optional): Whether to filter empty GT.
            Defaults to True.
        test_mode (bool, optional): Whether the dataset is in test mode.
            Defaults to False.
    """
47
48
49
50
51
52
    CLASSES = ('cabinet', 'bed', 'chair', 'sofa', 'table', 'door', 'window',
               'bookshelf', 'picture', 'counter', 'desk', 'curtain',
               'refrigerator', 'showercurtrain', 'toilet', 'sink', 'bathtub',
               'garbagebin')

    def __init__(self,
zhangwenwei's avatar
zhangwenwei committed
53
                 data_root,
54
55
                 ann_file,
                 pipeline=None,
liyinhao's avatar
liyinhao committed
56
                 classes=None,
57
                 modality=dict(use_camera=False, use_depth=True),
58
                 box_type_3d='Depth',
wuyuefeng's avatar
Votenet  
wuyuefeng committed
59
                 filter_empty_gt=True,
zhangwenwei's avatar
zhangwenwei committed
60
                 test_mode=False):
61
62
63
64
65
66
67
68
69
        super().__init__(
            data_root=data_root,
            ann_file=ann_file,
            pipeline=pipeline,
            classes=classes,
            modality=modality,
            box_type_3d=box_type_3d,
            filter_empty_gt=filter_empty_gt,
            test_mode=test_mode)
70
71
72
73
74
75
76
77
78
79
80
        assert 'use_camera' in self.modality and \
               'use_depth' in self.modality
        assert self.modality['use_camera'] or self.modality['use_depth']

    def get_data_info(self, index):
        """Get data info according to the given index.

        Args:
            index (int): Index of the sample data to get.

        Returns:
81
            dict: Data information that will be passed to the data
82
83
84
85
86
                preprocessing pipelines. It includes the following keys:

                - sample_idx (str): Sample index.
                - pts_filename (str): Filename of point clouds.
                - file_name (str): Filename of point clouds.
87
                - img_prefix (str, optional): Prefix of image files.
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
                - img_info (dict, optional): Image info.
                - ann_info (dict): Annotation info.
        """
        info = self.data_infos[index]
        sample_idx = info['point_cloud']['lidar_idx']
        pts_filename = osp.join(self.data_root, info['pts_path'])
        input_dict = dict(sample_idx=sample_idx)

        if self.modality['use_depth']:
            input_dict['pts_filename'] = pts_filename
            input_dict['file_name'] = pts_filename

        if self.modality['use_camera']:
            img_info = []
            for img_path in info['img_paths']:
                img_info.append(
                    dict(filename=osp.join(self.data_root, img_path)))
            intrinsic = info['intrinsics']
            axis_align_matrix = self._get_axis_align_matrix(info)
            depth2img = []
            for extrinsic in info['extrinsics']:
                depth2img.append(
                    intrinsic @ np.linalg.inv(axis_align_matrix @ extrinsic))

            input_dict['img_prefix'] = None
            input_dict['img_info'] = img_info
            input_dict['depth2img'] = depth2img

        if not self.test_mode:
            annos = self.get_ann_info(index)
            input_dict['ann_info'] = annos
            if self.filter_empty_gt and ~(annos['gt_labels_3d'] != -1).any():
                return None
        return input_dict
122

liyinhao's avatar
liyinhao committed
123
    def get_ann_info(self, index):
124
125
126
127
128
129
        """Get annotation info according to the given index.

        Args:
            index (int): Index of the annotation data to get.

        Returns:
zhangwenwei's avatar
zhangwenwei committed
130
            dict: annotation information consists of the following keys:
131

132
                - gt_bboxes_3d (:obj:`DepthInstance3DBoxes`):
133
                    3D ground truth bboxes
wangtai's avatar
wangtai committed
134
135
136
                - gt_labels_3d (np.ndarray): Labels of ground truths.
                - pts_instance_mask_path (str): Path of instance masks.
                - pts_semantic_mask_path (str): Path of semantic masks.
137
                - axis_align_matrix (np.ndarray): Transformation matrix for
138
                    global scene alignment.
139
        """
140
        # Use index to get the annos, thus the evalhook could also use this api
liyinhao's avatar
liyinhao committed
141
        info = self.data_infos[index]
142
        if info['annos']['gt_num'] != 0:
liyinhao's avatar
liyinhao committed
143
144
145
            gt_bboxes_3d = info['annos']['gt_boxes_upright_depth'].astype(
                np.float32)  # k, 6
            gt_labels_3d = info['annos']['class'].astype(np.long)
146
        else:
liyinhao's avatar
liyinhao committed
147
            gt_bboxes_3d = np.zeros((0, 6), dtype=np.float32)
liyinhao's avatar
liyinhao committed
148
            gt_labels_3d = np.zeros((0, ), dtype=np.long)
wuyuefeng's avatar
wuyuefeng committed
149
150
151
152
153
154
155
156

        # to target box structure
        gt_bboxes_3d = DepthInstance3DBoxes(
            gt_bboxes_3d,
            box_dim=gt_bboxes_3d.shape[-1],
            with_yaw=False,
            origin=(0.5, 0.5, 0.5)).convert_to(self.box_mode_3d)

zhangwenwei's avatar
zhangwenwei committed
157
        pts_instance_mask_path = osp.join(self.data_root,
liyinhao's avatar
liyinhao committed
158
                                          info['pts_instance_mask_path'])
zhangwenwei's avatar
zhangwenwei committed
159
        pts_semantic_mask_path = osp.join(self.data_root,
liyinhao's avatar
liyinhao committed
160
                                          info['pts_semantic_mask_path'])
161

162
163
        axis_align_matrix = self._get_axis_align_matrix(info)

164
165
        anns_results = dict(
            gt_bboxes_3d=gt_bboxes_3d,
zhangwenwei's avatar
zhangwenwei committed
166
            gt_labels_3d=gt_labels_3d,
167
            pts_instance_mask_path=pts_instance_mask_path,
168
169
            pts_semantic_mask_path=pts_semantic_mask_path,
            axis_align_matrix=axis_align_matrix)
170
        return anns_results
liyinhao's avatar
liyinhao committed
171

172
173
174
    def prepare_test_data(self, index):
        """Prepare data for testing.

175
        We should take axis_align_matrix from self.data_infos since we need
176
177
178
179
180
181
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
            to align point clouds.

        Args:
            index (int): Index for accessing the target data.

        Returns:
            dict: Testing data dict of the corresponding index.
        """
        input_dict = self.get_data_info(index)
        # take the axis_align_matrix from data_infos
        input_dict['ann_info'] = dict(
            axis_align_matrix=self._get_axis_align_matrix(
                self.data_infos[index]))
        self.pre_pipeline(input_dict)
        example = self.pipeline(input_dict)
        return example

    @staticmethod
    def _get_axis_align_matrix(info):
        """Get axis_align_matrix from info. If not exist, return identity mat.

        Args:
            info (dict): one data info term.

        Returns:
            np.ndarray: 4x4 transformation matrix.
        """
        if 'axis_align_matrix' in info['annos'].keys():
            return info['annos']['axis_align_matrix'].astype(np.float32)
        else:
            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)

211
212
213
214
215
216
217
218
219
    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]),
220
            dict(type='GlobalAlignment', rotation_axis=2),
221
222
223
224
225
226
227
228
229
            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):
230
231
232
233
234
        """Results visualization.

        Args:
            results (list[dict]): List of bounding boxes results.
            out_dir (str): Output directory of visualization result.
235
            show (bool): Visualize the results online.
236
237
            pipeline (list[dict], optional): raw data loading for showing.
                Default: None.
238
        """
liyinhao's avatar
liyinhao committed
239
        assert out_dir is not None, 'Expect out_dir, got none.'
240
        pipeline = self._get_pipeline(pipeline)
liyinhao's avatar
liyinhao committed
241
242
243
244
        for i, result in enumerate(results):
            data_info = self.data_infos[i]
            pts_path = data_info['pts_path']
            file_name = osp.split(pts_path)[-1].split('.')[0]
245
            points = self._extract_data(i, pipeline, 'points').numpy()
246
            gt_bboxes = self.get_ann_info(i)['gt_bboxes_3d'].tensor.numpy()
liyinhao's avatar
liyinhao committed
247
            pred_bboxes = result['boxes_3d'].tensor.numpy()
248
249
            show_result(points, gt_bboxes, pred_bboxes, out_dir, file_name,
                        show)
250
251
252


@DATASETS.register_module()
253
@SEG_DATASETS.register_module()
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
class ScanNetSegDataset(Custom3DSegDataset):
    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.
275
        ignore_index (int, optional): The label index to be ignored, e.g.
276
277
278
279
280
281
282
283
284
285
286
287
288
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
314
315
316
317
318
319
320
321
322
323
            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.
    """
    CLASSES = ('wall', 'floor', 'cabinet', 'bed', 'chair', 'sofa', 'table',
               'door', 'window', 'bookshelf', 'picture', 'counter', 'desk',
               'curtain', 'refrigerator', 'showercurtrain', 'toilet', 'sink',
               'bathtub', 'otherfurniture')

    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))

    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],
    ]

    def __init__(self,
                 data_root,
                 ann_file,
                 pipeline=None,
                 classes=None,
                 palette=None,
                 modality=None,
                 test_mode=False,
                 ignore_index=None,
324
                 scene_idxs=None):
325
326
327
328
329
330
331
332
333
334

        super().__init__(
            data_root=data_root,
            ann_file=ann_file,
            pipeline=pipeline,
            classes=classes,
            palette=palette,
            modality=modality,
            test_mode=test_mode,
            ignore_index=ignore_index,
335
            scene_idxs=scene_idxs)
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356

    def get_ann_info(self, index):
        """Get annotation info according to the given index.

        Args:
            index (int): Index of the annotation data to get.

        Returns:
            dict: annotation information consists of the following keys:

                - pts_semantic_mask_path (str): Path of semantic masks.
        """
        # Use index to get the annos, thus the evalhook could also use this api
        info = self.data_infos[index]

        pts_semantic_mask_path = osp.join(self.data_root,
                                          info['pts_semantic_mask_path'])

        anns_results = dict(pts_semantic_mask_path=pts_semantic_mask_path)
        return anns_results

357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
    def _build_default_pipeline(self):
        """Build the default pipeline for this dataset."""
        pipeline = [
            dict(
                type='LoadPointsFromFile',
                coord_type='DEPTH',
                shift_height=False,
                use_color=True,
                load_dim=6,
                use_dim=[0, 1, 2, 3, 4, 5]),
            dict(
                type='LoadAnnotations3D',
                with_bbox_3d=False,
                with_label_3d=False,
                with_mask_3d=False,
                with_seg_3d=True),
            dict(
                type='PointSegClassMapping',
375
376
                valid_cat_ids=self.VALID_CLASS_IDS,
                max_cat_id=np.max(self.ALL_CLASS_IDS)),
377
378
379
380
381
382
383
384
385
            dict(
                type='DefaultFormatBundle3D',
                with_label=False,
                class_names=self.CLASSES),
            dict(type='Collect3D', keys=['points', 'pts_semantic_mask'])
        ]
        return Compose(pipeline)

    def show(self, results, out_dir, show=True, pipeline=None):
386
387
388
389
390
391
        """Results visualization.

        Args:
            results (list[dict]): List of bounding boxes results.
            out_dir (str): Output directory of visualization result.
            show (bool): Visualize the results online.
392
393
            pipeline (list[dict], optional): raw data loading for showing.
                Default: None.
394
395
        """
        assert out_dir is not None, 'Expect out_dir, got none.'
396
        pipeline = self._get_pipeline(pipeline)
397
398
399
400
        for i, result in enumerate(results):
            data_info = self.data_infos[i]
            pts_path = data_info['pts_path']
            file_name = osp.split(pts_path)[-1].split('.')[0]
401
402
403
            points, gt_sem_mask = self._extract_data(
                i, pipeline, ['points', 'pts_semantic_mask'], load_annos=True)
            points = points.numpy()
404
405
406
407
408
            pred_sem_mask = result['semantic_mask'].numpy()
            show_seg_result(points, gt_sem_mask,
                            pred_sem_mask, out_dir, file_name,
                            np.array(self.PALETTE), self.ignore_index, show)

409
410
    def get_scene_idxs(self, scene_idxs):
        """Compute scene_idxs for data sampling.
411

412
        We sample more times for scenes with more points.
413
414
415
416
417
418
        """
        # 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')

419
        return super().get_scene_idxs(scene_idxs)
420
421
422
423
424
425
426

    def format_results(self, results, txtfile_prefix=None):
        r"""Format the results to txt file. Refer to `ScanNet documentation
        <http://kaldir.vc.in.tum.de/scannet_benchmark/documentation>`_.

        Args:
            outputs (list[dict]): Testing results of the dataset.
427
            txtfile_prefix (str): The prefix of saved files. It includes
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
                the file path and the prefix of filename, e.g., "a/b/prefix".
                If not specified, a temp file will be created. Default: None.

        Returns:
            tuple: (outputs, tmp_dir), outputs is the detection results,
                tmp_dir is the temporal directory created for saving submission
                files when ``submission_prefix`` is not specified.
        """
        import mmcv

        if txtfile_prefix is None:
            tmp_dir = tempfile.TemporaryDirectory()
            txtfile_prefix = osp.join(tmp_dir.name, 'results')
        else:
            tmp_dir = None
        mmcv.mkdir_or_exist(txtfile_prefix)

        # need to map network output to original label idx
        pred2label = np.zeros(len(self.VALID_CLASS_IDS)).astype(np.int)
        for original_label, output_idx in self.label_map.items():
            if output_idx != self.ignore_index:
                pred2label[output_idx] = original_label

        outputs = []
        for i, result in enumerate(results):
            info = self.data_infos[i]
            sample_idx = info['point_cloud']['lidar_idx']
            pred_sem_mask = result['semantic_mask'].numpy().astype(np.int)
            pred_label = pred2label[pred_sem_mask]
            curr_file = f'{txtfile_prefix}/{sample_idx}.txt'
            np.savetxt(curr_file, pred_label, fmt='%d')
            outputs.append(dict(seg_mask=pred_label))

        return outputs, tmp_dir