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

            - 'frame_based': Load all of the instances in the frame.
            - 'mv_image_based': Load all of the instances in the frame and need
56
57
58
59
60
61
62
63
64
65
              to convert to the FOV-based data type to support image-based
              detector.
            - '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.
        default_cam_key (str): The default camera for lidar to camera
            conversion. By default, KITTI: 'CAM2', Waymo: 'CAM_FRONT'.
            Defaults to 'CAM_FRONT'.
        use_pred_sample_idx (bool): In formating results, use the
            sample index from the prediction or from the load annotations.
66
            By default, KITTI: True, Waymo: False, Waymo has a conversion
67
68
            process, which needs to use the sample idx from load annotation.
            Defaults to False.
69
70
71
        collect_device (str): Device name used for collecting results
            from different ranks during distributed training. Must be 'cpu' or
            'gpu'. Defaults to 'cpu'.
72
        file_client_args (dict): File client for reading gt in waymo format.
73
            Defaults to ``dict(backend='disk')``.
74
75
76
77
        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.
78
    """
79
    num_cams = 5
80
81
82
83
84
85

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

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

111
        super(WaymoMetric, self).__init__(
112
113
114
115
116
117
118
            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,
119
120
            collect_device=collect_device,
            file_client_args=file_client_args)
121
122
123
124
125
126
127
        self.format_only = format_only
        if self.format_only:
            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.'

128
129
        self.default_prefix = 'Waymo metric'

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

        Args:
134
            results (List[dict]): The processed results of the whole dataset.
135
136
137
138
139
140

        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()
141
        self.classes = self.dataset_meta['classes']
142
143

        # load annotations
144
        self.data_infos = load(self.ann_file)['data_list']
145
146
        assert len(results) == len(self.data_infos), \
            'invalid list length of network outputs'
147
        # different from kitti, waymo do not need to convert the ann file
148
149
        # handle the mv_image_based load_mode
        if self.load_type == 'mv_image_based':
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
            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']

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

        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)

191
        metric_dict = {}
192
193
194
195
196
197

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

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

        if tmp_dir is not None:
            tmp_dir.cleanup()
207
        return metric_dict
208

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

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

        Returns:
223
            Dict[str, float]: Results of each evaluation metric.
224
225
226
227
228
229
230
231
232
        """

        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)
233
            ret_bytes = subprocess.check_output(eval_str, shell=True)
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
323
324
325
            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
326
327
        return ap_dict

328
329
330
331
332
333
334
    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]]:
335
        """Format the results to bin file.
336
337

        Args:
338
            results (List[dict]): Testing results of the dataset.
339
340
341
            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.
342
                Defaults to None.
343
344
345
            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.
346
347
348
                Defaults to None.
            classes (List[str], optional): A list of class name.
                Defaults to None.
349
350
351

        Returns:
            tuple: (result_dict, tmp_dir), result_dict is a dict containing
352
353
            the formatted result, tmp_dir is the temporal directory created
            for saving json files when jsonfile_prefix is not specified.
354
        """
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
        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)
371
372
373
374
375
376
377
378
379
380

        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.')
381
382
383
384
385

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

        return final_results, waymo_save_tmp_dir
398
399

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

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

        Returns:
409
            dict: Merged results.
410
411
412
413
414
415
416
417
        """
        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
418
        box_dict['sample_idx'] = cam0_info['image_id']
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
        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()
437
        nms_scores = scores.new_zeros(scores.shape[0], len(self.classes) + 1)
438
439
440
441
442
443
444
445
446
447
        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)
448
        box_preds_lidar = det['bboxes_3d']
449
450
451
        scores = det['scores_3d']
        labels = det['labels_3d']
        # box_preds_camera is in the cam0 system
452
453
        lidar2cam = cam0_info['images'][self.default_cam_key]['lidar2img']
        lidar2cam = np.array(lidar2cam).astype(np.float32)
454
        box_preds_camera = box_preds_lidar.convert_to(
455
            Box3DMode.CAM, lidar2cam, correct_yaw=True)
456
457
458
459
460
461
462
        # Note: bbox is meaningless in final evaluation, set to 0
        merged_box_dict = dict(
            bbox=np.zeros([box_preds_lidar.tensor.shape[0], 4]),
            box3d_camera=box_preds_camera.tensor.numpy(),
            box3d_lidar=box_preds_lidar.tensor.numpy(),
            scores=scores.numpy(),
            label_preds=labels.numpy(),
463
            sample_idx=box_dict['sample_idx'],
464
465
466
        )
        return merged_box_dict

467
468
469
470
471
472
473
    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]:
474
475
476
477
        """Convert 3D detection results to kitti format for evaluation and test
        submission.

        Args:
478
            net_outputs (List[dict]): List of dict storing the
479
                inferenced bounding boxes and scores.
480
481
            sample_idx_list (List[int]): List of input sample idx.
            class_names (List[str]): A list of class names.
482
483
484
485
486
487
            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:
488
            List[dict]: A list of dictionaries with the kitti format.
489
490
        """
        if submission_prefix is not None:
491
            mmengine.mkdir_or_exist(submission_prefix)
492
493
494
495

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

500
            if self.load_type == 'mv_image_based':
501
502
                if idx % self.num_cams == 0:
                    box_dict_per_frame = []
503
504
505
506
507
508
509
510
511
512
513
514
515
516
                    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'])
517
            if self.load_type == 'mv_image_based':
518
519
520
                box_dict_per_frame.append(box_dict)
                if (idx + 1) % self.num_cams != 0:
                    continue
521
522
523
                box_dict = self.merge_multi_view_boxes(box_dict_per_frame,
                                                       cam0_info)

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
595
596
597
            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']
598
599
            anno['sample_idx'] = np.array(
                [save_sample_idx] * len(anno['score']), dtype=np.int64)
600

601
            det_annos.append(anno)
602
603
604
605
606
607

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

        return det_annos
612

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

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

620
621
622
                - bboxes_3d (:obj:`BaseInstance3DBoxes`): 3D bounding boxes.
                - scores_3d (Tensor): Scores of boxes.
                - labels_3d (Tensor): Class labels of boxes.
623
624
625
626
627
628
629
            info (dict): Data info.

        Returns:
            dict: Valid predicted boxes.

                - bbox (np.ndarray): 2D bounding boxes.
                - box3d_camera (np.ndarray): 3D bounding boxes in
630
                  camera coordinate.
631
                - box3d_lidar (np.ndarray): 3D bounding boxes in
632
                  LiDAR coordinate.
633
634
635
636
637
638
639
640
                - scores (np.ndarray): Scores of boxes.
                - label_preds (np.ndarray): Class label predictions.
                - sample_idx (int): Sample index.
        """
        # TODO: refactor this function
        box_preds = box_dict['bboxes_3d']
        scores = box_dict['scores_3d']
        labels = box_dict['labels_3d']
641
        sample_idx = info['sample_idx']
642
643
644
645
646
647
648
649
650
651
        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)
652
        # Here default used 'CAM_FRONT' to compute metric. If you want to
653
        # use another camera, please modify it.
654
        if self.load_type in ['frame_based', 'fov_image_based']:
655
            cam_key = self.default_cam_key
656
        elif self.load_type == 'mv_image_based':
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
685
686
687
688
            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
689
        if self.load_type in ['frame_based']:
690
691
692
693
            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)
694
        elif self.load_type in ['mv_image_based', 'fov_image_based']:
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
            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),
                box3d_camera=box_preds_camera[valid_inds].tensor.numpy(),
                box3d_lidar=box_preds_lidar[valid_inds].tensor.numpy(),
                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)