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
159
160
161
162
163
164
    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.
        """
        class_sample_idx = {name: [] for name in self.CLASSES}
        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'])
        for name in gt_names:
            if name in self.CLASSES:
                class_sample_idx[name].append(idx)
        return class_sample_idx

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

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

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

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

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

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

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

        return input_dict

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

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

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

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

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

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

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

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

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

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

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

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

        Args:
            results (list[dict]): List of bounding boxes results.
            out_dir (str): Output directory of visualization result.
        """
liyinhao's avatar
liyinhao committed
499
        for i, result in enumerate(results):
liyinhao's avatar
liyinhao committed
500
501
            example = self.prepare_test_data(i)
            points = example['points'][0]._data.numpy()
liyinhao's avatar
liyinhao committed
502
503
504
            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
505
            # for now we convert points into depth mode
liyinhao's avatar
liyinhao committed
506
507
            points = points[..., [1, 0, 2]]
            points[..., 0] *= -1
liyinhao's avatar
liyinhao committed
508
            inds = result['pts_bbox']['scores_3d'] > 0.1
liyinhao's avatar
liyinhao committed
509
510
511
512
            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
513
            pred_bboxes = result['pts_bbox']['boxes_3d'][inds].tensor.numpy()
liyinhao's avatar
liyinhao committed
514
515
516
517
518
            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
519
520

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

    Args:
        detection (dict): Detection results.

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

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

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

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