waymo_metric.py 30.9 KB
Newer Older
1
2
3
# Copyright (c) OpenMMLab. All rights reserved.
import tempfile
from os import path as osp
4
from typing import Dict, List, Optional, Tuple, Union
5

6
import mmengine
7
8
import numpy as np
import torch
9
10
from mmengine import Config, load
from mmengine.logging import MMLogger, print_log
11
12
13

from mmdet3d.models.layers import box3d_multiclass_nms
from mmdet3d.registry import METRICS
14
15
16
from mmdet3d.structures import (Box3DMode, CameraInstance3DBoxes,
                                LiDARInstance3DBoxes, bbox3d2result,
                                points_cam2img, xywhr2xyxyr)
17
18
19
20
21
22
23
24
25
26
from .kitti_metric import KittiMetric


@METRICS.register_module()
class WaymoMetric(KittiMetric):
    """Waymo evaluation metric.

    Args:
        ann_file (str): The path of the annotation file in kitti format.
        waymo_bin_file (str): The path of the annotation file in waymo format.
27
28
        data_root (str): Path of dataset root. Used for storing waymo
            evaluation programs.
29
        split (str): The split of the evaluation set. Defaults to 'training'.
30
31
32
33
34
35
        metric (str or List[str]): Metrics to be evaluated. Defaults to 'mAP'.
        pcd_limit_range (List[float]): The range of point cloud used to filter
            invalid predicted boxes. Defaults to [-85, -85, -5, 85, 85, 5].
        convert_kitti_format (bool): Whether to convert the results to kitti
            format. Now, in order to be compatible with camera-based methods,
            defaults to True.
36
37
        prefix (str, optional): The prefix that will be added in the metric
            names to disambiguate homonymous metrics of different evaluators.
38
39
            If prefix is not provided in the argument, self.default_prefix will
            be used instead. Defaults to None.
40
        format_only (bool): Format the output results without perform
41
42
            evaluation. It is useful when you want to format the result to a
            specific format and submit it to the test server.
43
            Defaults to False.
44
45
46
47
48
        pklfile_prefix (str, optional): The prefix of pkl files, including the
            file path and the prefix of filename, e.g., "a/b/prefix". If not
            specified, a temp file will be created. Defaults to None.
        submission_prefix (str, optional): The prefix of submission data. If
            not specified, the submission data will not be generated.
49
50
            Defaults to None.
        load_type (str): Type of loading mode during training.
51
52
53

            - 'frame_based': Load all of the instances in the frame.
            - 'mv_image_based': Load all of the instances in the frame and need
54
55
              to convert to the FOV-based data type to support image-based
              detector.
56
57
58
            - 'fov_image_based': Only load the instances inside the default cam
              and need to convert to the FOV-based data type to support image-
              based detector.
59
60
61
        default_cam_key (str): The default camera for lidar to camera
            conversion. By default, KITTI: 'CAM2', Waymo: 'CAM_FRONT'.
            Defaults to 'CAM_FRONT'.
62
63
64
65
        use_pred_sample_idx (bool): In formating results, use the sample index
            from the prediction or from the load annotations. By default,
            KITTI: True, Waymo: False, Waymo has a conversion process, which
            needs to use the sample idx from load annotation.
66
            Defaults to False.
67
68
        collect_device (str): Device name used for collecting results from
            different ranks during distributed training. Must be 'cpu' or
69
            'gpu'. Defaults to 'cpu'.
70
71
        backend_args (dict, optional): Arguments to instantiate the
            corresponding backend. Defaults to None.
72
73
74
75
        idx2metainfo (str, optional): The file path of the metainfo in waymo.
            It stores the mapping from sample_idx to metainfo. The metainfo
            must contain the keys: 'idx2contextname' and 'idx2timestamp'.
            Defaults to None.
76
    """
77
    num_cams = 5
78
79
80
81
82
83

    def __init__(self,
                 ann_file: str,
                 waymo_bin_file: str,
                 data_root: str,
                 split: str = 'training',
84
                 metric: Union[str, List[str]] = 'mAP',
85
                 pcd_limit_range: List[float] = [-85, -85, -5, 85, 85, 5],
86
                 convert_kitti_format: bool = True,
87
                 prefix: Optional[str] = None,
88
89
90
                 format_only: bool = False,
                 pklfile_prefix: Optional[str] = None,
                 submission_prefix: Optional[str] = None,
91
                 load_type: str = 'frame_based',
92
93
94
                 default_cam_key: str = 'CAM_FRONT',
                 use_pred_sample_idx: bool = False,
                 collect_device: str = 'cpu',
95
                 backend_args: Optional[dict] = None,
96
                 idx2metainfo: Optional[str] = None) -> None:
97
98
99
        self.waymo_bin_file = waymo_bin_file
        self.data_root = data_root
        self.split = split
100
        self.load_type = load_type
101
        self.use_pred_sample_idx = use_pred_sample_idx
102
103
104
105
106
107
108
        self.convert_kitti_format = convert_kitti_format

        if idx2metainfo is not None:
            self.idx2metainfo = mmengine.load(idx2metainfo)
        else:
            self.idx2metainfo = None

109
        super(WaymoMetric, self).__init__(
110
111
112
113
114
115
116
            ann_file=ann_file,
            metric=metric,
            pcd_limit_range=pcd_limit_range,
            prefix=prefix,
            pklfile_prefix=pklfile_prefix,
            submission_prefix=submission_prefix,
            default_cam_key=default_cam_key,
117
            collect_device=collect_device,
118
            backend_args=backend_args)
119
120
        self.format_only = format_only
        if self.format_only:
121
122
123
            assert pklfile_prefix is not None, 'pklfile_prefix must be not '
            'None when format_only is True, otherwise the result files will '
            'be saved to a temp directory which will be cleaned up at the end.'
124

125
126
        self.default_prefix = 'Waymo metric'

127
    def compute_metrics(self, results: List[dict]) -> Dict[str, float]:
128
129
130
        """Compute the metrics from processed results.

        Args:
131
            results (List[dict]): The processed results of the whole dataset.
132
133
134
135
136
137

        Returns:
            Dict[str, float]: The computed metrics. The keys are the names of
            the metrics, and the values are corresponding results.
        """
        logger: MMLogger = MMLogger.get_current_instance()
138
        self.classes = self.dataset_meta['classes']
139
140

        # load annotations
141
        self.data_infos = load(self.ann_file)['data_list']
142
143
        assert len(results) == len(self.data_infos), \
            'invalid list length of network outputs'
144
        # different from kitti, waymo do not need to convert the ann file
145
146
        # handle the mv_image_based load_mode
        if self.load_type == 'mv_image_based':
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
            new_data_infos = []
            for info in self.data_infos:
                height = info['images'][self.default_cam_key]['height']
                width = info['images'][self.default_cam_key]['width']
                for (cam_key, img_info) in info['images'].items():
                    camera_info = dict()
                    camera_info['images'] = dict()
                    camera_info['images'][cam_key] = img_info
                    # TODO remove the check by updating the data info;
                    if 'height' not in img_info:
                        img_info['height'] = height
                        img_info['width'] = width
                    if 'cam_instances' in info \
                            and cam_key in info['cam_instances']:
                        camera_info['instances'] = info['cam_instances'][
                            cam_key]
                    else:
                        camera_info['instances'] = []
                    camera_info['ego2global'] = info['ego2global']
                    if 'image_sweeps' in info:
                        camera_info['image_sweeps'] = info['image_sweeps']

169
                    # TODO check if need to modify the sample idx
170
                    # TODO check when will use it except for evaluation.
171
                    camera_info['sample_idx'] = info['sample_idx']
172
173
                    new_data_infos.append(camera_info)
            self.data_infos = new_data_infos
174
175
176
177
178
179
180
181
182
183
184
185
186
187

        if self.pklfile_prefix is None:
            eval_tmp_dir = tempfile.TemporaryDirectory()
            pklfile_prefix = osp.join(eval_tmp_dir.name, 'results')
        else:
            eval_tmp_dir = None
            pklfile_prefix = self.pklfile_prefix

        result_dict, tmp_dir = self.format_results(
            results,
            pklfile_prefix=pklfile_prefix,
            submission_prefix=self.submission_prefix,
            classes=self.classes)

188
        metric_dict = {}
189
190
191
192
193
194

        if self.format_only:
            logger.info('results are saved in '
                        f'{osp.dirname(self.pklfile_prefix)}')
            return metric_dict

195
196
197
        for metric in self.metrics:
            ap_dict = self.waymo_evaluate(
                pklfile_prefix, metric=metric, logger=logger)
198
            metric_dict.update(ap_dict)
199
200
201
202
203
        if eval_tmp_dir is not None:
            eval_tmp_dir.cleanup()

        if tmp_dir is not None:
            tmp_dir.cleanup()
204
        return metric_dict
205

206
207
    def waymo_evaluate(self,
                       pklfile_prefix: str,
208
209
                       metric: Optional[str] = None,
                       logger: Optional[MMLogger] = None) -> Dict[str, float]:
210
211
212
213
214
        """Evaluation in Waymo protocol.

        Args:
            pklfile_prefix (str): The location that stored the prediction
                results.
215
            metric (str, optional): Metric to be evaluated. Defaults to None.
216
217
            logger (MMLogger, optional): Logger used for printing related
                information during evaluation. Defaults to None.
218
219

        Returns:
220
            Dict[str, float]: Results of each evaluation metric.
221
222
223
224
225
226
227
228
229
        """

        import subprocess

        if metric == 'mAP':
            eval_str = 'mmdet3d/evaluation/functional/waymo_utils/' + \
                f'compute_detection_metrics_main {pklfile_prefix}.bin ' + \
                f'{self.waymo_bin_file}'
            print(eval_str)
230
            ret_bytes = subprocess.check_output(eval_str, shell=True)
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
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
311
312
313
314
315
316
317
318
319
320
321
322
            ret_texts = ret_bytes.decode('utf-8')
            print_log(ret_texts, logger=logger)

            ap_dict = {
                'Vehicle/L1 mAP': 0,
                'Vehicle/L1 mAPH': 0,
                'Vehicle/L2 mAP': 0,
                'Vehicle/L2 mAPH': 0,
                'Pedestrian/L1 mAP': 0,
                'Pedestrian/L1 mAPH': 0,
                'Pedestrian/L2 mAP': 0,
                'Pedestrian/L2 mAPH': 0,
                'Sign/L1 mAP': 0,
                'Sign/L1 mAPH': 0,
                'Sign/L2 mAP': 0,
                'Sign/L2 mAPH': 0,
                'Cyclist/L1 mAP': 0,
                'Cyclist/L1 mAPH': 0,
                'Cyclist/L2 mAP': 0,
                'Cyclist/L2 mAPH': 0,
                'Overall/L1 mAP': 0,
                'Overall/L1 mAPH': 0,
                'Overall/L2 mAP': 0,
                'Overall/L2 mAPH': 0
            }
            mAP_splits = ret_texts.split('mAP ')
            mAPH_splits = ret_texts.split('mAPH ')
            for idx, key in enumerate(ap_dict.keys()):
                split_idx = int(idx / 2) + 1
                if idx % 2 == 0:  # mAP
                    ap_dict[key] = float(mAP_splits[split_idx].split(']')[0])
                else:  # mAPH
                    ap_dict[key] = float(mAPH_splits[split_idx].split(']')[0])
            ap_dict['Overall/L1 mAP'] = \
                (ap_dict['Vehicle/L1 mAP'] + ap_dict['Pedestrian/L1 mAP'] +
                    ap_dict['Cyclist/L1 mAP']) / 3
            ap_dict['Overall/L1 mAPH'] = \
                (ap_dict['Vehicle/L1 mAPH'] + ap_dict['Pedestrian/L1 mAPH'] +
                    ap_dict['Cyclist/L1 mAPH']) / 3
            ap_dict['Overall/L2 mAP'] = \
                (ap_dict['Vehicle/L2 mAP'] + ap_dict['Pedestrian/L2 mAP'] +
                    ap_dict['Cyclist/L2 mAP']) / 3
            ap_dict['Overall/L2 mAPH'] = \
                (ap_dict['Vehicle/L2 mAPH'] + ap_dict['Pedestrian/L2 mAPH'] +
                    ap_dict['Cyclist/L2 mAPH']) / 3
        elif metric == 'LET_mAP':
            eval_str = 'mmdet3d/evaluation/functional/waymo_utils/' + \
                f'compute_detection_let_metrics_main {pklfile_prefix}.bin ' + \
                f'{self.waymo_bin_file}'

            print(eval_str)
            ret_bytes = subprocess.check_output(eval_str, shell=True)
            ret_texts = ret_bytes.decode('utf-8')

            print_log(ret_texts, logger=logger)
            ap_dict = {
                'Vehicle mAPL': 0,
                'Vehicle mAP': 0,
                'Vehicle mAPH': 0,
                'Pedestrian mAPL': 0,
                'Pedestrian mAP': 0,
                'Pedestrian mAPH': 0,
                'Sign mAPL': 0,
                'Sign mAP': 0,
                'Sign mAPH': 0,
                'Cyclist mAPL': 0,
                'Cyclist mAP': 0,
                'Cyclist mAPH': 0,
                'Overall mAPL': 0,
                'Overall mAP': 0,
                'Overall mAPH': 0
            }
            mAPL_splits = ret_texts.split('mAPL ')
            mAP_splits = ret_texts.split('mAP ')
            mAPH_splits = ret_texts.split('mAPH ')
            for idx, key in enumerate(ap_dict.keys()):
                split_idx = int(idx / 3) + 1
                if idx % 3 == 0:  # mAPL
                    ap_dict[key] = float(mAPL_splits[split_idx].split(']')[0])
                elif idx % 3 == 1:  # mAP
                    ap_dict[key] = float(mAP_splits[split_idx].split(']')[0])
                else:  # mAPH
                    ap_dict[key] = float(mAPH_splits[split_idx].split(']')[0])
            ap_dict['Overall mAPL'] = \
                (ap_dict['Vehicle mAPL'] + ap_dict['Pedestrian mAPL'] +
                    ap_dict['Cyclist mAPL']) / 3
            ap_dict['Overall mAP'] = \
                (ap_dict['Vehicle mAP'] + ap_dict['Pedestrian mAP'] +
                    ap_dict['Cyclist mAP']) / 3
            ap_dict['Overall mAPH'] = \
                (ap_dict['Vehicle mAPH'] + ap_dict['Pedestrian mAPH'] +
                    ap_dict['Cyclist mAPH']) / 3
323
324
        return ap_dict

325
326
327
328
329
330
331
    def format_results(
        self,
        results: List[dict],
        pklfile_prefix: Optional[str] = None,
        submission_prefix: Optional[str] = None,
        classes: Optional[List[str]] = None
    ) -> Tuple[dict, Union[tempfile.TemporaryDirectory, None]]:
332
        """Format the results to bin file.
333
334

        Args:
335
            results (List[dict]): Testing results of the dataset.
336
337
338
            pklfile_prefix (str, optional): The prefix of pkl 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.
339
                Defaults to None.
340
341
342
            submission_prefix (str, optional): The prefix of submitted 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.
343
344
345
                Defaults to None.
            classes (List[str], optional): A list of class name.
                Defaults to None.
346
347

        Returns:
348
349
350
            tuple: (result_dict, tmp_dir), result_dict is a dict containing the
            formatted result, tmp_dir is the temporal directory created for
            saving json files when jsonfile_prefix is not specified.
351
        """
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
        waymo_save_tmp_dir = tempfile.TemporaryDirectory()
        waymo_results_save_dir = waymo_save_tmp_dir.name
        waymo_results_final_path = f'{pklfile_prefix}.bin'

        if self.convert_kitti_format:
            results_kitti_format, tmp_dir = super().format_results(
                results, pklfile_prefix, submission_prefix, classes)
            final_results = results_kitti_format['pred_instances_3d']
        else:
            final_results = results
            for i, res in enumerate(final_results):
                # Actually, `sample_idx` here is the filename without suffix.
                # It's for identitying the sample in formating.
                res['sample_idx'] = self.data_infos[i]['sample_idx']
                res['pred_instances_3d']['bboxes_3d'].limit_yaw(
                    offset=0.5, period=np.pi * 2)
368
369
370
371
372
373
374
375
376
377

        waymo_root = self.data_root
        if self.split == 'training':
            waymo_tfrecords_dir = osp.join(waymo_root, 'validation')
            prefix = '1'
        elif self.split == 'testing':
            waymo_tfrecords_dir = osp.join(waymo_root, 'testing')
            prefix = '2'
        else:
            raise ValueError('Not supported split value.')
378
379
380
381
382

        from ..functional.waymo_utils.prediction_to_waymo import \
            Prediction2Waymo
        converter = Prediction2Waymo(
            final_results,
383
384
385
386
            waymo_tfrecords_dir,
            waymo_results_save_dir,
            waymo_results_final_path,
            prefix,
387
            classes,
388
            backend_args=self.backend_args,
389
390
            from_kitti_format=self.convert_kitti_format,
            idx2metainfo=self.idx2metainfo)
391
392
        converter.convert()
        waymo_save_tmp_dir.cleanup()
393
394

        return final_results, waymo_save_tmp_dir
395
396

    def merge_multi_view_boxes(self, box_dict_per_frame: List[dict],
397
                               cam0_info: dict) -> dict:
398
        """Merge bounding boxes predicted from multi-view images.
399

400
        Args:
401
402
            box_dict_per_frame (List[dict]): The results of prediction for each
                camera.
403
            cam0_info (dict): Store the sample idx for the given frame.
404
405

        Returns:
406
            dict: Merged results.
407
408
409
410
411
412
413
414
        """
        box_dict = dict()
        # convert list[dict] to dict[list]
        for key in box_dict_per_frame[0].keys():
            box_dict[key] = list()
            for cam_idx in range(self.num_cams):
                box_dict[key].append(box_dict_per_frame[cam_idx][key])
        # merge each elements
415
        box_dict['sample_idx'] = cam0_info['image_id']
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
        for key in ['bbox', 'box3d_lidar', 'scores', 'label_preds']:
            box_dict[key] = np.concatenate(box_dict[key])

        # apply nms to box3d_lidar (box3d_camera are in different systems)
        # TODO: move this global setting into config
        nms_cfg = dict(
            use_rotate_nms=True,
            nms_across_levels=False,
            nms_pre=500,
            nms_thr=0.05,
            score_thr=0.001,
            min_bbox_size=0,
            max_per_frame=100)
        nms_cfg = Config(nms_cfg)
        lidar_boxes3d = LiDARInstance3DBoxes(
            torch.from_numpy(box_dict['box3d_lidar']).cuda())
        scores = torch.from_numpy(box_dict['scores']).cuda()
        labels = torch.from_numpy(box_dict['label_preds']).long().cuda()
434
        nms_scores = scores.new_zeros(scores.shape[0], len(self.classes) + 1)
435
436
437
438
439
440
441
442
443
444
        indices = labels.new_tensor(list(range(scores.shape[0])))
        nms_scores[indices, labels] = scores
        lidar_boxes3d_for_nms = xywhr2xyxyr(lidar_boxes3d.bev)
        boxes3d = lidar_boxes3d.tensor
        # generate attr scores from attr labels
        boxes3d, scores, labels = box3d_multiclass_nms(
            boxes3d, lidar_boxes3d_for_nms, nms_scores, nms_cfg.score_thr,
            nms_cfg.max_per_frame, nms_cfg)
        lidar_boxes3d = LiDARInstance3DBoxes(boxes3d)
        det = bbox3d2result(lidar_boxes3d, scores, labels)
445
        box_preds_lidar = det['bboxes_3d']
446
447
448
        scores = det['scores_3d']
        labels = det['labels_3d']
        # box_preds_camera is in the cam0 system
449
450
        lidar2cam = cam0_info['images'][self.default_cam_key]['lidar2img']
        lidar2cam = np.array(lidar2cam).astype(np.float32)
451
        box_preds_camera = box_preds_lidar.convert_to(
452
            Box3DMode.CAM, lidar2cam, correct_yaw=True)
453
454
455
        # Note: bbox is meaningless in final evaluation, set to 0
        merged_box_dict = dict(
            bbox=np.zeros([box_preds_lidar.tensor.shape[0], 4]),
456
457
            box3d_camera=box_preds_camera.numpy(),
            box3d_lidar=box_preds_lidar.numpy(),
458
459
            scores=scores.numpy(),
            label_preds=labels.numpy(),
460
            sample_idx=box_dict['sample_idx'],
461
462
463
        )
        return merged_box_dict

464
465
466
467
468
469
470
    def bbox2result_kitti(
            self,
            net_outputs: List[dict],
            sample_idx_list: List[int],
            class_names: List[str],
            pklfile_prefix: Optional[str] = None,
            submission_prefix: Optional[str] = None) -> List[dict]:
471
472
473
474
        """Convert 3D detection results to kitti format for evaluation and test
        submission.

        Args:
475
476
            net_outputs (List[dict]): List of dict storing the inferenced
                bounding boxes and scores.
477
478
            sample_idx_list (List[int]): List of input sample idx.
            class_names (List[str]): A list of class names.
479
480
481
482
483
484
            pklfile_prefix (str, optional): The prefix of pkl file.
                Defaults to None.
            submission_prefix (str, optional): The prefix of submission file.
                Defaults to None.

        Returns:
485
            List[dict]: A list of dictionaries with the kitti format.
486
487
        """
        if submission_prefix is not None:
488
            mmengine.mkdir_or_exist(submission_prefix)
489
490
491
492

        det_annos = []
        print('\nConverting prediction to KITTI format')
        for idx, pred_dicts in enumerate(
493
                mmengine.track_iter_progress(net_outputs)):
494
            sample_idx = sample_idx_list[idx]
495
496
            info = self.data_infos[sample_idx]

497
            if self.load_type == 'mv_image_based':
498
499
                if idx % self.num_cams == 0:
                    box_dict_per_frame = []
500
501
502
503
504
505
506
507
508
509
510
511
512
513
                    cam0_key = list(info['images'].keys())[0]
                    cam0_info = info
                    # Here in mono3d, we use the 'CAM_FRONT' "the first
                    # index in the camera" as the default image shape.
                    # If you want to another camera, please modify it.
                    image_shape = (info['images'][cam0_key]['height'],
                                   info['images'][cam0_key]['width'])
                box_dict = self.convert_valid_bboxes(pred_dicts, info)
            else:
                box_dict = self.convert_valid_bboxes(pred_dicts, info)
                # Here default used 'CAM_FRONT' to compute metric.
                # If you want to use another camera, please modify it.
                image_shape = (info['images'][self.default_cam_key]['height'],
                               info['images'][self.default_cam_key]['width'])
514
            if self.load_type == 'mv_image_based':
515
516
517
                box_dict_per_frame.append(box_dict)
                if (idx + 1) % self.num_cams != 0:
                    continue
518
519
520
                box_dict = self.merge_multi_view_boxes(box_dict_per_frame,
                                                       cam0_info)

521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
            anno = {
                'name': [],
                'truncated': [],
                'occluded': [],
                'alpha': [],
                'bbox': [],
                'dimensions': [],
                'location': [],
                'rotation_y': [],
                'score': []
            }
            if len(box_dict['bbox']) > 0:
                box_2d_preds = box_dict['bbox']
                box_preds = box_dict['box3d_camera']
                scores = box_dict['scores']
                box_preds_lidar = box_dict['box3d_lidar']
                label_preds = box_dict['label_preds']

                for box, box_lidar, bbox, score, label in zip(
                        box_preds, box_preds_lidar, box_2d_preds, scores,
                        label_preds):
                    bbox[2:] = np.minimum(bbox[2:], image_shape[::-1])
                    bbox[:2] = np.maximum(bbox[:2], [0, 0])
                    anno['name'].append(class_names[int(label)])
                    anno['truncated'].append(0.0)
                    anno['occluded'].append(0)
                    anno['alpha'].append(
                        -np.arctan2(-box_lidar[1], box_lidar[0]) + box[6])
                    anno['bbox'].append(bbox)
                    anno['dimensions'].append(box[3:6])
                    anno['location'].append(box[:3])
                    anno['rotation_y'].append(box[6])
                    anno['score'].append(score)

                anno = {k: np.stack(v) for k, v in anno.items()}
            else:
                anno = {
                    'name': np.array([]),
                    'truncated': np.array([]),
                    'occluded': np.array([]),
                    'alpha': np.array([]),
                    'bbox': np.zeros([0, 4]),
                    'dimensions': np.zeros([0, 3]),
                    'location': np.zeros([0, 3]),
                    'rotation_y': np.array([]),
                    'score': np.array([]),
                }

            if submission_prefix is not None:
                curr_file = f'{submission_prefix}/{sample_idx:06d}.txt'
                with open(curr_file, 'w') as f:
                    bbox = anno['bbox']
                    loc = anno['location']
                    dims = anno['dimensions']  # lhw -> hwl

                    for idx in range(len(bbox)):
                        print(
                            '{} -1 -1 {:.4f} {:.4f} {:.4f} {:.4f} '
                            '{:.4f} {:.4f} {:.4f} '
                            '{:.4f} {:.4f} {:.4f} {:.4f} {:.4f} {:.4f}'.format(
                                anno['name'][idx], anno['alpha'][idx],
                                bbox[idx][0], bbox[idx][1], bbox[idx][2],
                                bbox[idx][3], dims[idx][1], dims[idx][2],
                                dims[idx][0], loc[idx][0], loc[idx][1],
                                loc[idx][2], anno['rotation_y'][idx],
                                anno['score'][idx]),
                            file=f)
            if self.use_pred_sample_idx:
                save_sample_idx = sample_idx
            else:
                # use the sample idx in the info file
                # In waymo validation sample_idx in prediction is 000xxx
                # but in info file it is 1000xxx
                save_sample_idx = box_dict['sample_idx']
595
596
            anno['sample_idx'] = np.array(
                [save_sample_idx] * len(anno['score']), dtype=np.int64)
597

598
            det_annos.append(anno)
599
600
601
602
603
604

        if pklfile_prefix is not None:
            if not pklfile_prefix.endswith(('.pkl', '.pickle')):
                out = f'{pklfile_prefix}.pkl'
            else:
                out = pklfile_prefix
605
            mmengine.dump(det_annos, out)
606
607
608
            print(f'Result is saved to {out}.')

        return det_annos
609

610
    def convert_valid_bboxes(self, box_dict: dict, info: dict) -> dict:
611
        """Convert the predicted boxes into valid ones. Should handle the
612
        load_model (frame_based, mv_image_based, fov_image_based), separately.
613
614
615
616

        Args:
            box_dict (dict): Box dictionaries to be converted.

617
618
619
                - bboxes_3d (:obj:`BaseInstance3DBoxes`): 3D bounding boxes.
                - scores_3d (Tensor): Scores of boxes.
                - labels_3d (Tensor): Class labels of boxes.
620
621
622
623
624
            info (dict): Data info.

        Returns:
            dict: Valid predicted boxes.

625
626
627
628
629
630
631
            - bbox (np.ndarray): 2D bounding boxes.
            - box3d_camera (np.ndarray): 3D bounding boxes in camera
              coordinate.
            - box3d_lidar (np.ndarray): 3D bounding boxes in LiDAR coordinate.
            - scores (np.ndarray): Scores of boxes.
            - label_preds (np.ndarray): Class label predictions.
            - sample_idx (int): Sample index.
632
633
634
635
636
        """
        # TODO: refactor this function
        box_preds = box_dict['bboxes_3d']
        scores = box_dict['scores_3d']
        labels = box_dict['labels_3d']
637
        sample_idx = info['sample_idx']
638
639
640
641
642
643
644
645
646
647
        box_preds.limit_yaw(offset=0.5, period=np.pi * 2)

        if len(box_preds) == 0:
            return dict(
                bbox=np.zeros([0, 4]),
                box3d_camera=np.zeros([0, 7]),
                box3d_lidar=np.zeros([0, 7]),
                scores=np.zeros([0]),
                label_preds=np.zeros([0, 4]),
                sample_idx=sample_idx)
648
        # Here default used 'CAM_FRONT' to compute metric. If you want to
649
        # use another camera, please modify it.
650
        if self.load_type in ['frame_based', 'fov_image_based']:
651
            cam_key = self.default_cam_key
652
        elif self.load_type == 'mv_image_based':
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
            cam_key = list(info['images'].keys())[0]
        else:
            raise NotImplementedError

        lidar2cam = np.array(info['images'][cam_key]['lidar2cam']).astype(
            np.float32)
        P2 = np.array(info['images'][cam_key]['cam2img']).astype(np.float32)
        img_shape = (info['images'][cam_key]['height'],
                     info['images'][cam_key]['width'])
        P2 = box_preds.tensor.new_tensor(P2)

        if isinstance(box_preds, LiDARInstance3DBoxes):
            box_preds_camera = box_preds.convert_to(Box3DMode.CAM, lidar2cam)
            box_preds_lidar = box_preds
        elif isinstance(box_preds, CameraInstance3DBoxes):
            box_preds_camera = box_preds
            box_preds_lidar = box_preds.convert_to(Box3DMode.LIDAR,
                                                   np.linalg.inv(lidar2cam))

        box_corners = box_preds_camera.corners
        box_corners_in_image = points_cam2img(box_corners, P2)
        # box_corners_in_image: [N, 8, 2]
        minxy = torch.min(box_corners_in_image, dim=1)[0]
        maxxy = torch.max(box_corners_in_image, dim=1)[0]
        box_2d_preds = torch.cat([minxy, maxxy], dim=1)
        # Post-processing
        # check box_preds_camera
        image_shape = box_preds.tensor.new_tensor(img_shape)
        valid_cam_inds = ((box_2d_preds[:, 0] < image_shape[1]) &
                          (box_2d_preds[:, 1] < image_shape[0]) &
                          (box_2d_preds[:, 2] > 0) & (box_2d_preds[:, 3] > 0))
        # check box_preds_lidar
685
        if self.load_type in ['frame_based']:
686
687
688
689
            limit_range = box_preds.tensor.new_tensor(self.pcd_limit_range)
            valid_pcd_inds = ((box_preds_lidar.center > limit_range[:3]) &
                              (box_preds_lidar.center < limit_range[3:]))
            valid_inds = valid_pcd_inds.all(-1)
690
        elif self.load_type in ['mv_image_based', 'fov_image_based']:
691
692
693
694
695
696
            valid_inds = valid_cam_inds

        if valid_inds.sum() > 0:
            return dict(
                bbox=box_2d_preds[valid_inds, :].numpy(),
                pred_box_type_3d=type(box_preds),
697
698
                box3d_camera=box_preds_camera[valid_inds].numpy(),
                box3d_lidar=box_preds_lidar[valid_inds].numpy(),
699
700
701
702
703
704
705
706
707
708
709
710
                scores=scores[valid_inds].numpy(),
                label_preds=labels[valid_inds].numpy(),
                sample_idx=sample_idx)
        else:
            return dict(
                bbox=np.zeros([0, 4]),
                pred_box_type_3d=type(box_preds),
                box3d_camera=np.zeros([0, 7]),
                box3d_lidar=np.zeros([0, 7]),
                scores=np.zeros([0]),
                label_preds=np.zeros([0]),
                sample_idx=sample_idx)