nuscenes_dataset.py 21.8 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):
zhangwenwei's avatar
zhangwenwei committed
16
    """NuScenes Dataset.
wangtai's avatar
wangtai committed
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51

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

    Please refer to `<https://www.nuscenes.org/download>`_for data
    downloading. It is recommended to symlink the dataset root to
    $MMDETECTION3D/data and organize them as the doc shows.

    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`.
            Defaults to 'LiDAR' in this dataset. Available options includes

            - 'LiDAR': box in LiDAR coordinates
            - 'Depth': box in depth coordinates, usually for indoor dataset
            - 'Camera': box in camera coordinates
        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'.
    """
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,
zhangwenwei's avatar
zhangwenwei committed
115
                 eval_version='detection_cvpr_2019'):
zhangwenwei's avatar
zhangwenwei committed
116
        self.load_interval = load_interval
zhangwenwei's avatar
zhangwenwei committed
117
118
119
120
121
122
        super().__init__(
            data_root=data_root,
            ann_file=ann_file,
            pipeline=pipeline,
            classes=classes,
            modality=modality,
123
124
            box_type_3d=box_type_3d,
            filter_empty_gt=filter_empty_gt,
zhangwenwei's avatar
zhangwenwei committed
125
            test_mode=test_mode)
zhangwenwei's avatar
zhangwenwei committed
126
127
128
129
130
131

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

zhangwenwei's avatar
zhangwenwei committed
141
    def load_annotations(self, ann_file):
142
143
144
145
146
147
148
149
        """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
150
151
152
153
154
155
        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
156

zhangwenwei's avatar
zhangwenwei committed
157
    def get_data_info(self, index):
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
        """Get data info according to the given index.

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

        Returns:
            dict: Standard input_dict consists of the
                data information.

                - 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
        """
zhangwenwei's avatar
zhangwenwei committed
176
        info = self.data_infos[index]
zhangwenwei's avatar
zhangwenwei committed
177

zhangwenwei's avatar
zhangwenwei committed
178
        # standard protocal modified from SECOND.Pytorch
zhangwenwei's avatar
zhangwenwei committed
179
180
        input_dict = dict(
            sample_idx=info['token'],
zhangwenwei's avatar
zhangwenwei committed
181
182
183
            pts_filename=info['lidar_path'],
            sweeps=info['sweeps'],
            timestamp=info['timestamp'] / 1e6,
zhangwenwei's avatar
zhangwenwei committed
184
185
186
187
188
189
        )

        if self.modality['use_camera']:
            image_paths = []
            lidar2img_rts = []
            for cam_type, cam_info in info['cams'].items():
zhangwenwei's avatar
zhangwenwei committed
190
                image_paths.append(cam_info['data_path'])
zhangwenwei's avatar
zhangwenwei committed
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
                # 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
206
                    img_filename=image_paths,
zhangwenwei's avatar
zhangwenwei committed
207
208
209
                    lidar2img=lidar2img_rts,
                ))

zhangwenwei's avatar
zhangwenwei committed
210
        if not self.test_mode:
zhangwenwei's avatar
zhangwenwei committed
211
            annos = self.get_ann_info(index)
zhangwenwei's avatar
zhangwenwei committed
212
            input_dict['ann_info'] = annos
zhangwenwei's avatar
zhangwenwei committed
213
214
215
216

        return input_dict

    def get_ann_info(self, index):
217
218
219
220
221
222
223
224
225
226
227
228
229
230
        """Get annotation info according to the given index.

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

        Returns:
            dict: Standard annotation dictionary
                consists of the data information.

                - gt_bboxes_3d (:obj:``LiDARInstance3DBoxes``):
                    3D ground truth bboxes
                - gt_labels_3d (np.ndarray): labels of ground truths
                - gt_names (list[str]): class names of ground truths
        """
zhangwenwei's avatar
zhangwenwei committed
231
        info = self.data_infos[index]
zhangwenwei's avatar
zhangwenwei committed
232
233
234
235
        # filter out bbox containing no points
        mask = info['num_lidar_pts'] > 0
        gt_bboxes_3d = info['gt_boxes'][mask]
        gt_names_3d = info['gt_names'][mask]
zhangwenwei's avatar
zhangwenwei committed
236
237
238
239
240
241
242
        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
243
244
245
246
247
248
249

        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
250
        # the nuscenes box center is [0.5, 0.5, 0.5], we change it to be
wuyuefeng's avatar
wuyuefeng committed
251
        # the same as KITTI (0.5, 0.5, 0)
zhangwenwei's avatar
zhangwenwei committed
252
253
254
        gt_bboxes_3d = LiDARInstance3DBoxes(
            gt_bboxes_3d,
            box_dim=gt_bboxes_3d.shape[-1],
wuyuefeng's avatar
wuyuefeng committed
255
            origin=(0.5, 0.5, 0.5)).convert_to(self.box_mode_3d)
zhangwenwei's avatar
zhangwenwei committed
256

zhangwenwei's avatar
zhangwenwei committed
257
258
        anns_results = dict(
            gt_bboxes_3d=gt_bboxes_3d,
zhangwenwei's avatar
zhangwenwei committed
259
            gt_labels_3d=gt_labels_3d,
liyinhao's avatar
liyinhao committed
260
            gt_names=gt_names_3d)
zhangwenwei's avatar
zhangwenwei committed
261
262
263
        return anns_results

    def _format_bbox(self, results, jsonfile_prefix=None):
264
265
266
267
268
269
270
271
272
273
274
        """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
275
        nusc_annos = {}
zhangwenwei's avatar
zhangwenwei committed
276
        mapped_class_names = self.CLASSES
zhangwenwei's avatar
zhangwenwei committed
277

zhangwenwei's avatar
zhangwenwei committed
278
        print('Start to convert detection format...')
zhangwenwei's avatar
zhangwenwei committed
279
        for sample_id, det in enumerate(mmcv.track_iter_progress(results)):
zhangwenwei's avatar
zhangwenwei committed
280
            annos = []
zhangwenwei's avatar
zhangwenwei committed
281
282
283
284
            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
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
                                             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
311
                    sample_token=sample_token,
zhangwenwei's avatar
zhangwenwei committed
312
313
314
315
316
317
318
319
                    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
320
            nusc_annos[sample_token] = annos
zhangwenwei's avatar
zhangwenwei committed
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
        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'):
337
338
339
340
341
342
343
344
345
346
347
348
349
        """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
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
        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 = {
            'v1.0-mini': 'mini_train',
            '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
372
        metric_prefix = f'{result_name}_NuScenes'
zhangwenwei's avatar
zhangwenwei committed
373
        for name in self.CLASSES:
zhangwenwei's avatar
zhangwenwei committed
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
            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
389
            results (list[dict]): Testing results of the dataset.
zhangwenwei's avatar
zhangwenwei committed
390
391
392
393
394
            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:
wangtai's avatar
wangtai committed
395
396
397
            tuple (dict, str): 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
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
        """
        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
415
                print(f'\nFormating bboxes of {name}')
zhangwenwei's avatar
zhangwenwei committed
416
417
418
419
420
421
422
423
424
425
426
                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
427
428
429
                 result_names=['pts_bbox'],
                 show=False,
                 out_dir=None):
zhangwenwei's avatar
zhangwenwei committed
430
431
432
        """Evaluation in nuScenes protocol.

        Args:
wangtai's avatar
wangtai committed
433
            results (list[dict]): Testing results of the dataset.
zhangwenwei's avatar
zhangwenwei committed
434
435
436
437
438
439
            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
440
441
442
443
            show (bool): Whether to visualize.
                Default: False.
            out_dir (str): Path to save the visualization results.
                Default: None.
zhangwenwei's avatar
zhangwenwei committed
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460

        Returns:
            dict[str: float]
        """
        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
461
462
463

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

liyinhao's avatar
liyinhao committed
466
    def show(self, results, out_dir):
467
468
469
470
471
472
        """Results visualization.

        Args:
            results (list[dict]): List of bounding boxes results.
            out_dir (str): Output directory of visualization result.
        """
liyinhao's avatar
liyinhao committed
473
        for i, result in enumerate(results):
liyinhao's avatar
liyinhao committed
474
475
            example = self.prepare_test_data(i)
            points = example['points'][0]._data.numpy()
liyinhao's avatar
liyinhao committed
476
477
478
            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
479
            # for now we convert points into depth mode
liyinhao's avatar
liyinhao committed
480
481
            points = points[..., [1, 0, 2]]
            points[..., 0] *= -1
liyinhao's avatar
liyinhao committed
482
            inds = result['pts_bbox']['scores_3d'] > 0.1
liyinhao's avatar
liyinhao committed
483
484
485
486
            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
487
            pred_bboxes = result['pts_bbox']['boxes_3d'][inds].tensor.numpy()
liyinhao's avatar
liyinhao committed
488
489
490
491
492
            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
493
494

def output_to_nusc_box(detection):
495
496
497
498
499
500
501
502
503
504
505
506
    """Convert the output to the box class in the nuScenes.

    Args:
        detection (dict): Detection results.

            - boxes_3d (:obj:``BaseInstance3DBoxes``): detection bbox
            - scores_3d (torch.Tensor): detection scores
            - labels_3d (torch.Tensor): predicted box labels

    Returns:
        list[:obj:``NuScenesBox``]: List of standard NuScenesBoxes.
    """
507
    box3d = detection['boxes_3d']
zhangwenwei's avatar
zhangwenwei committed
508
509
    scores = detection['scores_3d'].numpy()
    labels = detection['labels_3d'].numpy()
510
511
512
513

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

zhangwenwei's avatar
zhangwenwei committed
518
    box_list = []
519
520
521
    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
522
523
524
525
526
        # 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(
527
528
            box_gravity_center[i],
            box_dims[i],
zhangwenwei's avatar
zhangwenwei committed
529
530
531
532
533
534
535
536
537
538
539
540
541
            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'):
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
    """Convert the box from ego to global coordinate.

    Args:
        info (dict): Info for a specific sample data, including the
            calibration information.
        boxes (list[:obj:``NuScenesBox``]): List of predicted NuScenesBoxes.
        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
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
    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