nuscenes_dataset.py 24.6 KB
Newer Older
zhangwenwei's avatar
zhangwenwei committed
1
2
3
import mmcv
import numpy as np
import pyquaternion
zhangwenwei's avatar
zhangwenwei committed
4
import tempfile
zhangwenwei's avatar
zhangwenwei committed
5
from nuscenes.utils.data_classes import Box as NuScenesBox
zhangwenwei's avatar
zhangwenwei committed
6
from os import path as osp
zhangwenwei's avatar
zhangwenwei committed
7
8

from mmdet.datasets import DATASETS
liyinhao's avatar
liyinhao committed
9
from ..core import show_result
10
from ..core.bbox import Box3DMode, Coord3DMode, LiDARInstance3DBoxes
zhangwenwei's avatar
zhangwenwei committed
11
from .custom_3d import Custom3DDataset
12
from .pipelines import Compose
zhangwenwei's avatar
zhangwenwei committed
13
14


15
@DATASETS.register_module()
zhangwenwei's avatar
zhangwenwei committed
16
class NuScenesDataset(Custom3DDataset):
wangtai's avatar
wangtai committed
17
    r"""NuScenes Dataset.
wangtai's avatar
wangtai committed
18
19
20

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

zhangwenwei's avatar
zhangwenwei committed
21
22
    Please refer to `NuScenes Dataset <https://www.nuscenes.org/download>`_
    for data downloading.
wangtai's avatar
wangtai committed
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39

    Args:
        ann_file (str): Path of annotation file.
        pipeline (list[dict], optional): Pipeline used for data processing.
            Defaults to None.
        data_root (str): Path of dataset root.
        classes (tuple[str], optional): Classes used in the dataset.
            Defaults to None.
        load_interval (int, optional): Interval of loading the dataset. It is
            used to uniformly sample the dataset. Defaults to 1.
        with_velocity (bool, optional): Whether include velocity prediction
            into the experiments. Defaults to True.
        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`.
yinchimaoliang's avatar
yinchimaoliang committed
40
            Defaults to 'LiDAR' 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.
wangtai's avatar
wangtai committed
44
45
46
47
48
49
        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.
        eval_version (bool, optional): Configuration version of evaluation.
            Defaults to  'detection_cvpr_2019'.
yinchimaoliang's avatar
yinchimaoliang committed
50
51
        use_valid_flag (bool): Whether to use `use_valid_flag` key in the info
            file as mask to filter gt_boxes and gt_names. Defaults to False.
wangtai's avatar
wangtai committed
52
    """
zhangwenwei's avatar
zhangwenwei committed
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
    NameMapping = {
        'movable_object.barrier': 'barrier',
        'vehicle.bicycle': 'bicycle',
        'vehicle.bus.bendy': 'bus',
        'vehicle.bus.rigid': 'bus',
        'vehicle.car': 'car',
        'vehicle.construction': 'construction_vehicle',
        'vehicle.motorcycle': 'motorcycle',
        'human.pedestrian.adult': 'pedestrian',
        'human.pedestrian.child': 'pedestrian',
        'human.pedestrian.construction_worker': 'pedestrian',
        'human.pedestrian.police_officer': 'pedestrian',
        'movable_object.trafficcone': 'traffic_cone',
        'vehicle.trailer': 'trailer',
        'vehicle.truck': 'truck'
    }
    DefaultAttribute = {
        'car': 'vehicle.parked',
        'pedestrian': 'pedestrian.moving',
        'trailer': 'vehicle.parked',
        'truck': 'vehicle.parked',
        'bus': 'vehicle.moving',
        'motorcycle': 'cycle.without_rider',
        'construction_vehicle': 'vehicle.parked',
        'bicycle': 'cycle.without_rider',
        'barrier': '',
        'traffic_cone': '',
    }
    AttrMapping = {
        'cycle.with_rider': 0,
        'cycle.without_rider': 1,
        'pedestrian.moving': 2,
        'pedestrian.standing': 3,
        'pedestrian.sitting_lying_down': 4,
        'vehicle.moving': 5,
        'vehicle.parked': 6,
        'vehicle.stopped': 7,
    }
    AttrMapping_rev = [
        'cycle.with_rider',
        'cycle.without_rider',
        'pedestrian.moving',
        'pedestrian.standing',
        'pedestrian.sitting_lying_down',
        'vehicle.moving',
        'vehicle.parked',
        'vehicle.stopped',
    ]
    CLASSES = ('car', 'truck', 'trailer', 'bus', 'construction_vehicle',
               'bicycle', 'motorcycle', 'pedestrian', 'traffic_cone',
               'barrier')

    def __init__(self,
                 ann_file,
                 pipeline=None,
zhangwenwei's avatar
zhangwenwei committed
108
109
                 data_root=None,
                 classes=None,
zhangwenwei's avatar
zhangwenwei committed
110
111
112
                 load_interval=1,
                 with_velocity=True,
                 modality=None,
113
114
115
                 box_type_3d='LiDAR',
                 filter_empty_gt=True,
                 test_mode=False,
yinchimaoliang's avatar
yinchimaoliang committed
116
117
                 eval_version='detection_cvpr_2019',
                 use_valid_flag=False):
zhangwenwei's avatar
zhangwenwei committed
118
        self.load_interval = load_interval
yinchimaoliang's avatar
yinchimaoliang committed
119
        self.use_valid_flag = use_valid_flag
zhangwenwei's avatar
zhangwenwei committed
120
121
122
123
124
125
        super().__init__(
            data_root=data_root,
            ann_file=ann_file,
            pipeline=pipeline,
            classes=classes,
            modality=modality,
126
127
            box_type_3d=box_type_3d,
            filter_empty_gt=filter_empty_gt,
zhangwenwei's avatar
zhangwenwei committed
128
            test_mode=test_mode)
zhangwenwei's avatar
zhangwenwei committed
129
130
131
132
133

        self.with_velocity = with_velocity
        self.eval_version = eval_version
        from nuscenes.eval.detection.config import config_factory
        self.eval_detection_configs = config_factory(self.eval_version)
zhangwenwei's avatar
zhangwenwei committed
134
135
        if self.modality is None:
            self.modality = dict(
zhangwenwei's avatar
zhangwenwei committed
136
137
138
139
140
141
142
                use_camera=False,
                use_lidar=True,
                use_radar=False,
                use_map=False,
                use_external=False,
            )

yinchimaoliang's avatar
yinchimaoliang committed
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
    def get_cat_ids(self, idx):
        """Get category distribution of single scene.

        Args:
            idx (int): Index of the data_info.

        Returns:
            dict[list]: for each category, if the current scene
                contains such boxes, store a list containing idx,
                otherwise, store empty list.
        """
        info = self.data_infos[idx]
        if self.use_valid_flag:
            mask = info['valid_flag']
            gt_names = set(info['gt_names'][mask])
        else:
            gt_names = set(info['gt_names'])
160
161

        cat_ids = []
yinchimaoliang's avatar
yinchimaoliang committed
162
163
        for name in gt_names:
            if name in self.CLASSES:
164
165
                cat_ids.append(self.cat2id[name])
        return cat_ids
yinchimaoliang's avatar
yinchimaoliang committed
166

zhangwenwei's avatar
zhangwenwei committed
167
    def load_annotations(self, ann_file):
168
169
170
171
172
173
174
175
        """Load annotations from ann_file.

        Args:
            ann_file (str): Path of the annotation file.

        Returns:
            list[dict]: List of annotations sorted by timestamps.
        """
zhangwenwei's avatar
zhangwenwei committed
176
177
178
179
180
181
        data = mmcv.load(ann_file)
        data_infos = list(sorted(data['infos'], key=lambda e: e['timestamp']))
        data_infos = data_infos[::self.load_interval]
        self.metadata = data['metadata']
        self.version = self.metadata['version']
        return data_infos
zhangwenwei's avatar
zhangwenwei committed
182

zhangwenwei's avatar
zhangwenwei committed
183
    def get_data_info(self, index):
184
185
186
187
188
189
        """Get data info according to the given index.

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

        Returns:
zhangwenwei's avatar
zhangwenwei committed
190
191
            dict: Data information that will be passed to the data \
                preprocessing pipelines. It includes the following keys:
192

wangtai's avatar
wangtai committed
193
194
195
196
197
198
199
200
                - sample_idx (str): Sample index.
                - pts_filename (str): Filename of point clouds.
                - sweeps (list[dict]): Infos of sweeps.
                - timestamp (float): Sample timestamp.
                - img_filename (str, optional): Image filename.
                - lidar2img (list[np.ndarray], optional): Transformations \
                    from lidar to different cameras.
                - ann_info (dict): Annotation info.
201
        """
zhangwenwei's avatar
zhangwenwei committed
202
        info = self.data_infos[index]
zhangwenwei's avatar
zhangwenwei committed
203
        # standard protocal modified from SECOND.Pytorch
zhangwenwei's avatar
zhangwenwei committed
204
205
        input_dict = dict(
            sample_idx=info['token'],
zhangwenwei's avatar
zhangwenwei committed
206
207
208
            pts_filename=info['lidar_path'],
            sweeps=info['sweeps'],
            timestamp=info['timestamp'] / 1e6,
zhangwenwei's avatar
zhangwenwei committed
209
210
211
212
213
214
        )

        if self.modality['use_camera']:
            image_paths = []
            lidar2img_rts = []
            for cam_type, cam_info in info['cams'].items():
zhangwenwei's avatar
zhangwenwei committed
215
                image_paths.append(cam_info['data_path'])
zhangwenwei's avatar
zhangwenwei committed
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
                # obtain lidar to image transformation matrix
                lidar2cam_r = np.linalg.inv(cam_info['sensor2lidar_rotation'])
                lidar2cam_t = cam_info[
                    'sensor2lidar_translation'] @ lidar2cam_r.T
                lidar2cam_rt = np.eye(4)
                lidar2cam_rt[:3, :3] = lidar2cam_r.T
                lidar2cam_rt[3, :3] = -lidar2cam_t
                intrinsic = cam_info['cam_intrinsic']
                viewpad = np.eye(4)
                viewpad[:intrinsic.shape[0], :intrinsic.shape[1]] = intrinsic
                lidar2img_rt = (viewpad @ lidar2cam_rt.T)
                lidar2img_rts.append(lidar2img_rt)

            input_dict.update(
                dict(
zhangwenwei's avatar
zhangwenwei committed
231
                    img_filename=image_paths,
zhangwenwei's avatar
zhangwenwei committed
232
233
234
                    lidar2img=lidar2img_rts,
                ))

zhangwenwei's avatar
zhangwenwei committed
235
        if not self.test_mode:
zhangwenwei's avatar
zhangwenwei committed
236
            annos = self.get_ann_info(index)
zhangwenwei's avatar
zhangwenwei committed
237
            input_dict['ann_info'] = annos
zhangwenwei's avatar
zhangwenwei committed
238
239
240
241

        return input_dict

    def get_ann_info(self, index):
242
243
244
245
246
247
        """Get annotation info according to the given index.

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

        Returns:
wangtai's avatar
wangtai committed
248
            dict: Annotation information consists of the following keys:
249

zhangwenwei's avatar
zhangwenwei committed
250
                - gt_bboxes_3d (:obj:`LiDARInstance3DBoxes`): \
251
                    3D ground truth bboxes
wangtai's avatar
wangtai committed
252
253
                - gt_labels_3d (np.ndarray): Labels of ground truths.
                - gt_names (list[str]): Class names of ground truths.
254
        """
zhangwenwei's avatar
zhangwenwei committed
255
        info = self.data_infos[index]
zhangwenwei's avatar
zhangwenwei committed
256
        # filter out bbox containing no points
yinchimaoliang's avatar
yinchimaoliang committed
257
258
259
260
        if self.use_valid_flag:
            mask = info['valid_flag']
        else:
            mask = info['num_lidar_pts'] > 0
zhangwenwei's avatar
zhangwenwei committed
261
262
        gt_bboxes_3d = info['gt_boxes'][mask]
        gt_names_3d = info['gt_names'][mask]
zhangwenwei's avatar
zhangwenwei committed
263
264
265
266
267
268
269
        gt_labels_3d = []
        for cat in gt_names_3d:
            if cat in self.CLASSES:
                gt_labels_3d.append(self.CLASSES.index(cat))
            else:
                gt_labels_3d.append(-1)
        gt_labels_3d = np.array(gt_labels_3d)
zhangwenwei's avatar
zhangwenwei committed
270
271
272
273
274
275
276

        if self.with_velocity:
            gt_velocity = info['gt_velocity'][mask]
            nan_mask = np.isnan(gt_velocity[:, 0])
            gt_velocity[nan_mask] = [0.0, 0.0]
            gt_bboxes_3d = np.concatenate([gt_bboxes_3d, gt_velocity], axis=-1)

wangtai's avatar
wangtai committed
277
        # the nuscenes box center is [0.5, 0.5, 0.5], we change it to be
wuyuefeng's avatar
wuyuefeng committed
278
        # the same as KITTI (0.5, 0.5, 0)
zhangwenwei's avatar
zhangwenwei committed
279
280
281
        gt_bboxes_3d = LiDARInstance3DBoxes(
            gt_bboxes_3d,
            box_dim=gt_bboxes_3d.shape[-1],
wuyuefeng's avatar
wuyuefeng committed
282
            origin=(0.5, 0.5, 0.5)).convert_to(self.box_mode_3d)
zhangwenwei's avatar
zhangwenwei committed
283

zhangwenwei's avatar
zhangwenwei committed
284
285
        anns_results = dict(
            gt_bboxes_3d=gt_bboxes_3d,
zhangwenwei's avatar
zhangwenwei committed
286
            gt_labels_3d=gt_labels_3d,
liyinhao's avatar
liyinhao committed
287
            gt_names=gt_names_3d)
zhangwenwei's avatar
zhangwenwei committed
288
289
290
        return anns_results

    def _format_bbox(self, results, jsonfile_prefix=None):
291
292
293
294
295
296
297
298
299
300
301
        """Convert the results to the standard format.

        Args:
            results (list[dict]): Testing results of the dataset.
            jsonfile_prefix (str): The prefix of the output jsonfile.
                You can specify the output directory/filename by
                modifying the jsonfile_prefix. Default: None.

        Returns:
            str: Path of the output json file.
        """
zhangwenwei's avatar
zhangwenwei committed
302
        nusc_annos = {}
zhangwenwei's avatar
zhangwenwei committed
303
        mapped_class_names = self.CLASSES
zhangwenwei's avatar
zhangwenwei committed
304

zhangwenwei's avatar
zhangwenwei committed
305
        print('Start to convert detection format...')
zhangwenwei's avatar
zhangwenwei committed
306
        for sample_id, det in enumerate(mmcv.track_iter_progress(results)):
zhangwenwei's avatar
zhangwenwei committed
307
            annos = []
zhangwenwei's avatar
zhangwenwei committed
308
309
310
311
            boxes = output_to_nusc_box(det)
            sample_token = self.data_infos[sample_id]['token']
            boxes = lidar_nusc_box_to_global(self.data_infos[sample_id], boxes,
                                             mapped_class_names,
zhangwenwei's avatar
zhangwenwei committed
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
                                             self.eval_detection_configs,
                                             self.eval_version)
            for i, box in enumerate(boxes):
                name = mapped_class_names[box.label]
                if np.sqrt(box.velocity[0]**2 + box.velocity[1]**2) > 0.2:
                    if name in [
                            'car',
                            'construction_vehicle',
                            'bus',
                            'truck',
                            'trailer',
                    ]:
                        attr = 'vehicle.moving'
                    elif name in ['bicycle', 'motorcycle']:
                        attr = 'cycle.with_rider'
                    else:
                        attr = NuScenesDataset.DefaultAttribute[name]
                else:
                    if name in ['pedestrian']:
                        attr = 'pedestrian.standing'
                    elif name in ['bus']:
                        attr = 'vehicle.stopped'
                    else:
                        attr = NuScenesDataset.DefaultAttribute[name]

                nusc_anno = dict(
zhangwenwei's avatar
zhangwenwei committed
338
                    sample_token=sample_token,
zhangwenwei's avatar
zhangwenwei committed
339
340
341
342
343
344
345
346
                    translation=box.center.tolist(),
                    size=box.wlh.tolist(),
                    rotation=box.orientation.elements.tolist(),
                    velocity=box.velocity[:2].tolist(),
                    detection_name=name,
                    detection_score=box.score,
                    attribute_name=attr)
                annos.append(nusc_anno)
zhangwenwei's avatar
zhangwenwei committed
347
            nusc_annos[sample_token] = annos
zhangwenwei's avatar
zhangwenwei committed
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
        nusc_submissions = {
            'meta': self.modality,
            'results': nusc_annos,
        }

        mmcv.mkdir_or_exist(jsonfile_prefix)
        res_path = osp.join(jsonfile_prefix, 'results_nusc.json')
        print('Results writes to', res_path)
        mmcv.dump(nusc_submissions, res_path)
        return res_path

    def _evaluate_single(self,
                         result_path,
                         logger=None,
                         metric='bbox',
                         result_name='pts_bbox'):
364
365
366
367
368
369
370
371
372
373
374
375
376
        """Evaluation for a single model in nuScenes protocol.

        Args:
            result_path (str): Path of the result file.
            logger (logging.Logger | str | None): Logger used for printing
                related information during evaluation. Default: None.
            metric (str): Metric name used for evaluation. Default: 'bbox'.
            result_name (str): Result name in the metric prefix.
                Default: 'pts_bbox'.

        Returns:
            dict: Dictionary of evaluation details.
        """
zhangwenwei's avatar
zhangwenwei committed
377
378
379
380
381
382
383
        from nuscenes import NuScenes
        from nuscenes.eval.detection.evaluate import NuScenesEval

        output_dir = osp.join(*osp.split(result_path)[:-1])
        nusc = NuScenes(
            version=self.version, dataroot=self.data_root, verbose=False)
        eval_set_map = {
384
            'v1.0-mini': 'mini_val',
zhangwenwei's avatar
zhangwenwei committed
385
386
387
388
389
390
391
392
393
394
395
396
397
398
            'v1.0-trainval': 'val',
        }
        nusc_eval = NuScenesEval(
            nusc,
            config=self.eval_detection_configs,
            result_path=result_path,
            eval_set=eval_set_map[self.version],
            output_dir=output_dir,
            verbose=False)
        nusc_eval.main(render_curves=False)

        # record metrics
        metrics = mmcv.load(osp.join(output_dir, 'metrics_summary.json'))
        detail = dict()
wangtai's avatar
wangtai committed
399
        metric_prefix = f'{result_name}_NuScenes'
zhangwenwei's avatar
zhangwenwei committed
400
        for name in self.CLASSES:
zhangwenwei's avatar
zhangwenwei committed
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
            for k, v in metrics['label_aps'][name].items():
                val = float('{:.4f}'.format(v))
                detail['{}/{}_AP_dist_{}'.format(metric_prefix, name, k)] = val
            for k, v in metrics['label_tp_errors'][name].items():
                val = float('{:.4f}'.format(v))
                detail['{}/{}_{}'.format(metric_prefix, name, k)] = val

        detail['{}/NDS'.format(metric_prefix)] = metrics['nd_score']
        detail['{}/mAP'.format(metric_prefix)] = metrics['mean_ap']
        return detail

    def format_results(self, results, jsonfile_prefix=None):
        """Format the results to json (standard format for COCO evaluation).

        Args:
wangtai's avatar
wangtai committed
416
            results (list[dict]): Testing results of the dataset.
zhangwenwei's avatar
zhangwenwei committed
417
418
419
420
421
            jsonfile_prefix (str | None): The prefix of json files. It includes
                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:
zhangwenwei's avatar
zhangwenwei committed
422
423
424
425
            tuple: Returns (result_files, tmp_dir), where `result_files` is a \
                dict containing the json filepaths, `tmp_dir` is the temporal \
                directory created for saving json files when \
                `jsonfile_prefix` is not specified.
zhangwenwei's avatar
zhangwenwei committed
426
427
428
429
430
431
432
433
434
435
436
437
        """
        assert isinstance(results, list), 'results must be a list'
        assert len(results) == len(self), (
            'The length of results is not equal to the dataset len: {} != {}'.
            format(len(results), len(self)))

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

438
439
440
441
442
443
444
        # currently the output prediction results could be in two formats
        # 1. list of dict('boxes_3d': ..., 'scores_3d': ..., 'labels_3d': ...)
        # 2. list of dict('pts_bbox' or 'img_bbox':
        #     dict('boxes_3d': ..., 'scores_3d': ..., 'labels_3d': ...))
        # this is a workaround to enable evaluation of both formats on nuScenes
        # refer to https://github.com/open-mmlab/mmdetection3d/issues/449
        if not ('pts_bbox' in results[0] or 'img_bbox' in results[0]):
zhangwenwei's avatar
zhangwenwei committed
445
446
            result_files = self._format_bbox(results, jsonfile_prefix)
        else:
447
            # should take the inner dict out of 'pts_bbox' or 'img_bbox' dict
zhangwenwei's avatar
zhangwenwei committed
448
449
            result_files = dict()
            for name in results[0]:
zhangwenwei's avatar
zhangwenwei committed
450
                print(f'\nFormating bboxes of {name}')
zhangwenwei's avatar
zhangwenwei committed
451
452
453
454
455
456
457
458
459
460
461
                results_ = [out[name] for out in results]
                tmp_file_ = osp.join(jsonfile_prefix, name)
                result_files.update(
                    {name: self._format_bbox(results_, tmp_file_)})
        return result_files, tmp_dir

    def evaluate(self,
                 results,
                 metric='bbox',
                 logger=None,
                 jsonfile_prefix=None,
liyinhao's avatar
liyinhao committed
462
463
                 result_names=['pts_bbox'],
                 show=False,
464
465
                 out_dir=None,
                 pipeline=None):
zhangwenwei's avatar
zhangwenwei committed
466
467
468
        """Evaluation in nuScenes protocol.

        Args:
wangtai's avatar
wangtai committed
469
            results (list[dict]): Testing results of the dataset.
zhangwenwei's avatar
zhangwenwei committed
470
471
472
473
474
475
            metric (str | list[str]): Metrics to be evaluated.
            logger (logging.Logger | str | None): Logger used for printing
                related information during evaluation. Default: None.
            jsonfile_prefix (str | None): The prefix of json files. It includes
                the file path and the prefix of filename, e.g., "a/b/prefix".
                If not specified, a temp file will be created. Default: None.
liyinhao's avatar
liyinhao committed
476
477
478
479
            show (bool): Whether to visualize.
                Default: False.
            out_dir (str): Path to save the visualization results.
                Default: None.
480
481
            pipeline (list[dict], optional): raw data loading for showing.
                Default: None.
zhangwenwei's avatar
zhangwenwei committed
482
483

        Returns:
wangtai's avatar
wangtai committed
484
            dict[str, float]: Results of each evaluation metric.
zhangwenwei's avatar
zhangwenwei committed
485
486
487
488
489
490
491
492
493
494
495
496
497
498
        """
        result_files, tmp_dir = self.format_results(results, jsonfile_prefix)

        if isinstance(result_files, dict):
            results_dict = dict()
            for name in result_names:
                print('Evaluating bboxes of {}'.format(name))
                ret_dict = self._evaluate_single(result_files[name])
            results_dict.update(ret_dict)
        elif isinstance(result_files, str):
            results_dict = self._evaluate_single(result_files)

        if tmp_dir is not None:
            tmp_dir.cleanup()
liyinhao's avatar
liyinhao committed
499
500

        if show:
501
            self.show(results, out_dir, pipeline=pipeline)
zhangwenwei's avatar
zhangwenwei committed
502
503
        return results_dict

504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
    def _build_default_pipeline(self):
        """Build the default pipeline for this dataset."""
        pipeline = [
            dict(
                type='LoadPointsFromFile',
                coord_type='LIDAR',
                load_dim=5,
                use_dim=5,
                file_client_args=dict(backend='disk')),
            dict(
                type='LoadPointsFromMultiSweeps',
                sweeps_num=10,
                file_client_args=dict(backend='disk')),
            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):
526
527
528
529
530
        """Results visualization.

        Args:
            results (list[dict]): List of bounding boxes results.
            out_dir (str): Output directory of visualization result.
531
532
533
            show (bool): Visualize the results online.
            pipeline (list[dict], optional): raw data loading for showing.
                Default: None.
534
        """
535
536
        assert out_dir is not None, 'Expect out_dir, got none.'
        pipeline = self._get_pipeline(pipeline)
liyinhao's avatar
liyinhao committed
537
        for i, result in enumerate(results):
538
539
            if 'pts_bbox' in result.keys():
                result = result['pts_bbox']
liyinhao's avatar
liyinhao committed
540
541
542
            data_info = self.data_infos[i]
            pts_path = data_info['lidar_path']
            file_name = osp.split(pts_path)[-1].split('.')[0]
543
            points = self._extract_data(i, pipeline, 'points').numpy()
liyinhao's avatar
liyinhao committed
544
            # for now we convert points into depth mode
545
546
            points = Coord3DMode.convert_point(points, Coord3DMode.LIDAR,
                                               Coord3DMode.DEPTH)
547
            inds = result['scores_3d'] > 0.1
548
            gt_bboxes = self.get_ann_info(i)['gt_bboxes_3d'].tensor.numpy()
549
550
551
552
553
554
555
            show_gt_bboxes = Box3DMode.convert(gt_bboxes, Box3DMode.LIDAR,
                                               Box3DMode.DEPTH)
            pred_bboxes = result['boxes_3d'][inds].tensor.numpy()
            show_pred_bboxes = Box3DMode.convert(pred_bboxes, Box3DMode.LIDAR,
                                                 Box3DMode.DEPTH)
            show_result(points, show_gt_bboxes, show_pred_bboxes, out_dir,
                        file_name, show)
liyinhao's avatar
liyinhao committed
556

zhangwenwei's avatar
zhangwenwei committed
557
558

def output_to_nusc_box(detection):
559
560
561
562
563
    """Convert the output to the box class in the nuScenes.

    Args:
        detection (dict): Detection results.

wangtai's avatar
wangtai committed
564
565
566
            - boxes_3d (:obj:`BaseInstance3DBoxes`): Detection bbox.
            - scores_3d (torch.Tensor): Detection scores.
            - labels_3d (torch.Tensor): Predicted box labels.
567
568

    Returns:
zhangwenwei's avatar
zhangwenwei committed
569
        list[:obj:`NuScenesBox`]: List of standard NuScenesBoxes.
570
    """
571
    box3d = detection['boxes_3d']
zhangwenwei's avatar
zhangwenwei committed
572
573
    scores = detection['scores_3d'].numpy()
    labels = detection['labels_3d'].numpy()
574
575
576
577

    box_gravity_center = box3d.gravity_center.numpy()
    box_dims = box3d.dims.numpy()
    box_yaw = box3d.yaw.numpy()
zhangwenwei's avatar
zhangwenwei committed
578
579
    # TODO: check whether this is necessary
    # with dir_offset & dir_limit in the head
580
581
    box_yaw = -box_yaw - np.pi / 2

zhangwenwei's avatar
zhangwenwei committed
582
    box_list = []
583
584
585
    for i in range(len(box3d)):
        quat = pyquaternion.Quaternion(axis=[0, 0, 1], radians=box_yaw[i])
        velocity = (*box3d.tensor[i, 7:9], 0.0)
zhangwenwei's avatar
zhangwenwei committed
586
587
588
589
590
        # velo_val = np.linalg.norm(box3d[i, 7:9])
        # velo_ori = box3d[i, 6]
        # velocity = (
        # velo_val * np.cos(velo_ori), velo_val * np.sin(velo_ori), 0.0)
        box = NuScenesBox(
591
592
            box_gravity_center[i],
            box_dims[i],
zhangwenwei's avatar
zhangwenwei committed
593
594
595
596
597
598
599
600
601
602
603
604
605
            quat,
            label=labels[i],
            score=scores[i],
            velocity=velocity)
        box_list.append(box)
    return box_list


def lidar_nusc_box_to_global(info,
                             boxes,
                             classes,
                             eval_configs,
                             eval_version='detection_cvpr_2019'):
606
607
608
609
610
    """Convert the box from ego to global coordinate.

    Args:
        info (dict): Info for a specific sample data, including the
            calibration information.
zhangwenwei's avatar
zhangwenwei committed
611
        boxes (list[:obj:`NuScenesBox`]): List of predicted NuScenesBoxes.
612
613
614
615
616
617
618
619
620
        classes (list[str]): Mapped classes in the evaluation.
        eval_configs (object): Evaluation configuration object.
        eval_version (str): Evaluation version.
            Default: 'detection_cvpr_2019'

    Returns:
        list: List of standard NuScenesBoxes in the global
            coordinate.
    """
zhangwenwei's avatar
zhangwenwei committed
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
    box_list = []
    for box in boxes:
        # Move box to ego vehicle coord system
        box.rotate(pyquaternion.Quaternion(info['lidar2ego_rotation']))
        box.translate(np.array(info['lidar2ego_translation']))
        # filter det in ego.
        cls_range_map = eval_configs.class_range
        radius = np.linalg.norm(box.center[:2], 2)
        det_range = cls_range_map[classes[box.label]]
        if radius > det_range:
            continue
        # Move box to global coord system
        box.rotate(pyquaternion.Quaternion(info['ego2global_rotation']))
        box.translate(np.array(info['ego2global_translation']))
        box_list.append(box)
    return box_list