kitti_dataset.py 30.4 KB
Newer Older
dingchang's avatar
dingchang committed
1
# Copyright (c) OpenMMLab. All rights reserved.
zhangwenwei's avatar
zhangwenwei committed
2
3
4
import copy
import mmcv
import numpy as np
zhangwenwei's avatar
zhangwenwei committed
5
6
import os
import tempfile
zhangwenwei's avatar
zhangwenwei committed
7
import torch
zhangwenwei's avatar
zhangwenwei committed
8
from mmcv.utils import print_log
zhangwenwei's avatar
zhangwenwei committed
9
from os import path as osp
zhangwenwei's avatar
zhangwenwei committed
10

zhangwenwei's avatar
zhangwenwei committed
11
from mmdet.datasets import DATASETS
12
from ..core import show_multi_modality_result, show_result
13
from ..core.bbox import (Box3DMode, CameraInstance3DBoxes, Coord3DMode,
14
                         LiDARInstance3DBoxes, points_cam2img)
zhangwenwei's avatar
zhangwenwei committed
15
from .custom_3d import Custom3DDataset
16
from .pipelines import Compose
zhangwenwei's avatar
zhangwenwei committed
17
18


19
@DATASETS.register_module()
zhangwenwei's avatar
zhangwenwei committed
20
class KittiDataset(Custom3DDataset):
zhangwenwei's avatar
zhangwenwei committed
21
    r"""KITTI Dataset.
wangtai's avatar
wangtai committed
22

zhangwenwei's avatar
zhangwenwei committed
23
24
    This class serves as the API for experiments on the `KITTI Dataset
    <http://www.cvlibs.net/datasets/kitti/eval_object.php?obj_benchmark=3d>`_.
wangtai's avatar
wangtai committed
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42

    Args:
        data_root (str): Path of dataset root.
        ann_file (str): Path of annotation file.
        split (str): Split of input data.
        pts_prefix (str, optional): Prefix of points files.
            Defaults to 'velodyne'.
        pipeline (list[dict], optional): Pipeline used for data processing.
            Defaults to None.
        classes (tuple[str], optional): Classes used in the dataset.
            Defaults to None.
        modality (dict, optional): Modality to specify the sensor data used
            as input. Defaults to None.
        box_type_3d (str, optional): Type of 3D box of this dataset.
            Based on the `box_type_3d`, the dataset will encapsulate the box
            to its original format then converted them to `box_type_3d`.
            Defaults to 'LiDAR' in this dataset. Available options includes

wangtai's avatar
wangtai committed
43
44
45
            - 'LiDAR': Box in LiDAR coordinates.
            - 'Depth': Box in depth coordinates, usually for indoor dataset.
            - 'Camera': Box in camera coordinates.
wangtai's avatar
wangtai committed
46
47
48
49
        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.
Wenwei Zhang's avatar
Wenwei Zhang committed
50
51
        pcd_limit_range (list): The range of point cloud used to filter
            invalid predicted boxes. Default: [0, -40, -3, 70.4, 40, 0.0].
wangtai's avatar
wangtai committed
52
    """
zhangwenwei's avatar
zhangwenwei committed
53
54
55
    CLASSES = ('car', 'pedestrian', 'cyclist')

    def __init__(self,
zhangwenwei's avatar
zhangwenwei committed
56
                 data_root,
zhangwenwei's avatar
zhangwenwei committed
57
58
                 ann_file,
                 split,
zhangwenwei's avatar
zhangwenwei committed
59
                 pts_prefix='velodyne',
zhangwenwei's avatar
zhangwenwei committed
60
                 pipeline=None,
zhangwenwei's avatar
zhangwenwei committed
61
                 classes=None,
zhangwenwei's avatar
zhangwenwei committed
62
                 modality=None,
63
64
                 box_type_3d='LiDAR',
                 filter_empty_gt=True,
Wenwei Zhang's avatar
Wenwei Zhang committed
65
66
                 test_mode=False,
                 pcd_limit_range=[0, -40, -3, 70.4, 40, 0.0]):
zhangwenwei's avatar
zhangwenwei committed
67
68
69
70
71
72
        super().__init__(
            data_root=data_root,
            ann_file=ann_file,
            pipeline=pipeline,
            classes=classes,
            modality=modality,
73
74
            box_type_3d=box_type_3d,
            filter_empty_gt=filter_empty_gt,
zhangwenwei's avatar
zhangwenwei committed
75
76
            test_mode=test_mode)

Wenwei Zhang's avatar
Wenwei Zhang committed
77
        self.split = split
zhangwenwei's avatar
zhangwenwei committed
78
        self.root_split = os.path.join(self.data_root, split)
zhangwenwei's avatar
zhangwenwei committed
79
        assert self.modality is not None
Wenwei Zhang's avatar
Wenwei Zhang committed
80
        self.pcd_limit_range = pcd_limit_range
zhangwenwei's avatar
zhangwenwei committed
81
        self.pts_prefix = pts_prefix
zhangwenwei's avatar
zhangwenwei committed
82

zhangwenwei's avatar
zhangwenwei committed
83
    def _get_pts_filename(self, idx):
84
85
86
87
88
89
90
91
        """Get point cloud filename according to the given index.

        Args:
            index (int): Index of the point cloud file to get.

        Returns:
            str: Name of the point cloud file.
        """
zhangwenwei's avatar
zhangwenwei committed
92
93
94
        pts_filename = osp.join(self.root_split, self.pts_prefix,
                                f'{idx:06d}.bin')
        return pts_filename
zhangwenwei's avatar
zhangwenwei committed
95

zhangwenwei's avatar
zhangwenwei committed
96
    def get_data_info(self, index):
97
98
99
100
101
102
        """Get data info according to the given index.

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

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

wangtai's avatar
wangtai committed
106
107
108
109
110
111
112
                - sample_idx (str): Sample index.
                - pts_filename (str): Filename of point clouds.
                - img_prefix (str | None): Prefix of image files.
                - img_info (dict): Image info.
                - lidar2img (list[np.ndarray], optional): Transformations \
                    from lidar to different cameras.
                - ann_info (dict): Annotation info.
113
        """
zhangwenwei's avatar
zhangwenwei committed
114
        info = self.data_infos[index]
zhangwenwei's avatar
zhangwenwei committed
115
        sample_idx = info['image']['image_idx']
zhangwenwei's avatar
zhangwenwei committed
116
        img_filename = os.path.join(self.data_root,
zhangwenwei's avatar
zhangwenwei committed
117
118
                                    info['image']['image_path'])

zhangwenwei's avatar
zhangwenwei committed
119
120
121
122
123
124
        # TODO: consider use torch.Tensor only
        rect = info['calib']['R0_rect'].astype(np.float32)
        Trv2c = info['calib']['Tr_velo_to_cam'].astype(np.float32)
        P2 = info['calib']['P2'].astype(np.float32)
        lidar2img = P2 @ rect @ Trv2c

zhangwenwei's avatar
zhangwenwei committed
125
        pts_filename = self._get_pts_filename(sample_idx)
zhangwenwei's avatar
zhangwenwei committed
126
127
        input_dict = dict(
            sample_idx=sample_idx,
zhangwenwei's avatar
zhangwenwei committed
128
            pts_filename=pts_filename,
zhangwenwei's avatar
zhangwenwei committed
129
130
            img_prefix=None,
            img_info=dict(filename=img_filename),
zhangwenwei's avatar
zhangwenwei committed
131
132
133
            lidar2img=lidar2img)

        if not self.test_mode:
zhangwenwei's avatar
zhangwenwei committed
134
            annos = self.get_ann_info(index)
zhangwenwei's avatar
zhangwenwei committed
135
            input_dict['ann_info'] = annos
zhangwenwei's avatar
zhangwenwei committed
136
137
138
139

        return input_dict

    def get_ann_info(self, index):
140
141
142
143
144
145
        """Get annotation info according to the given index.

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

        Returns:
zhangwenwei's avatar
zhangwenwei committed
146
            dict: annotation information consists of the following keys:
147

zhangwenwei's avatar
zhangwenwei committed
148
                - gt_bboxes_3d (:obj:`LiDARInstance3DBoxes`): \
wangtai's avatar
wangtai committed
149
150
151
152
153
                    3D ground truth bboxes.
                - gt_labels_3d (np.ndarray): Labels of ground truths.
                - gt_bboxes (np.ndarray): 2D ground truth bboxes.
                - gt_labels (np.ndarray): Labels of ground truths.
                - gt_names (list[str]): Class names of ground truths.
154
                - difficulty (int): kitti difficulty.
155
        """
zhangwenwei's avatar
zhangwenwei committed
156
        # Use index to get the annos, thus the evalhook could also use this api
zhangwenwei's avatar
zhangwenwei committed
157
        info = self.data_infos[index]
zhangwenwei's avatar
zhangwenwei committed
158
159
160
        rect = info['calib']['R0_rect'].astype(np.float32)
        Trv2c = info['calib']['Tr_velo_to_cam'].astype(np.float32)

161
        difficulty = info['annos']['difficulty']
zhangwenwei's avatar
zhangwenwei committed
162
163
        annos = info['annos']
        # we need other objects to avoid collision when sample
164
        annos = self.remove_dontcare(annos)
zhangwenwei's avatar
zhangwenwei committed
165
166
167
168
169
170
        loc = annos['location']
        dims = annos['dimensions']
        rots = annos['rotation_y']
        gt_names = annos['name']
        gt_bboxes_3d = np.concatenate([loc, dims, rots[..., np.newaxis]],
                                      axis=1).astype(np.float32)
171
172
173

        # convert gt_bboxes_3d to velodyne coordinates
        gt_bboxes_3d = CameraInstance3DBoxes(gt_bboxes_3d).convert_to(
174
            self.box_mode_3d, np.linalg.inv(rect @ Trv2c))
zhangwenwei's avatar
zhangwenwei committed
175
176
177
178
179
180
181
182
183
184
185
186
        gt_bboxes = annos['bbox']

        selected = self.drop_arrays_by_name(gt_names, ['DontCare'])
        gt_bboxes = gt_bboxes[selected].astype('float32')
        gt_names = gt_names[selected]

        gt_labels = []
        for cat in gt_names:
            if cat in self.CLASSES:
                gt_labels.append(self.CLASSES.index(cat))
            else:
                gt_labels.append(-1)
Wenwei Zhang's avatar
Wenwei Zhang committed
187
        gt_labels = np.array(gt_labels).astype(np.int64)
zhangwenwei's avatar
zhangwenwei committed
188
        gt_labels_3d = copy.deepcopy(gt_labels)
zhangwenwei's avatar
zhangwenwei committed
189
190
191

        anns_results = dict(
            gt_bboxes_3d=gt_bboxes_3d,
zhangwenwei's avatar
zhangwenwei committed
192
            gt_labels_3d=gt_labels_3d,
zhangwenwei's avatar
zhangwenwei committed
193
            bboxes=gt_bboxes,
liyinhao's avatar
liyinhao committed
194
            labels=gt_labels,
195
196
            gt_names=gt_names,
            difficulty=difficulty)
zhangwenwei's avatar
zhangwenwei committed
197
198
199
        return anns_results

    def drop_arrays_by_name(self, gt_names, used_classes):
200
201
202
203
204
205
206
207
208
        """Drop irrelevant ground truths by name.

        Args:
            gt_names (list[str]): Names of ground truths.
            used_classes (list[str]): Classes of interest.

        Returns:
            np.ndarray: Indices of ground truths that will be dropped.
        """
zhangwenwei's avatar
zhangwenwei committed
209
210
211
212
213
        inds = [i for i, x in enumerate(gt_names) if x not in used_classes]
        inds = np.array(inds, dtype=np.int64)
        return inds

    def keep_arrays_by_name(self, gt_names, used_classes):
214
215
216
217
218
219
220
221
222
        """Keep useful ground truths by name.

        Args:
            gt_names (list[str]): Names of ground truths.
            used_classes (list[str]): Classes of interest.

        Returns:
            np.ndarray: Indices of ground truths that will be keeped.
        """
zhangwenwei's avatar
zhangwenwei committed
223
224
225
226
        inds = [i for i, x in enumerate(gt_names) if x in used_classes]
        inds = np.array(inds, dtype=np.int64)
        return inds

227
    def remove_dontcare(self, ann_info):
228
229
230
231
232
233
234
235
236
        """Remove annotations that do not need to be cared.

        Args:
            ann_info (dict): Dict of annotation infos. The ``'DontCare'``
                annotations will be removed according to ann_file['name'].

        Returns:
            dict: Annotations after filtering.
        """
237
238
239
240
241
242
243
244
245
        img_filtered_annotations = {}
        relevant_annotation_indices = [
            i for i, x in enumerate(ann_info['name']) if x != 'DontCare'
        ]
        for key in ann_info.keys():
            img_filtered_annotations[key] = (
                ann_info[key][relevant_annotation_indices])
        return img_filtered_annotations

246
247
248
249
    def format_results(self,
                       outputs,
                       pklfile_prefix=None,
                       submission_prefix=None):
250
251
252
253
254
255
256
257
258
259
260
261
262
        """Format the results to pkl file.

        Args:
            outputs (list[dict]): Testing results of the dataset.
            pklfile_prefix (str | None): 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. Default: None.
            submission_prefix (str | None): 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.
                Default: None.

        Returns:
zhangwenwei's avatar
zhangwenwei committed
263
264
            tuple: (result_files, tmp_dir), result_files is a dict containing \
                the json filepaths, tmp_dir is the temporal directory created \
265
266
                for saving json files when jsonfile_prefix is not specified.
        """
267
268
269
270
271
272
        if pklfile_prefix is None:
            tmp_dir = tempfile.TemporaryDirectory()
            pklfile_prefix = osp.join(tmp_dir.name, 'results')
        else:
            tmp_dir = None

zhangwenwei's avatar
zhangwenwei committed
273
        if not isinstance(outputs[0], dict):
zhangwenwei's avatar
zhangwenwei committed
274
            result_files = self.bbox2result_kitti2d(outputs, self.CLASSES,
zhangwenwei's avatar
zhangwenwei committed
275
                                                    pklfile_prefix,
276
                                                    submission_prefix)
zhangwenwei's avatar
zhangwenwei committed
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
        elif 'pts_bbox' in outputs[0] or 'img_bbox' in outputs[0]:
            result_files = dict()
            for name in outputs[0]:
                results_ = [out[name] for out in outputs]
                pklfile_prefix_ = pklfile_prefix + name
                if submission_prefix is not None:
                    submission_prefix_ = submission_prefix + name
                else:
                    submission_prefix_ = None
                if 'img' in name:
                    result_files = self.bbox2result_kitti2d(
                        results_, self.CLASSES, pklfile_prefix_,
                        submission_prefix_)
                else:
                    result_files_ = self.bbox2result_kitti(
                        results_, self.CLASSES, pklfile_prefix_,
                        submission_prefix_)
                result_files[name] = result_files_
zhangwenwei's avatar
zhangwenwei committed
295
        else:
zhangwenwei's avatar
zhangwenwei committed
296
            result_files = self.bbox2result_kitti(outputs, self.CLASSES,
297
298
                                                  pklfile_prefix,
                                                  submission_prefix)
zhangwenwei's avatar
zhangwenwei committed
299
        return result_files, tmp_dir
zhangwenwei's avatar
zhangwenwei committed
300

301
302
303
304
305
    def evaluate(self,
                 results,
                 metric=None,
                 logger=None,
                 pklfile_prefix=None,
liyinhao's avatar
liyinhao committed
306
307
                 submission_prefix=None,
                 show=False,
308
309
                 out_dir=None,
                 pipeline=None):
310
311
312
        """Evaluation in KITTI protocol.

        Args:
wangtai's avatar
wangtai committed
313
            results (list[dict]): Testing results of the dataset.
314
315
316
317
318
319
            metric (str | list[str]): Metrics to be evaluated.
            logger (logging.Logger | str | None): Logger used for printing
                related information during evaluation. Default: None.
            pklfile_prefix (str | None): 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. Default: None.
320
            submission_prefix (str | None): The prefix of submission data.
321
                If not specified, the submission data will not be generated.
liyinhao's avatar
liyinhao committed
322
323
324
325
            show (bool): Whether to visualize.
                Default: False.
            out_dir (str): Path to save the visualization results.
                Default: None.
326
327
            pipeline (list[dict], optional): raw data loading for showing.
                Default: None.
328
329

        Returns:
wangtai's avatar
wangtai committed
330
            dict[str, float]: Results of each evaluation metric.
331
332
        """
        result_files, tmp_dir = self.format_results(results, pklfile_prefix)
zhangwenwei's avatar
zhangwenwei committed
333
        from mmdet3d.core.evaluation import kitti_eval
zhangwenwei's avatar
zhangwenwei committed
334
        gt_annos = [info['annos'] for info in self.data_infos]
zhangwenwei's avatar
zhangwenwei committed
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352

        if isinstance(result_files, dict):
            ap_dict = dict()
            for name, result_files_ in result_files.items():
                eval_types = ['bbox', 'bev', '3d']
                if 'img' in name:
                    eval_types = ['bbox']
                ap_result_str, ap_dict_ = kitti_eval(
                    gt_annos,
                    result_files_,
                    self.CLASSES,
                    eval_types=eval_types)
                for ap_type, ap in ap_dict_.items():
                    ap_dict[f'{name}/{ap_type}'] = float('{:.4f}'.format(ap))

                print_log(
                    f'Results of {name}:\n' + ap_result_str, logger=logger)

zhangwenwei's avatar
zhangwenwei committed
353
        else:
zhangwenwei's avatar
zhangwenwei committed
354
355
356
357
358
359
360
361
            if metric == 'img_bbox':
                ap_result_str, ap_dict = kitti_eval(
                    gt_annos, result_files, self.CLASSES, eval_types=['bbox'])
            else:
                ap_result_str, ap_dict = kitti_eval(gt_annos, result_files,
                                                    self.CLASSES)
            print_log('\n' + ap_result_str, logger=logger)

362
363
        if tmp_dir is not None:
            tmp_dir.cleanup()
liyinhao's avatar
liyinhao committed
364
        if show:
365
            self.show(results, out_dir, pipeline=pipeline)
366
        return ap_dict
367
368
369
370
371
372

    def bbox2result_kitti(self,
                          net_outputs,
                          class_names,
                          pklfile_prefix=None,
                          submission_prefix=None):
373
374
375
376
377
378
379
380
381
382
383
384
385
        """Convert 3D detection results to kitti format for evaluation and test
        submission.

        Args:
            net_outputs (list[np.ndarray]): List of array storing the \
                inferenced bounding boxes and scores.
            class_names (list[String]): A list of class names.
            pklfile_prefix (str | None): The prefix of pkl file.
            submission_prefix (str | None): The prefix of submission file.

        Returns:
            list[dict]: A list of dictionaries with the kitti format.
        """
Wenwei Zhang's avatar
Wenwei Zhang committed
386
387
        assert len(net_outputs) == len(self.data_infos), \
            'invalid list length of network outputs'
388
389
        if submission_prefix is not None:
            mmcv.mkdir_or_exist(submission_prefix)
zhangwenwei's avatar
zhangwenwei committed
390
391

        det_annos = []
zhangwenwei's avatar
zhangwenwei committed
392
        print('\nConverting prediction to KITTI format')
zhangwenwei's avatar
zhangwenwei committed
393
394
395
        for idx, pred_dicts in enumerate(
                mmcv.track_iter_progress(net_outputs)):
            annos = []
zhangwenwei's avatar
zhangwenwei committed
396
            info = self.data_infos[idx]
zhangwenwei's avatar
zhangwenwei committed
397
            sample_idx = info['image']['image_idx']
zhangwenwei's avatar
zhangwenwei committed
398
            image_shape = info['image']['image_shape'][:2]
zhangwenwei's avatar
zhangwenwei committed
399
            box_dict = self.convert_valid_bboxes(pred_dicts, info)
xiliu8006's avatar
xiliu8006 committed
400
401
402
403
404
405
406
407
408
409
410
            anno = {
                'name': [],
                'truncated': [],
                'occluded': [],
                'alpha': [],
                'bbox': [],
                'dimensions': [],
                'location': [],
                'rotation_y': [],
                'score': []
            }
zhangwenwei's avatar
zhangwenwei committed
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
            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()}
                annos.append(anno)
            else:
xiliu8006's avatar
xiliu8006 committed
437
                anno = {
zhangwenwei's avatar
zhangwenwei committed
438
439
440
441
442
443
444
445
446
                    '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([]),
xiliu8006's avatar
xiliu8006 committed
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
                }
                annos.append(anno)

            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)

zhangwenwei's avatar
zhangwenwei committed
470
471
            annos[-1]['sample_idx'] = np.array(
                [sample_idx] * len(annos[-1]['score']), dtype=np.int64)
zhangwenwei's avatar
zhangwenwei committed
472
473
474

            det_annos += annos

475
476
477
        if pklfile_prefix is not None:
            if not pklfile_prefix.endswith(('.pkl', '.pickle')):
                out = f'{pklfile_prefix}.pkl'
zhangwenwei's avatar
zhangwenwei committed
478
            mmcv.dump(det_annos, out)
Wenwei Zhang's avatar
Wenwei Zhang committed
479
            print(f'Result is saved to {out}.')
zhangwenwei's avatar
zhangwenwei committed
480
481
482
483
484
485

        return det_annos

    def bbox2result_kitti2d(self,
                            net_outputs,
                            class_names,
486
487
                            pklfile_prefix=None,
                            submission_prefix=None):
zhangwenwei's avatar
zhangwenwei committed
488
489
        """Convert 2D detection results to kitti format for evaluation and test
        submission.
zhangwenwei's avatar
zhangwenwei committed
490
491

        Args:
492
493
494
            net_outputs (list[np.ndarray]): List of array storing the \
                inferenced bounding boxes and scores.
            class_names (list[String]): A list of class names.
495
496
            pklfile_prefix (str | None): The prefix of pkl file.
            submission_prefix (str | None): The prefix of submission file.
zhangwenwei's avatar
zhangwenwei committed
497

498
        Returns:
499
            list[dict]: A list of dictionaries have the kitti format
zhangwenwei's avatar
zhangwenwei committed
500
        """
Wenwei Zhang's avatar
Wenwei Zhang committed
501
502
        assert len(net_outputs) == len(self.data_infos), \
            'invalid list length of network outputs'
zhangwenwei's avatar
zhangwenwei committed
503
        det_annos = []
zhangwenwei's avatar
zhangwenwei committed
504
        print('\nConverting prediction to KITTI format')
zhangwenwei's avatar
zhangwenwei committed
505
506
507
508
509
510
511
512
513
514
515
516
517
        for i, bboxes_per_sample in enumerate(
                mmcv.track_iter_progress(net_outputs)):
            annos = []
            anno = dict(
                name=[],
                truncated=[],
                occluded=[],
                alpha=[],
                bbox=[],
                dimensions=[],
                location=[],
                rotation_y=[],
                score=[])
zhangwenwei's avatar
zhangwenwei committed
518
            sample_idx = self.data_infos[i]['image']['image_idx']
zhangwenwei's avatar
zhangwenwei committed
519
520
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

            num_example = 0
            for label in range(len(bboxes_per_sample)):
                bbox = bboxes_per_sample[label]
                for i in range(bbox.shape[0]):
                    anno['name'].append(class_names[int(label)])
                    anno['truncated'].append(0.0)
                    anno['occluded'].append(0)
                    anno['alpha'].append(0.0)
                    anno['bbox'].append(bbox[i, :4])
                    # set dimensions (height, width, length) to zero
                    anno['dimensions'].append(
                        np.zeros(shape=[3], dtype=np.float32))
                    # set the 3D translation to (-1000, -1000, -1000)
                    anno['location'].append(
                        np.ones(shape=[3], dtype=np.float32) * (-1000.0))
                    anno['rotation_y'].append(0.0)
                    anno['score'].append(bbox[i, 4])
                    num_example += 1

            if num_example == 0:
                annos.append(
                    dict(
                        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([]),
                    ))
            else:
                anno = {k: np.stack(v) for k, v in anno.items()}
                annos.append(anno)

            annos[-1]['sample_idx'] = np.array(
                [sample_idx] * num_example, dtype=np.int64)
            det_annos += annos

560
561
562
563
564
565
566
567
        if pklfile_prefix is not None:
            # save file in pkl format
            pklfile_path = (
                pklfile_prefix[:-4] if pklfile_prefix.endswith(
                    ('.pkl', '.pickle')) else pklfile_prefix)
            mmcv.dump(det_annos, pklfile_path)

        if submission_prefix is not None:
zhangwenwei's avatar
zhangwenwei committed
568
            # save file in submission format
569
570
            mmcv.mkdir_or_exist(submission_prefix)
            print(f'Saving KITTI submission to {submission_prefix}')
zhangwenwei's avatar
zhangwenwei committed
571
            for i, anno in enumerate(det_annos):
zhangwenwei's avatar
zhangwenwei committed
572
                sample_idx = self.data_infos[i]['image']['image_idx']
573
                cur_det_file = f'{submission_prefix}/{sample_idx:06d}.txt'
zhangwenwei's avatar
zhangwenwei committed
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
                with open(cur_det_file, 'w') as f:
                    bbox = anno['bbox']
                    loc = anno['location']
                    dims = anno['dimensions'][::-1]  # 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],  # 4 float
                                *dims[idx],  # 3 float
                                *loc[idx],  # 3 float
                                anno['rotation_y'][idx],
                                anno['score'][idx]),
                            file=f,
                        )
591
            print(f'Result is saved to {submission_prefix}')
zhangwenwei's avatar
zhangwenwei committed
592
593
594
595

        return det_annos

    def convert_valid_bboxes(self, box_dict, info):
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
        """Convert the predicted boxes into valid ones.

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

                - boxes_3d (:obj:`LiDARInstance3DBoxes`): 3D bounding boxes.
                - scores_3d (torch.Tensor): Scores of boxes.
                - labels_3d (torch.Tensor): Class labels of boxes.
            info (dict): Data info.

        Returns:
            dict: Valid predicted boxes.

                - 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.
        """
zhangwenwei's avatar
zhangwenwei committed
618
        # TODO: refactor this function
619
620
621
        box_preds = box_dict['boxes_3d']
        scores = box_dict['scores_3d']
        labels = box_dict['labels_3d']
zhangwenwei's avatar
zhangwenwei committed
622
        sample_idx = info['image']['image_idx']
623
624
625
        # TODO: remove the hack of yaw
        box_preds.tensor[:, -1] = box_preds.tensor[:, -1] - np.pi
        box_preds.limit_yaw(offset=0.5, period=np.pi * 2)
zhangwenwei's avatar
zhangwenwei committed
626

627
        if len(box_preds) == 0:
zhangwenwei's avatar
zhangwenwei committed
628
            return dict(
629
630
631
632
633
634
                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)
zhangwenwei's avatar
zhangwenwei committed
635
636
637
638
639

        rect = info['calib']['R0_rect'].astype(np.float32)
        Trv2c = info['calib']['Tr_velo_to_cam'].astype(np.float32)
        P2 = info['calib']['P2'].astype(np.float32)
        img_shape = info['image']['image_shape']
640
641
642
643
644
        P2 = box_preds.tensor.new_tensor(P2)

        box_preds_camera = box_preds.convert_to(Box3DMode.CAM, rect @ Trv2c)

        box_corners = box_preds_camera.corners
zhangwenwei's avatar
zhangwenwei committed
645
        box_corners_in_image = points_cam2img(box_corners, P2)
zhangwenwei's avatar
zhangwenwei committed
646
647
648
649
650
        # 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
651
652
        # check box_preds_camera
        image_shape = box_preds.tensor.new_tensor(img_shape)
twang's avatar
twang committed
653
654
655
        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))
656
657
658
659
        # check box_preds
        limit_range = box_preds.tensor.new_tensor(self.pcd_limit_range)
        valid_pcd_inds = ((box_preds.center > limit_range[:3]) &
                          (box_preds.center < limit_range[3:]))
zhangwenwei's avatar
zhangwenwei committed
660
661
662
663
664
        valid_inds = valid_cam_inds & valid_pcd_inds.all(-1)

        if valid_inds.sum() > 0:
            return dict(
                bbox=box_2d_preds[valid_inds, :].numpy(),
665
666
667
668
                box3d_camera=box_preds_camera[valid_inds].tensor.numpy(),
                box3d_lidar=box_preds[valid_inds].tensor.numpy(),
                scores=scores[valid_inds].numpy(),
                label_preds=labels[valid_inds].numpy(),
669
                sample_idx=sample_idx)
zhangwenwei's avatar
zhangwenwei committed
670
671
        else:
            return dict(
672
673
674
675
676
                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]),
677
                sample_idx=sample_idx)
liyinhao's avatar
liyinhao committed
678

679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
    def _build_default_pipeline(self):
        """Build the default pipeline for this dataset."""
        pipeline = [
            dict(
                type='LoadPointsFromFile',
                coord_type='LIDAR',
                load_dim=4,
                use_dim=4,
                file_client_args=dict(backend='disk')),
            dict(
                type='DefaultFormatBundle3D',
                class_names=self.CLASSES,
                with_label=False),
            dict(type='Collect3D', keys=['points'])
        ]
        if self.modality['use_camera']:
            pipeline.insert(0, dict(type='LoadImageFromFile'))
        return Compose(pipeline)

    def show(self, results, out_dir, show=True, pipeline=None):
699
700
701
        """Results visualization.

        Args:
wangtai's avatar
wangtai committed
702
            results (list[dict]): List of bounding boxes results.
703
            out_dir (str): Output directory of visualization result.
704
            show (bool): Visualize the results online.
705
706
            pipeline (list[dict], optional): raw data loading for showing.
                Default: None.
707
        """
liyinhao's avatar
liyinhao committed
708
        assert out_dir is not None, 'Expect out_dir, got none.'
709
        pipeline = self._get_pipeline(pipeline)
liyinhao's avatar
liyinhao committed
710
        for i, result in enumerate(results):
711
712
            if 'pts_bbox' in result.keys():
                result = result['pts_bbox']
liyinhao's avatar
liyinhao committed
713
714
715
            data_info = self.data_infos[i]
            pts_path = data_info['point_cloud']['velodyne_path']
            file_name = osp.split(pts_path)[-1].split('.')[0]
716
717
718
            points, img_metas, img = self._extract_data(
                i, pipeline, ['points', 'img_metas', 'img'])
            points = points.numpy()
liyinhao's avatar
liyinhao committed
719
            # for now we convert points into depth mode
720
721
            points = Coord3DMode.convert_point(points, Coord3DMode.LIDAR,
                                               Coord3DMode.DEPTH)
722
723
724
            gt_bboxes = self.get_ann_info(i)['gt_bboxes_3d'].tensor.numpy()
            show_gt_bboxes = Box3DMode.convert(gt_bboxes, Box3DMode.LIDAR,
                                               Box3DMode.DEPTH)
liyinhao's avatar
liyinhao committed
725
            pred_bboxes = result['boxes_3d'].tensor.numpy()
726
727
728
729
730
731
            show_pred_bboxes = Box3DMode.convert(pred_bboxes, Box3DMode.LIDAR,
                                                 Box3DMode.DEPTH)
            show_result(points, show_gt_bboxes, show_pred_bboxes, out_dir,
                        file_name, show)

            # multi-modality visualization
732
733
734
735
            if self.modality['use_camera'] and 'lidar2img' in img_metas.keys():
                img = img.numpy()
                # need to transpose channel to first dim
                img = img.transpose(1, 2, 0)
736
737
738
739
740
741
742
743
                show_pred_bboxes = LiDARInstance3DBoxes(
                    pred_bboxes, origin=(0.5, 0.5, 0))
                show_gt_bboxes = LiDARInstance3DBoxes(
                    gt_bboxes, origin=(0.5, 0.5, 0))
                show_multi_modality_result(
                    img,
                    show_gt_bboxes,
                    show_pred_bboxes,
744
                    img_metas['lidar2img'],
745
746
                    out_dir,
                    file_name,
747
748
                    box_mode='lidar',
                    show=show)