[Fix] Fix some bugs in waymo conversion and dataset parsing (#1892)

* add the code of generating cam_sync_labels in waymo dataset

* fix key error in waymo converation

* fix waymo convert bug and refactor parse_ann_info in WaymoDataset

* add image_path key in waymo sweeps

* polish code

* add sweeps in data_prefix in waymo config

* fix some bugs in waymo

* resolve some comments and fix configs of pointpillars

* add 2 TODOs

* fix lint
Co-authored-by: lianqing <lianqing1997@gmail.com>

configs/_base_/datasets/waymoD5-3d-3class.py

@@ -2,13 +2,18 @@
 # D5 in the config name means the whole dataset is divided into 5 folds
 # We only use one fold for efficient experiments
 dataset_type = 'WaymoDataset'
+# data_root = 's3://openmmlab/datasets/detection3d/waymo/kitti_format/'
 data_root = 'data/waymo/kitti_format/'
 file_client_args = dict(backend='disk')
 # Uncomment the following if use ceph or other file clients.
 # See https://mmcv.readthedocs.io/en/latest/api.html#mmcv.fileio.FileClient
 # for more details.
 # file_client_args = dict(
-#     backend='petrel', path_mapping=dict(data='s3://waymo_data/'))
+#     backend='petrel',
+#     path_mapping={
+#         './data/waymo': 's3://openmmlab/datasets/detection3d/waymo',
+#         'data/waymo': 's3://openmmlab/datasets/detection3d/waymo'
+#     })
 
 class_names = ['Car', 'Pedestrian', 'Cyclist']
 metainfo = dict(CLASSES=class_names)
@@ -33,7 +38,7 @@ db_sampler = dict(
 train_pipeline = [
     dict(type='LoadPointsFromFile', coord_type='LIDAR', load_dim=6, use_dim=5),
     dict(type='LoadAnnotations3D', with_bbox_3d=True, with_label_3d=True),
-    dict(type='ObjectSample', db_sampler=db_sampler),
+    # dict(type='ObjectSample', db_sampler=db_sampler),
     dict(
         type='RandomFlip3D',
         sync_2d=False,
@@ -51,7 +56,12 @@ train_pipeline = [
         keys=['points', 'gt_bboxes_3d', 'gt_labels_3d'])
 ]
 test_pipeline = [
-    dict(type='LoadPointsFromFile', coord_type='LIDAR', load_dim=6, use_dim=5),
+    dict(
+        type='LoadPointsFromFile',
+        coord_type='LIDAR',
+        load_dim=6,
+        use_dim=5,
+        file_client_args=file_client_args),
     dict(
         type='MultiScaleFlipAug3D',
         img_scale=(1333, 800),
@@ -98,7 +108,8 @@ train_dataloader = dict(
             # and box_type_3d='Depth' in sunrgbd and scannet dataset.
             box_type_3d='LiDAR',
             # load one frame every five frames
-            load_interval=5)))
+            load_interval=5,
+            file_client_args=file_client_args)))
 val_dataloader = dict(
     batch_size=1,
     num_workers=1,
@@ -114,7 +125,8 @@ val_dataloader = dict(
         modality=input_modality,
         test_mode=True,
         metainfo=metainfo,
-        box_type_3d='LiDAR'))
+        box_type_3d='LiDAR',
+        file_client_args=file_client_args))
 
 test_dataloader = dict(
     batch_size=1,
@@ -131,13 +143,15 @@ test_dataloader = dict(
         modality=input_modality,
         test_mode=True,
         metainfo=metainfo,
-        box_type_3d='LiDAR'))
+        box_type_3d='LiDAR',
+        file_client_args=file_client_args))
 
 val_evaluator = dict(
     type='WaymoMetric',
     ann_file='./data/waymo/kitti_format/waymo_infos_val.pkl',
     waymo_bin_file='./data/waymo/waymo_format/gt.bin',
-    data_root='./data/waymo/waymo_format')
+    data_root='./data/waymo/waymo_format',
+    file_client_args=file_client_args)
 test_evaluator = val_evaluator
 
 vis_backends = [dict(type='LocalVisBackend')]

configs/_base_/datasets/waymoD5-3d-car.py

@@ -86,7 +86,8 @@ train_dataloader = dict(
             type=dataset_type,
             data_root=data_root,
             ann_file='waymo_infos_train.pkl',
-            data_prefix=dict(pts='training/velodyne'),
+            data_prefix=dict(
+                pts='training/velodyne', sweeps='training/velodyne'),
             pipeline=train_pipeline,
             modality=input_modality,
             test_mode=False,
@@ -105,7 +106,7 @@ val_dataloader = dict(
     dataset=dict(
         type=dataset_type,
         data_root=data_root,
-        data_prefix=dict(pts='training/velodyne'),
+        data_prefix=dict(pts='training/velodyne', sweeps='training/velodyne'),
         ann_file='waymo_infos_val.pkl',
         pipeline=eval_pipeline,
         modality=input_modality,
@@ -122,7 +123,7 @@ test_dataloader = dict(
     dataset=dict(
         type=dataset_type,
         data_root=data_root,
-        data_prefix=dict(pts='training/velodyne'),
+        data_prefix=dict(pts='training/velodyne', sweeps='training/velodyne'),
         ann_file='waymo_infos_val.pkl',
         pipeline=eval_pipeline,
         modality=input_modality,

configs/_base_/default_runtime.py

@@ -4,7 +4,7 @@ default_hooks = dict(
     timer=dict(type='IterTimerHook'),
     logger=dict(type='LoggerHook', interval=50),
     param_scheduler=dict(type='ParamSchedulerHook'),
-    checkpoint=dict(type='CheckpointHook', interval=1),
+    checkpoint=dict(type='CheckpointHook', interval=-1),
     sampler_seed=dict(type='DistSamplerSeedHook'),
     visualization=dict(type='Det3DVisualizationHook'))
 

configs/pointpillars/pointpillars_hv_secfpn_sbn-all_16xb2-2x_waymo-3d-3class.py

@@ -6,7 +6,7 @@ _base_ = [
 ]
 
 # data settings
-data = dict(train=dict(dataset=dict(load_interval=1)))
+train_dataloader = dict(dataset=dict(dataset=dict(load_interval=1)))
 # Default setting for scaling LR automatically
 #   - `enable` means enable scaling LR automatically
 #       or not by default.

configs/pointpillars/pointpillars_hv_secfpn_sbn-all_16xb2-2x_waymo-3d-car.py

@@ -6,7 +6,7 @@ _base_ = [
 ]
 
 # data settings
-data = dict(train=dict(dataset=dict(load_interval=1)))
+train_dataloader = dict(dataset=dict(dataset=dict(load_interval=1)))
 
 # model settings
 model = dict(

mmdet3d/datasets/waymo_dataset.py

@@ -130,10 +130,9 @@ class WaymoDataset(KittiDataset):
         ann_info = Det3DDataset.parse_ann_info(self, info)
         if ann_info is None:
             # empty instance
-            anns_results = {}
-            anns_results['gt_bboxes_3d'] = np.zeros((0, 7), dtype=np.float32)
-            anns_results['gt_labels_3d'] = np.zeros(0, dtype=np.int64)
-            return anns_results
+            ann_info = {}
+            ann_info['gt_bboxes_3d'] = np.zeros((0, 7), dtype=np.float32)
+            ann_info['gt_labels_3d'] = np.zeros(0, dtype=np.int64)
 
         ann_info = self._remove_dontcare(ann_info)
         # in kitti, lidar2cam = R0_rect @ Tr_velo_to_cam
@@ -158,12 +157,14 @@ class WaymoDataset(KittiDataset):
                 origin=(0.5, 0.5, 0.5))
 
         else:
+            # in waymo, lidar2cam = R0_rect @ Tr_velo_to_cam
+            # convert gt_bboxes_3d to velodyne coordinates with `lidar2cam`
             lidar2cam = np.array(
                 info['images'][self.default_cam_key]['lidar2cam'])
-
             gt_bboxes_3d = CameraInstance3DBoxes(
                 ann_info['gt_bboxes_3d']).convert_to(self.box_mode_3d,
                                                      np.linalg.inv(lidar2cam))
+        ann_info['gt_bboxes_3d'] = gt_bboxes_3d
 
         anns_results = dict(
             gt_bboxes_3d=gt_bboxes_3d,
@@ -220,7 +221,7 @@ class WaymoDataset(KittiDataset):
 
                 # TODO check if need to modify the sample id
                 # TODO check when will use it except for evaluation.
-                camera_info['sample_id'] = info['sample_id']
+                camera_info['sample_idx'] = info['sample_idx']
 
                 if not self.test_mode:
                     # used in training

mmdet3d/evaluation/metrics/waymo_metric.py

@@ -93,7 +93,7 @@ class WaymoMetric(KittiMetric):
         """Compute the metrics from processed results.
 
         Args:
-            results (list): The processed results of each batch.
+            results (list): The processed results of the whole dataset.
 
         Returns:
             Dict[str, float]: The computed metrics. The keys are the names of
@@ -360,7 +360,7 @@ class WaymoMetric(KittiMetric):
             for cam_idx in range(self.num_cams):
                 box_dict[key].append(box_dict_per_frame[cam_idx][key])
         # merge each elements
-        box_dict['sample_id'] = cam0_info['image_id']
+        box_dict['sample_idx'] = cam0_info['image_id']
         for key in ['bbox', 'box3d_lidar', 'scores', 'label_preds']:
             box_dict[key] = np.concatenate(box_dict[key])
 

mmdet3d/models/detectors/imvoxelnet.py

@@ -9,8 +9,6 @@ from mmdet3d.models.layers.fusion_layers.point_fusion import point_sample
 from mmdet3d.registry import MODELS, TASK_UTILS
 from mmdet3d.structures.det3d_data_sample import SampleList
 from mmdet3d.utils import ConfigType, OptConfigType, OptInstanceList
-from mmdet.models.detectors import BaseDetector
-
 
 
 @MODELS.register_module()

mmdet3d/structures/bbox_3d/cam_box3d.py

@@ -291,13 +291,15 @@ class CameraInstance3DBoxes(BaseInstance3DBoxes):
                 The conversion from ``src`` coordinates to ``dst`` coordinates
                 usually comes along the change of sensors, e.g., from camera
                 to LiDAR. This requires a transformation matrix.
-            correct_yaw (bool): If convert the yaw angle to the target
+            correct_yaw (bool): Whether to convert the yaw angle to the target
                 coordinate. Defaults to False.
         Returns:
             :obj:`BaseInstance3DBoxes`:
                 The converted box of the same type in the ``dst`` mode.
         """
         from .box_3d_mode import Box3DMode
+
+        # TODO: always set correct_yaw=True
         return Box3DMode.convert(
             box=self,
             src=Box3DMode.CAM,

tools/create_data.py

@@ -215,14 +215,14 @@ def waymo_data_prep(root_path,
     info_train_path = osp.join(out_dir, f'{info_prefix}_infos_train.pkl')
     info_val_path = osp.join(out_dir, f'{info_prefix}_infos_val.pkl')
     info_trainval_path = osp.join(out_dir, f'{info_prefix}_infos_trainval.pkl')
-    update_pkl_infos('kitti', out_dir=out_dir, pkl_path=info_train_path)
-    update_pkl_infos('kitti', out_dir=out_dir, pkl_path=info_val_path)
-    update_pkl_infos('kitti', out_dir=out_dir, pkl_path=info_trainval_path)
+    update_pkl_infos('waymo', out_dir=out_dir, pkl_path=info_train_path)
+    update_pkl_infos('waymo', out_dir=out_dir, pkl_path=info_val_path)
+    update_pkl_infos('waymo', out_dir=out_dir, pkl_path=info_trainval_path)
     GTDatabaseCreater(
         'WaymoDataset',
         out_dir,
         info_prefix,
-        f'{out_dir}/{info_prefix}_infos_train.pkl',
+        f'{info_prefix}_infos_train.pkl',
         relative_path=False,
         with_mask=False,
         num_worker=workers).create()

tools/dataset_converters/create_gt_database.py

@@ -196,7 +196,8 @@ def create_groundtruth_database(dataset_class_name,
         file_client_args = dict(backend='disk')
         dataset_cfg.update(
             test_mode=False,
-            split='training',
+            data_prefix=dict(
+                pts='training/velodyne', img='', sweeps='training/velodyne'),
             modality=dict(
                 use_lidar=True,
                 use_depth=False,
@@ -407,7 +408,9 @@ class GTDatabaseCreater:
         image_idx = example['sample_idx']
         points = example['points'].tensor.numpy()
         gt_boxes_3d = annos['gt_bboxes_3d'].tensor.numpy()
-        names = annos['gt_names']
+        names = [
+            self.dataset.metainfo['CLASSES'][i] for i in annos['gt_labels_3d']
+        ]
         group_dict = dict()
         if 'group_ids' in annos:
             group_ids = annos['group_ids']
@@ -510,7 +513,8 @@ class GTDatabaseCreater:
             file_client_args = dict(backend='disk')
             dataset_cfg.update(
                 test_mode=False,
-                split='training',
+                data_prefix=dict(
+                    pts='training/velodyne_reduced', img='training/image_2'),
                 modality=dict(
                     use_lidar=True,
                     use_depth=False,
@@ -534,6 +538,9 @@ class GTDatabaseCreater:
         elif self.dataset_class_name == 'NuScenesDataset':
             dataset_cfg.update(
                 use_valid_flag=True,
+                data_prefix=dict(
+                    pts='samples/LIDAR_TOP', img='',
+                    sweeps='sweeps/LIDAR_TOP'),
                 pipeline=[
                     dict(
                         type='LoadPointsFromFile',
@@ -556,7 +563,10 @@ class GTDatabaseCreater:
             file_client_args = dict(backend='disk')
             dataset_cfg.update(
                 test_mode=False,
-                split='training',
+                data_prefix=dict(
+                    pts='training/velodyne',
+                    img='',
+                    sweeps='training/velodyne'),
                 modality=dict(
                     use_lidar=True,
                     use_depth=False,
@@ -577,8 +587,8 @@ class GTDatabaseCreater:
                         file_client_args=file_client_args)
                 ])
 
-        dataset = build_dataset(dataset_cfg)
-        self.pipeline = dataset.pipeline
+        self.dataset = build_dataset(dataset_cfg)
+        self.pipeline = self.dataset.pipeline
         if self.database_save_path is None:
             self.database_save_path = osp.join(
                 self.data_path, f'{self.info_prefix}_gt_database')
@@ -595,13 +605,15 @@ class GTDatabaseCreater:
                 self.file2id.update({info['file_name']: i})
 
         def loop_dataset(i):
-            input_dict = dataset.get_data_info(i)
-            dataset.pre_pipeline(input_dict)
+            input_dict = self.dataset.get_data_info(i)
+            input_dict['box_type_3d'] = self.dataset.box_type_3d
+            input_dict['box_mode_3d'] = self.dataset.box_mode_3d
             return input_dict
 
         multi_db_infos = mmengine.track_parallel_progress(
             self.create_single,
-            ((loop_dataset(i) for i in range(len(dataset))), len(dataset)),
+            ((loop_dataset(i)
+              for i in range(len(self.dataset))), len(self.dataset)),
             self.num_worker)
         print('Make global unique group id')
         group_counter_offset = 0

tools/dataset_converters/kitti_data_utils.py

@@ -46,8 +46,9 @@ def get_image_path(idx,
                    relative_path=True,
                    exist_check=True,
                    info_type='image_2',
+                   file_tail='.png',
                    use_prefix_id=False):
-    return get_kitti_info_path(idx, prefix, info_type, '.png', training,
+    return get_kitti_info_path(idx, prefix, info_type, file_tail, training,
                                relative_path, exist_check, use_prefix_id)
 
 
@@ -378,6 +379,7 @@ class WaymoInfoGatherer:
             self.training,
             self.relative_path,
             info_type='image_0',
+            file_tail='.jpg',
             use_prefix_id=True)
         if self.with_imageshape:
             img_path = image_info['image_path']
@@ -443,6 +445,7 @@ class WaymoInfoGatherer:
             else:
                 rect_4x4 = R0_rect
 
+            # TODO: naming Tr_velo_to_cam or Tr_velo_to_cam0
             Tr_velo_to_cam = np.array([
                 float(info) for info in lines[6].split(' ')[1:13]
             ]).reshape([3, 4])
@@ -521,6 +524,14 @@ class WaymoInfoGatherer:
                             relative_path=False,
                             use_prefix_id=True)) as f:
                     prev_info['timestamp'] = np.int64(f.read())
+                prev_info['image_path'] = get_image_path(
+                    prev_idx,
+                    self.path,
+                    self.training,
+                    self.relative_path,
+                    info_type='image_0',
+                    file_tail='.jpg',
+                    use_prefix_id=True)
                 prev_pose_path = get_pose_path(
                     prev_idx,
                     self.path,

tools/dataset_converters/nuscenes_converter.py

@@ -177,6 +177,7 @@ def _fill_trainval_infos(nusc,
 
         info = {
             'lidar_path': lidar_path,
+            'num_features': 5,
             'token': sample['token'],
             'sweeps': [],
             'cams': dict(),

tools/dataset_converters/update_infos_to_v2.py

@@ -285,7 +285,8 @@ def update_nuscenes_infos(pkl_path, out_dir):
         temp_data_info['ego2global'] = convert_quaternion_to_matrix(
             ori_info_dict['ego2global_rotation'],
             ori_info_dict['ego2global_translation'])
-        temp_data_info['lidar_points']['num_pts_feats'] = 5
+        temp_data_info['lidar_points']['num_pts_feats'] = ori_info_dict.get(
+            'num_features', 5)
         temp_data_info['lidar_points']['lidar_path'] = ori_info_dict[
             'lidar_path'].split('/')[-1]
         temp_data_info['lidar_points'][
@@ -515,7 +516,7 @@ def update_s3dis_infos(pkl_path, out_dir):
     converted_list = []
     for i, ori_info_dict in enumerate(mmengine.track_iter_progress(data_list)):
         temp_data_info = get_empty_standard_data_info()
-        temp_data_info['sample_id'] = i
+        temp_data_info['sample_idx'] = i
         temp_data_info['lidar_points']['num_pts_feats'] = ori_info_dict[
             'point_cloud']['num_features']
         temp_data_info['lidar_points']['lidar_path'] = ori_info_dict[
@@ -830,7 +831,7 @@ def update_waymo_infos(pkl_path, out_dir):
 
         if 'plane' in ori_info_dict:
             temp_data_info['plane'] = ori_info_dict['plane']
-        temp_data_info['sample_id'] = ori_info_dict['image']['image_idx']
+        temp_data_info['sample_idx'] = ori_info_dict['image']['image_idx']
 
         # calib matrix
         for cam_idx, cam_key in enumerate(camera_types):

tools/dataset_converters/waymo_converter.py

@@ -33,8 +33,10 @@ class Waymo2KITTI(object):
         prefix (str): Prefix of filename. In general, 0 for training, 1 for
             validation and 2 for testing.
         workers (int, optional): Number of workers for the parallel process.
-        test_mode (bool, optional): Whether in the test_mode. Default: False.
-        save_cam_sync_labels (bool, Optional): Whether to save cam sync labels.
+            Defaults to 64.
+        test_mode (bool, optional): Whether in the test_mode.
+            Defaults to False.
+        save_cam_sync_labels (bool, optional): Whether to save cam sync labels.
             Defaults to True.
     """
 
@@ -149,7 +151,7 @@ class Waymo2KITTI(object):
         return len(self.tfrecord_pathnames)
 
     def save_image(self, frame, file_idx, frame_idx):
-        """Parse and save the images in png format.
+        """Parse and save the images in jpg format.
 
         Args:
             frame (:obj:`Frame`): Open dataset frame proto.
@@ -159,7 +161,7 @@ class Waymo2KITTI(object):
         for img in frame.images:
             img_path = f'{self.image_save_dir}{str(img.name - 1)}/' + \
                 f'{self.prefix}{str(file_idx).zfill(3)}' + \
-                f'{str(frame_idx).zfill(3)}.png'
+                f'{str(frame_idx).zfill(3)}.jpg'
             img = mmcv.imfrombytes(img.image)
             mmcv.imwrite(img, img_path)
 
@@ -227,7 +229,7 @@ class Waymo2KITTI(object):
             file_idx (int): Current file index.
             frame_idx (int): Current frame index.
         """
-        range_images, camera_projections, range_image_top_pose = \
+        range_images, camera_projections, seg_labels, range_image_top_pose = \
             parse_range_image_and_camera_projection(frame)
 
         # First return