kitti_dataset.py 31.3 KB
Newer Older
dingchang's avatar
dingchang committed
1
# Copyright (c) OpenMMLab. All rights reserved.
zhangwenwei's avatar
zhangwenwei committed
2
import copy
zhangwenwei's avatar
zhangwenwei committed
3
4
import os
import tempfile
5
6
7
8
from os import path as osp

import mmcv
import numpy as np
zhangwenwei's avatar
zhangwenwei committed
9
import torch
zhangwenwei's avatar
zhangwenwei committed
10
from mmcv.utils import print_log
zhangwenwei's avatar
zhangwenwei committed
11

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


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

zhangwenwei's avatar
zhangwenwei committed
24
25
    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
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43

    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
44
45
46
            - 'LiDAR': Box in LiDAR coordinates.
            - 'Depth': Box in depth coordinates, usually for indoor dataset.
            - 'Camera': Box in camera coordinates.
wangtai's avatar
wangtai committed
47
48
49
50
        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.
51
52
53
        pcd_limit_range (list, optional): 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
54
    """
zhangwenwei's avatar
zhangwenwei committed
55
56
57
    CLASSES = ('car', 'pedestrian', 'cyclist')

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

Wenwei Zhang's avatar
Wenwei Zhang committed
81
        self.split = split
zhangwenwei's avatar
zhangwenwei committed
82
        self.root_split = os.path.join(self.data_root, split)
zhangwenwei's avatar
zhangwenwei committed
83
        assert self.modality is not None
Wenwei Zhang's avatar
Wenwei Zhang committed
84
        self.pcd_limit_range = pcd_limit_range
zhangwenwei's avatar
zhangwenwei committed
85
        self.pts_prefix = pts_prefix
zhangwenwei's avatar
zhangwenwei committed
86

zhangwenwei's avatar
zhangwenwei committed
87
    def _get_pts_filename(self, idx):
88
89
90
91
92
93
94
95
        """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
96
97
98
        pts_filename = osp.join(self.root_split, self.pts_prefix,
                                f'{idx:06d}.bin')
        return pts_filename
zhangwenwei's avatar
zhangwenwei committed
99

zhangwenwei's avatar
zhangwenwei committed
100
    def get_data_info(self, index):
101
102
103
104
105
106
        """Get data info according to the given index.

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

        Returns:
107
            dict: Data information that will be passed to the data
zhangwenwei's avatar
zhangwenwei committed
108
                preprocessing pipelines. It includes the following keys:
109

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

zhangwenwei's avatar
zhangwenwei committed
123
124
125
126
127
128
        # 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
129
        pts_filename = self._get_pts_filename(sample_idx)
zhangwenwei's avatar
zhangwenwei committed
130
131
        input_dict = dict(
            sample_idx=sample_idx,
zhangwenwei's avatar
zhangwenwei committed
132
            pts_filename=pts_filename,
zhangwenwei's avatar
zhangwenwei committed
133
134
            img_prefix=None,
            img_info=dict(filename=img_filename),
zhangwenwei's avatar
zhangwenwei committed
135
136
137
            lidar2img=lidar2img)

        if not self.test_mode:
zhangwenwei's avatar
zhangwenwei committed
138
            annos = self.get_ann_info(index)
zhangwenwei's avatar
zhangwenwei committed
139
            input_dict['ann_info'] = annos
zhangwenwei's avatar
zhangwenwei committed
140
141
142
143

        return input_dict

    def get_ann_info(self, index):
144
145
146
147
148
149
        """Get annotation info according to the given index.

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

        Returns:
zhangwenwei's avatar
zhangwenwei committed
150
            dict: annotation information consists of the following keys:
151

152
                - gt_bboxes_3d (:obj:`LiDARInstance3DBoxes`):
wangtai's avatar
wangtai committed
153
154
155
156
157
                    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.
158
159
                - difficulty (int): Difficulty defined by KITTI.
                    0, 1, 2 represent xxxxx respectively.
160
        """
zhangwenwei's avatar
zhangwenwei committed
161
        # Use index to get the annos, thus the evalhook could also use this api
zhangwenwei's avatar
zhangwenwei committed
162
        info = self.data_infos[index]
zhangwenwei's avatar
zhangwenwei committed
163
164
165
        rect = info['calib']['R0_rect'].astype(np.float32)
        Trv2c = info['calib']['Tr_velo_to_cam'].astype(np.float32)

166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
        if 'plane' in info:
            # convert ground plane to velodyne coordinates
            reverse = np.linalg.inv(rect @ Trv2c)

            (plane_norm_cam,
             plane_off_cam) = (info['plane'][:3],
                               -info['plane'][:3] * info['plane'][3])
            plane_norm_lidar = \
                (reverse[:3, :3] @ plane_norm_cam[:, None])[:, 0]
            plane_off_lidar = (
                reverse[:3, :3] @ plane_off_cam[:, None][:, 0] +
                reverse[:3, 3])
            plane_lidar = np.zeros_like(plane_norm_lidar, shape=(4, ))
            plane_lidar[:3] = plane_norm_lidar
            plane_lidar[3] = -plane_norm_lidar.T @ plane_off_lidar
        else:
            plane_lidar = None

184
        difficulty = info['annos']['difficulty']
zhangwenwei's avatar
zhangwenwei committed
185
186
        annos = info['annos']
        # we need other objects to avoid collision when sample
187
        annos = self.remove_dontcare(annos)
zhangwenwei's avatar
zhangwenwei committed
188
189
190
191
192
193
        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)
194
195
196

        # convert gt_bboxes_3d to velodyne coordinates
        gt_bboxes_3d = CameraInstance3DBoxes(gt_bboxes_3d).convert_to(
197
            self.box_mode_3d, np.linalg.inv(rect @ Trv2c))
zhangwenwei's avatar
zhangwenwei committed
198
199
200
201
202
203
204
205
206
207
208
209
        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
210
        gt_labels = np.array(gt_labels).astype(np.int64)
zhangwenwei's avatar
zhangwenwei committed
211
        gt_labels_3d = copy.deepcopy(gt_labels)
zhangwenwei's avatar
zhangwenwei committed
212
213
214

        anns_results = dict(
            gt_bboxes_3d=gt_bboxes_3d,
zhangwenwei's avatar
zhangwenwei committed
215
            gt_labels_3d=gt_labels_3d,
zhangwenwei's avatar
zhangwenwei committed
216
            bboxes=gt_bboxes,
liyinhao's avatar
liyinhao committed
217
            labels=gt_labels,
218
            gt_names=gt_names,
219
            plane=plane_lidar,
220
            difficulty=difficulty)
zhangwenwei's avatar
zhangwenwei committed
221
222
223
        return anns_results

    def drop_arrays_by_name(self, gt_names, used_classes):
224
225
226
227
228
229
230
231
232
        """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
233
234
235
236
237
        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):
238
239
240
241
242
243
244
245
246
        """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
247
248
249
250
        inds = [i for i, x in enumerate(gt_names) if x in used_classes]
        inds = np.array(inds, dtype=np.int64)
        return inds

251
    def remove_dontcare(self, ann_info):
252
253
254
255
256
257
258
259
260
        """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.
        """
261
262
263
264
265
266
267
268
269
        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

270
271
272
273
    def format_results(self,
                       outputs,
                       pklfile_prefix=None,
                       submission_prefix=None):
274
275
276
277
        """Format the results to pkl file.

        Args:
            outputs (list[dict]): Testing results of the dataset.
278
            pklfile_prefix (str): The prefix of pkl files. It includes
279
280
                the file path and the prefix of filename, e.g., "a/b/prefix".
                If not specified, a temp file will be created. Default: None.
281
            submission_prefix (str): The prefix of submitted files. It
282
283
284
285
286
                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:
287
288
            tuple: (result_files, tmp_dir), result_files is a dict containing
                the json filepaths, tmp_dir is the temporal directory created
289
290
                for saving json files when jsonfile_prefix is not specified.
        """
291
292
293
294
295
296
        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
297
        if not isinstance(outputs[0], dict):
zhangwenwei's avatar
zhangwenwei committed
298
            result_files = self.bbox2result_kitti2d(outputs, self.CLASSES,
zhangwenwei's avatar
zhangwenwei committed
299
                                                    pklfile_prefix,
300
                                                    submission_prefix)
zhangwenwei's avatar
zhangwenwei committed
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
        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
319
        else:
zhangwenwei's avatar
zhangwenwei committed
320
            result_files = self.bbox2result_kitti(outputs, self.CLASSES,
321
322
                                                  pklfile_prefix,
                                                  submission_prefix)
zhangwenwei's avatar
zhangwenwei committed
323
        return result_files, tmp_dir
zhangwenwei's avatar
zhangwenwei committed
324

325
326
327
328
329
    def evaluate(self,
                 results,
                 metric=None,
                 logger=None,
                 pklfile_prefix=None,
liyinhao's avatar
liyinhao committed
330
331
                 submission_prefix=None,
                 show=False,
332
333
                 out_dir=None,
                 pipeline=None):
334
335
336
        """Evaluation in KITTI protocol.

        Args:
wangtai's avatar
wangtai committed
337
            results (list[dict]): Testing results of the dataset.
338
339
340
            metric (str | list[str], optional): Metrics to be evaluated.
                Default: None.
            logger (logging.Logger | str, optional): Logger used for printing
341
                related information during evaluation. Default: None.
342
            pklfile_prefix (str, optional): The prefix of pkl files, including
343
344
                the file path and the prefix of filename, e.g., "a/b/prefix".
                If not specified, a temp file will be created. Default: None.
345
            submission_prefix (str, optional): The prefix of submission data.
346
                If not specified, the submission data will not be generated.
347
                Default: None.
348
            show (bool, optional): Whether to visualize.
liyinhao's avatar
liyinhao committed
349
                Default: False.
350
            out_dir (str, optional): Path to save the visualization results.
liyinhao's avatar
liyinhao committed
351
                Default: None.
352
353
            pipeline (list[dict], optional): raw data loading for showing.
                Default: None.
354
355

        Returns:
wangtai's avatar
wangtai committed
356
            dict[str, float]: Results of each evaluation metric.
357
358
        """
        result_files, tmp_dir = self.format_results(results, pklfile_prefix)
zhangwenwei's avatar
zhangwenwei committed
359
        from mmdet3d.core.evaluation import kitti_eval
zhangwenwei's avatar
zhangwenwei committed
360
        gt_annos = [info['annos'] for info in self.data_infos]
zhangwenwei's avatar
zhangwenwei committed
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378

        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
379
        else:
zhangwenwei's avatar
zhangwenwei committed
380
381
382
383
384
385
386
387
            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)

388
389
        if tmp_dir is not None:
            tmp_dir.cleanup()
390
391
        if show or out_dir:
            self.show(results, out_dir, show=show, pipeline=pipeline)
392
        return ap_dict
393
394
395
396
397
398

    def bbox2result_kitti(self,
                          net_outputs,
                          class_names,
                          pklfile_prefix=None,
                          submission_prefix=None):
399
400
401
402
        """Convert 3D detection results to kitti format for evaluation and test
        submission.

        Args:
403
            net_outputs (list[np.ndarray]): List of array storing the
404
405
                inferenced bounding boxes and scores.
            class_names (list[String]): A list of class names.
406
407
            pklfile_prefix (str): The prefix of pkl file.
            submission_prefix (str): The prefix of submission file.
408
409
410
411

        Returns:
            list[dict]: A list of dictionaries with the kitti format.
        """
Wenwei Zhang's avatar
Wenwei Zhang committed
412
413
        assert len(net_outputs) == len(self.data_infos), \
            'invalid list length of network outputs'
414
415
        if submission_prefix is not None:
            mmcv.mkdir_or_exist(submission_prefix)
zhangwenwei's avatar
zhangwenwei committed
416
417

        det_annos = []
zhangwenwei's avatar
zhangwenwei committed
418
        print('\nConverting prediction to KITTI format')
zhangwenwei's avatar
zhangwenwei committed
419
420
421
        for idx, pred_dicts in enumerate(
                mmcv.track_iter_progress(net_outputs)):
            annos = []
zhangwenwei's avatar
zhangwenwei committed
422
            info = self.data_infos[idx]
zhangwenwei's avatar
zhangwenwei committed
423
            sample_idx = info['image']['image_idx']
zhangwenwei's avatar
zhangwenwei committed
424
            image_shape = info['image']['image_shape'][:2]
zhangwenwei's avatar
zhangwenwei committed
425
            box_dict = self.convert_valid_bboxes(pred_dicts, info)
xiliu8006's avatar
xiliu8006 committed
426
427
428
429
430
431
432
433
434
435
436
            anno = {
                'name': [],
                'truncated': [],
                'occluded': [],
                'alpha': [],
                'bbox': [],
                'dimensions': [],
                'location': [],
                'rotation_y': [],
                'score': []
            }
zhangwenwei's avatar
zhangwenwei committed
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
            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
463
                anno = {
zhangwenwei's avatar
zhangwenwei committed
464
465
466
467
468
469
470
471
472
                    '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
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
                }
                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
496
497
            annos[-1]['sample_idx'] = np.array(
                [sample_idx] * len(annos[-1]['score']), dtype=np.int64)
zhangwenwei's avatar
zhangwenwei committed
498
499
500

            det_annos += annos

501
502
503
        if pklfile_prefix is not None:
            if not pklfile_prefix.endswith(('.pkl', '.pickle')):
                out = f'{pklfile_prefix}.pkl'
zhangwenwei's avatar
zhangwenwei committed
504
            mmcv.dump(det_annos, out)
Wenwei Zhang's avatar
Wenwei Zhang committed
505
            print(f'Result is saved to {out}.')
zhangwenwei's avatar
zhangwenwei committed
506
507
508
509
510
511

        return det_annos

    def bbox2result_kitti2d(self,
                            net_outputs,
                            class_names,
512
513
                            pklfile_prefix=None,
                            submission_prefix=None):
zhangwenwei's avatar
zhangwenwei committed
514
515
        """Convert 2D detection results to kitti format for evaluation and test
        submission.
zhangwenwei's avatar
zhangwenwei committed
516
517

        Args:
518
            net_outputs (list[np.ndarray]): List of array storing the
519
520
                inferenced bounding boxes and scores.
            class_names (list[String]): A list of class names.
521
522
            pklfile_prefix (str): The prefix of pkl file.
            submission_prefix (str): The prefix of submission file.
zhangwenwei's avatar
zhangwenwei committed
523

524
        Returns:
525
            list[dict]: A list of dictionaries have the kitti format
zhangwenwei's avatar
zhangwenwei committed
526
        """
Wenwei Zhang's avatar
Wenwei Zhang committed
527
528
        assert len(net_outputs) == len(self.data_infos), \
            'invalid list length of network outputs'
zhangwenwei's avatar
zhangwenwei committed
529
        det_annos = []
zhangwenwei's avatar
zhangwenwei committed
530
        print('\nConverting prediction to KITTI format')
zhangwenwei's avatar
zhangwenwei committed
531
532
533
534
535
536
537
538
539
540
541
542
543
        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
544
            sample_idx = self.data_infos[i]['image']['image_idx']
zhangwenwei's avatar
zhangwenwei committed
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

            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

586
587
588
589
590
591
592
593
        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
594
            # save file in submission format
595
596
            mmcv.mkdir_or_exist(submission_prefix)
            print(f'Saving KITTI submission to {submission_prefix}')
zhangwenwei's avatar
zhangwenwei committed
597
            for i, anno in enumerate(det_annos):
zhangwenwei's avatar
zhangwenwei committed
598
                sample_idx = self.data_infos[i]['image']['image_idx']
599
                cur_det_file = f'{submission_prefix}/{sample_idx:06d}.txt'
zhangwenwei's avatar
zhangwenwei committed
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
                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,
                        )
617
            print(f'Result is saved to {submission_prefix}')
zhangwenwei's avatar
zhangwenwei committed
618
619
620
621

        return det_annos

    def convert_valid_bboxes(self, box_dict, info):
622
623
624
625
626
627
628
629
630
631
632
633
634
635
        """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.
636
                - box3d_camera (np.ndarray): 3D bounding boxes in
637
                    camera coordinate.
638
                - box3d_lidar (np.ndarray): 3D bounding boxes in
639
640
641
642
643
                    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
644
        # TODO: refactor this function
645
646
647
        box_preds = box_dict['boxes_3d']
        scores = box_dict['scores_3d']
        labels = box_dict['labels_3d']
zhangwenwei's avatar
zhangwenwei committed
648
        sample_idx = info['image']['image_idx']
649
        box_preds.limit_yaw(offset=0.5, period=np.pi * 2)
zhangwenwei's avatar
zhangwenwei committed
650

651
        if len(box_preds) == 0:
zhangwenwei's avatar
zhangwenwei committed
652
            return dict(
653
654
655
656
657
658
                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
659
660
661
662
663

        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']
664
665
666
667
668
        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
669
        box_corners_in_image = points_cam2img(box_corners, P2)
zhangwenwei's avatar
zhangwenwei committed
670
671
672
673
674
        # 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
675
676
        # check box_preds_camera
        image_shape = box_preds.tensor.new_tensor(img_shape)
twang's avatar
twang committed
677
678
679
        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))
680
681
682
683
        # 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
684
685
686
687
688
        valid_inds = valid_cam_inds & valid_pcd_inds.all(-1)

        if valid_inds.sum() > 0:
            return dict(
                bbox=box_2d_preds[valid_inds, :].numpy(),
689
690
691
692
                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(),
693
                sample_idx=sample_idx)
zhangwenwei's avatar
zhangwenwei committed
694
695
        else:
            return dict(
696
697
698
699
700
                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]),
701
                sample_idx=sample_idx)
liyinhao's avatar
liyinhao committed
702

703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
    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):
723
724
725
        """Results visualization.

        Args:
wangtai's avatar
wangtai committed
726
            results (list[dict]): List of bounding boxes results.
727
            out_dir (str): Output directory of visualization result.
728
729
            show (bool): Whether to visualize the results online.
                Default: False.
730
731
            pipeline (list[dict], optional): raw data loading for showing.
                Default: None.
732
        """
liyinhao's avatar
liyinhao committed
733
        assert out_dir is not None, 'Expect out_dir, got none.'
734
        pipeline = self._get_pipeline(pipeline)
liyinhao's avatar
liyinhao committed
735
        for i, result in enumerate(results):
736
737
            if 'pts_bbox' in result.keys():
                result = result['pts_bbox']
liyinhao's avatar
liyinhao committed
738
739
740
            data_info = self.data_infos[i]
            pts_path = data_info['point_cloud']['velodyne_path']
            file_name = osp.split(pts_path)[-1].split('.')[0]
741
742
743
            points, img_metas, img = self._extract_data(
                i, pipeline, ['points', 'img_metas', 'img'])
            points = points.numpy()
liyinhao's avatar
liyinhao committed
744
            # for now we convert points into depth mode
745
746
            points = Coord3DMode.convert_point(points, Coord3DMode.LIDAR,
                                               Coord3DMode.DEPTH)
747
748
749
            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
750
            pred_bboxes = result['boxes_3d'].tensor.numpy()
751
752
753
754
755
756
            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
757
758
759
760
            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)
761
762
763
764
765
766
767
768
                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,
769
                    img_metas['lidar2img'],
770
771
                    out_dir,
                    file_name,
772
773
                    box_mode='lidar',
                    show=show)