nuscenes_dataset.py 22.9 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
10
from ..core import show_result
from ..core.bbox import Box3DMode, LiDARInstance3DBoxes
zhangwenwei's avatar
zhangwenwei committed
11
from .custom_3d import Custom3DDataset
zhangwenwei's avatar
zhangwenwei committed
12
13


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

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

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

    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
39
            Defaults to 'LiDAR' in this dataset. Available options includes.
wangtai's avatar
wangtai committed
40
41
42
            - 'LiDAR': Box in LiDAR coordinates.
            - 'Depth': Box in depth coordinates, usually for indoor dataset.
            - 'Camera': Box in camera coordinates.
wangtai's avatar
wangtai committed
43
44
45
46
47
48
        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
49
50
        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
51
    """
zhangwenwei's avatar
zhangwenwei committed
52
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
    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
107
108
                 data_root=None,
                 classes=None,
zhangwenwei's avatar
zhangwenwei committed
109
110
111
                 load_interval=1,
                 with_velocity=True,
                 modality=None,
112
113
114
                 box_type_3d='LiDAR',
                 filter_empty_gt=True,
                 test_mode=False,
yinchimaoliang's avatar
yinchimaoliang committed
115
116
                 eval_version='detection_cvpr_2019',
                 use_valid_flag=False):
zhangwenwei's avatar
zhangwenwei committed
117
        self.load_interval = load_interval
yinchimaoliang's avatar
yinchimaoliang committed
118
        self.use_valid_flag = use_valid_flag
zhangwenwei's avatar
zhangwenwei committed
119
120
121
122
123
124
        super().__init__(
            data_root=data_root,
            ann_file=ann_file,
            pipeline=pipeline,
            classes=classes,
            modality=modality,
125
126
            box_type_3d=box_type_3d,
            filter_empty_gt=filter_empty_gt,
zhangwenwei's avatar
zhangwenwei committed
127
            test_mode=test_mode)
zhangwenwei's avatar
zhangwenwei committed
128
129
130
131
132

        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
133
134
        if self.modality is None:
            self.modality = dict(
zhangwenwei's avatar
zhangwenwei committed
135
136
137
138
139
140
141
                use_camera=False,
                use_lidar=True,
                use_radar=False,
                use_map=False,
                use_external=False,
            )

yinchimaoliang's avatar
yinchimaoliang committed
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
    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'])
159
160

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

zhangwenwei's avatar
zhangwenwei committed
166
    def load_annotations(self, ann_file):
167
168
169
170
171
172
173
174
        """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
175
176
177
178
179
180
        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
181

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

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

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

wangtai's avatar
wangtai committed
192
193
194
195
196
197
198
199
                - 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.
200
        """
zhangwenwei's avatar
zhangwenwei committed
201
        info = self.data_infos[index]
zhangwenwei's avatar
zhangwenwei committed
202
        # standard protocal modified from SECOND.Pytorch
zhangwenwei's avatar
zhangwenwei committed
203
204
        input_dict = dict(
            sample_idx=info['token'],
zhangwenwei's avatar
zhangwenwei committed
205
206
207
            pts_filename=info['lidar_path'],
            sweeps=info['sweeps'],
            timestamp=info['timestamp'] / 1e6,
zhangwenwei's avatar
zhangwenwei committed
208
209
210
211
212
213
        )

        if self.modality['use_camera']:
            image_paths = []
            lidar2img_rts = []
            for cam_type, cam_info in info['cams'].items():
zhangwenwei's avatar
zhangwenwei committed
214
                image_paths.append(cam_info['data_path'])
zhangwenwei's avatar
zhangwenwei committed
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
                # 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
230
                    img_filename=image_paths,
zhangwenwei's avatar
zhangwenwei committed
231
232
233
                    lidar2img=lidar2img_rts,
                ))

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

        return input_dict

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

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

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

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

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

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

    def _format_bbox(self, results, jsonfile_prefix=None):
290
291
292
293
294
295
296
297
298
299
300
        """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
301
        nusc_annos = {}
zhangwenwei's avatar
zhangwenwei committed
302
        mapped_class_names = self.CLASSES
zhangwenwei's avatar
zhangwenwei committed
303

zhangwenwei's avatar
zhangwenwei committed
304
        print('Start to convert detection format...')
zhangwenwei's avatar
zhangwenwei committed
305
        for sample_id, det in enumerate(mmcv.track_iter_progress(results)):
zhangwenwei's avatar
zhangwenwei committed
306
            annos = []
zhangwenwei's avatar
zhangwenwei committed
307
308
309
310
            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
311
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
                                             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
337
                    sample_token=sample_token,
zhangwenwei's avatar
zhangwenwei committed
338
339
340
341
342
343
344
345
                    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
346
            nusc_annos[sample_token] = annos
zhangwenwei's avatar
zhangwenwei committed
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
        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'):
363
364
365
366
367
368
369
370
371
372
373
374
375
        """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
376
377
378
379
380
381
382
        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 = {
383
            'v1.0-mini': 'mini_val',
zhangwenwei's avatar
zhangwenwei committed
384
385
386
387
388
389
390
391
392
393
394
395
396
397
            '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
398
        metric_prefix = f'{result_name}_NuScenes'
zhangwenwei's avatar
zhangwenwei committed
399
        for name in self.CLASSES:
zhangwenwei's avatar
zhangwenwei committed
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
            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
415
            results (list[dict]): Testing results of the dataset.
zhangwenwei's avatar
zhangwenwei committed
416
417
418
419
420
            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
421
422
423
424
            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
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
        """
        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

        if not isinstance(results[0], dict):
            result_files = self._format_bbox(results, jsonfile_prefix)
        else:
            result_files = dict()
            for name in results[0]:
zhangwenwei's avatar
zhangwenwei committed
442
                print(f'\nFormating bboxes of {name}')
zhangwenwei's avatar
zhangwenwei committed
443
444
445
446
447
448
449
450
451
452
453
                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
454
455
456
                 result_names=['pts_bbox'],
                 show=False,
                 out_dir=None):
zhangwenwei's avatar
zhangwenwei committed
457
458
459
        """Evaluation in nuScenes protocol.

        Args:
wangtai's avatar
wangtai committed
460
            results (list[dict]): Testing results of the dataset.
zhangwenwei's avatar
zhangwenwei committed
461
462
463
464
465
466
            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
467
468
469
470
            show (bool): Whether to visualize.
                Default: False.
            out_dir (str): Path to save the visualization results.
                Default: None.
zhangwenwei's avatar
zhangwenwei committed
471
472

        Returns:
wangtai's avatar
wangtai committed
473
            dict[str, float]: Results of each evaluation metric.
zhangwenwei's avatar
zhangwenwei committed
474
475
476
477
478
479
480
481
482
483
484
485
486
487
        """
        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
488
489
490

        if show:
            self.show(results, out_dir)
zhangwenwei's avatar
zhangwenwei committed
491
492
        return results_dict

liyinhao's avatar
liyinhao committed
493
    def show(self, results, out_dir):
494
495
496
497
498
499
        """Results visualization.

        Args:
            results (list[dict]): List of bounding boxes results.
            out_dir (str): Output directory of visualization result.
        """
liyinhao's avatar
liyinhao committed
500
        for i, result in enumerate(results):
liyinhao's avatar
liyinhao committed
501
502
            example = self.prepare_test_data(i)
            points = example['points'][0]._data.numpy()
liyinhao's avatar
liyinhao committed
503
504
505
            data_info = self.data_infos[i]
            pts_path = data_info['lidar_path']
            file_name = osp.split(pts_path)[-1].split('.')[0]
liyinhao's avatar
liyinhao committed
506
            # for now we convert points into depth mode
liyinhao's avatar
liyinhao committed
507
508
            points = points[..., [1, 0, 2]]
            points[..., 0] *= -1
liyinhao's avatar
liyinhao committed
509
            inds = result['pts_bbox']['scores_3d'] > 0.1
liyinhao's avatar
liyinhao committed
510
511
512
513
            gt_bboxes = self.get_ann_info(i)['gt_bboxes_3d'].tensor
            gt_bboxes = Box3DMode.convert(gt_bboxes, Box3DMode.LIDAR,
                                          Box3DMode.DEPTH)
            gt_bboxes[..., 2] += gt_bboxes[..., 5] / 2
liyinhao's avatar
liyinhao committed
514
            pred_bboxes = result['pts_bbox']['boxes_3d'][inds].tensor.numpy()
liyinhao's avatar
liyinhao committed
515
516
517
518
519
            pred_bboxes = Box3DMode.convert(pred_bboxes, Box3DMode.LIDAR,
                                            Box3DMode.DEPTH)
            pred_bboxes[..., 2] += pred_bboxes[..., 5] / 2
            show_result(points, gt_bboxes, pred_bboxes, out_dir, file_name)

zhangwenwei's avatar
zhangwenwei committed
520
521

def output_to_nusc_box(detection):
522
523
524
525
526
    """Convert the output to the box class in the nuScenes.

    Args:
        detection (dict): Detection results.

wangtai's avatar
wangtai committed
527
528
529
            - boxes_3d (:obj:`BaseInstance3DBoxes`): Detection bbox.
            - scores_3d (torch.Tensor): Detection scores.
            - labels_3d (torch.Tensor): Predicted box labels.
530
531

    Returns:
zhangwenwei's avatar
zhangwenwei committed
532
        list[:obj:`NuScenesBox`]: List of standard NuScenesBoxes.
533
    """
534
    box3d = detection['boxes_3d']
zhangwenwei's avatar
zhangwenwei committed
535
536
    scores = detection['scores_3d'].numpy()
    labels = detection['labels_3d'].numpy()
537
538
539
540

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

zhangwenwei's avatar
zhangwenwei committed
545
    box_list = []
546
547
548
    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
549
550
551
552
553
        # 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(
554
555
            box_gravity_center[i],
            box_dims[i],
zhangwenwei's avatar
zhangwenwei committed
556
557
558
559
560
561
562
563
564
565
566
567
568
            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'):
569
570
571
572
573
    """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
574
        boxes (list[:obj:`NuScenesBox`]): List of predicted NuScenesBoxes.
575
576
577
578
579
580
581
582
583
        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
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
    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