sunrgbd_dataset.py 10.9 KB
Newer Older
dingchang's avatar
dingchang committed
1
# Copyright (c) OpenMMLab. All rights reserved.
2
from collections import OrderedDict
zhangwenwei's avatar
zhangwenwei committed
3
from os import path as osp
liyinhao's avatar
liyinhao committed
4

5
6
import numpy as np

7
from mmdet3d.core import show_multi_modality_result, show_result
wuyuefeng's avatar
wuyuefeng committed
8
from mmdet3d.core.bbox import DepthInstance3DBoxes
9
from mmdet.core import eval_map
10
from .builder import DATASETS
zhangwenwei's avatar
zhangwenwei committed
11
from .custom_3d import Custom3DDataset
12
from .pipelines import Compose
liyinhao's avatar
liyinhao committed
13
14
15


@DATASETS.register_module()
zhangwenwei's avatar
zhangwenwei committed
16
class SUNRGBDDataset(Custom3DDataset):
zhangwenwei's avatar
zhangwenwei committed
17
    r"""SUNRGBD Dataset.
liyinhao's avatar
liyinhao committed
18

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

zhangwenwei's avatar
zhangwenwei committed
21
22
    See the `download page <http://rgbd.cs.princeton.edu/challenge.html>`_
    for data downloading.
wangtai's avatar
wangtai committed
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37

    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
38
39
40
            - 'LiDAR': Box in LiDAR coordinates.
            - 'Depth': Box in depth coordinates, usually for indoor dataset.
            - 'Camera': Box in camera coordinates.
wangtai's avatar
wangtai committed
41
42
43
44
45
        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.
    """
liyinhao's avatar
liyinhao committed
46
47
48
49
    CLASSES = ('bed', 'table', 'sofa', 'chair', 'toilet', 'desk', 'dresser',
               'night_stand', 'bookshelf', 'bathtub')

    def __init__(self,
zhangwenwei's avatar
zhangwenwei committed
50
                 data_root,
liyinhao's avatar
liyinhao committed
51
52
                 ann_file,
                 pipeline=None,
liyinhao's avatar
liyinhao committed
53
                 classes=None,
54
                 modality=dict(use_camera=True, use_lidar=True),
55
                 box_type_3d='Depth',
wuyuefeng's avatar
Votenet  
wuyuefeng committed
56
                 filter_empty_gt=True,
57
58
                 test_mode=False,
                 **kwargs):
59
60
61
62
63
64
65
66
        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,
67
68
            test_mode=test_mode,
            **kwargs)
69
70
        assert 'use_camera' in self.modality and \
            'use_lidar' in self.modality
71
72
73
74
75
76
77
78
79
        assert self.modality['use_camera'] or self.modality['use_lidar']

    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:
80
            dict: Data information that will be passed to the data
81
82
83
84
85
                preprocessing pipelines. It includes the following keys:

                - sample_idx (str): Sample index.
                - pts_filename (str, optional): Filename of point clouds.
                - file_name (str, optional): Filename of point clouds.
86
                - img_prefix (str, optional): Prefix of image files.
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
                - img_info (dict, optional): Image info.
                - calib (dict, optional): Camera calibration info.
                - ann_info (dict): Annotation info.
        """
        info = self.data_infos[index]
        sample_idx = info['point_cloud']['lidar_idx']
        assert info['point_cloud']['lidar_idx'] == info['image']['image_idx']
        input_dict = dict(sample_idx=sample_idx)

        if self.modality['use_lidar']:
            pts_filename = osp.join(self.data_root, info['pts_path'])
            input_dict['pts_filename'] = pts_filename
            input_dict['file_name'] = pts_filename

        if self.modality['use_camera']:
102
103
104
            img_filename = osp.join(
                osp.join(self.data_root, 'sunrgbd_trainval'),
                info['image']['image_path'])
105
106
107
            input_dict['img_prefix'] = None
            input_dict['img_info'] = dict(filename=img_filename)
            calib = info['calib']
108
109
110
111
112
113
            rt_mat = calib['Rt']
            # follow Coord3DMode.convert_point
            rt_mat = np.array([[1, 0, 0], [0, 0, -1], [0, 1, 0]
                               ]) @ rt_mat.transpose(1, 0)
            depth2img = calib['K'] @ rt_mat
            input_dict['depth2img'] = depth2img
114
115
116
117
118
119
120

        if not self.test_mode:
            annos = self.get_ann_info(index)
            input_dict['ann_info'] = annos
            if self.filter_empty_gt and len(annos['gt_bboxes_3d']) == 0:
                return None
        return input_dict
liyinhao's avatar
liyinhao committed
121

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

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

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

131
                - gt_bboxes_3d (:obj:`DepthInstance3DBoxes`):
132
                    3D ground truth bboxes
wangtai's avatar
wangtai committed
133
134
135
                - 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.
136
        """
liyinhao's avatar
liyinhao committed
137
        # Use index to get the annos, thus the evalhook could also use this api
liyinhao's avatar
liyinhao committed
138
        info = self.data_infos[index]
liyinhao's avatar
liyinhao committed
139
        if info['annos']['gt_num'] != 0:
liyinhao's avatar
liyinhao committed
140
141
            gt_bboxes_3d = info['annos']['gt_boxes_upright_depth'].astype(
                np.float32)  # k, 6
WRH's avatar
WRH committed
142
            gt_labels_3d = info['annos']['class'].astype(np.int64)
liyinhao's avatar
liyinhao committed
143
        else:
liyinhao's avatar
liyinhao committed
144
            gt_bboxes_3d = np.zeros((0, 7), dtype=np.float32)
WRH's avatar
WRH committed
145
            gt_labels_3d = np.zeros((0, ), dtype=np.int64)
liyinhao's avatar
liyinhao committed
146

wuyuefeng's avatar
wuyuefeng committed
147
148
149
150
        # to target box structure
        gt_bboxes_3d = DepthInstance3DBoxes(
            gt_bboxes_3d, origin=(0.5, 0.5, 0.5)).convert_to(self.box_mode_3d)

liyinhao's avatar
liyinhao committed
151
        anns_results = dict(
liyinhao's avatar
liyinhao committed
152
            gt_bboxes_3d=gt_bboxes_3d, gt_labels_3d=gt_labels_3d)
153
154
155
156
157
158
159
160
161

        if self.modality['use_camera']:
            if info['annos']['gt_num'] != 0:
                gt_bboxes_2d = info['annos']['bbox'].astype(np.float32)
            else:
                gt_bboxes_2d = np.zeros((0, 4), dtype=np.float32)
            anns_results['bboxes'] = gt_bboxes_2d
            anns_results['labels'] = gt_labels_3d

liyinhao's avatar
liyinhao committed
162
        return anns_results
liyinhao's avatar
liyinhao committed
163

164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
    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]),
            dict(
                type='DefaultFormatBundle3D',
                class_names=self.CLASSES,
                with_label=False),
            dict(type='Collect3D', keys=['points'])
        ]
        if self.modality['use_camera']:
            pipeline.insert(0, dict(type='LoadImageFromFile'))
        return Compose(pipeline)

    def show(self, results, out_dir, show=True, pipeline=None):
184
185
186
187
188
        """Results visualization.

        Args:
            results (list[dict]): List of bounding boxes results.
            out_dir (str): Output directory of visualization result.
189
            show (bool): Visualize the results online.
190
191
            pipeline (list[dict], optional): raw data loading for showing.
                Default: None.
192
        """
liyinhao's avatar
liyinhao committed
193
        assert out_dir is not None, 'Expect out_dir, got none.'
194
        pipeline = self._get_pipeline(pipeline)
liyinhao's avatar
liyinhao committed
195
196
197
198
        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]
199
200
            points, img_metas, img = self._extract_data(
                i, pipeline, ['points', 'img_metas', 'img'])
201
202
            # scale colors to [0, 255]
            points = points.numpy()
liyinhao's avatar
liyinhao committed
203
            points[:, 3:] *= 255
204
205

            gt_bboxes = self.get_ann_info(i)['gt_bboxes_3d'].tensor.numpy()
liyinhao's avatar
liyinhao committed
206
            pred_bboxes = result['boxes_3d'].tensor.numpy()
207
208
209
210
            show_result(points, gt_bboxes.copy(), pred_bboxes.copy(), out_dir,
                        file_name, show)

            # multi-modality visualization
211
            if self.modality['use_camera']:
212
213
214
                img = img.numpy()
                # need to transpose channel to first dim
                img = img.transpose(1, 2, 0)
215
216
217
218
219
220
221
222
                pred_bboxes = DepthInstance3DBoxes(
                    pred_bboxes, origin=(0.5, 0.5, 0))
                gt_bboxes = DepthInstance3DBoxes(
                    gt_bboxes, origin=(0.5, 0.5, 0))
                show_multi_modality_result(
                    img,
                    gt_bboxes,
                    pred_bboxes,
223
                    None,
224
225
                    out_dir,
                    file_name,
226
                    box_mode='depth',
227
                    img_metas=img_metas,
228
                    show=show)
229
230
231
232
233
234
235
236

    def evaluate(self,
                 results,
                 metric=None,
                 iou_thr=(0.25, 0.5),
                 iou_thr_2d=(0.5, ),
                 logger=None,
                 show=False,
237
238
239
240
241
                 out_dir=None,
                 pipeline=None):
        """Evaluate.

        Evaluation in indoor protocol.
242

243
244
        Args:
            results (list[dict]): List of results.
245
246
247
248
249
250
251
            metric (str | list[str], optional): Metrics to be evaluated.
                Default: None.
            iou_thr (list[float], optional): AP IoU thresholds for 3D
                evaluation. Default: (0.25, 0.5).
            iou_thr_2d (list[float], optional): AP IoU thresholds for 2D
                evaluation. Default: (0.5, ).
            show (bool, optional): Whether to visualize.
252
                Default: False.
253
            out_dir (str, optional): Path to save the visualization results.
254
255
256
257
258
259
260
                Default: None.
            pipeline (list[dict], optional): raw data loading for showing.
                Default: None.

        Returns:
            dict: Evaluation results.
        """
261
262
263
        # evaluate 3D detection performance
        if isinstance(results[0], dict):
            return super().evaluate(results, metric, iou_thr, logger, show,
264
                                    out_dir, pipeline)
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
        # evaluate 2D detection performance
        else:
            eval_results = OrderedDict()
            annotations = [self.get_ann_info(i) for i in range(len(self))]
            iou_thr_2d = (iou_thr_2d) if isinstance(iou_thr_2d,
                                                    float) else iou_thr_2d
            for iou_thr_2d_single in iou_thr_2d:
                mean_ap, _ = eval_map(
                    results,
                    annotations,
                    scale_ranges=None,
                    iou_thr=iou_thr_2d_single,
                    dataset=self.CLASSES,
                    logger=logger)
                eval_results['mAP_' + str(iou_thr_2d_single)] = mean_ap
            return eval_results