[Feature]: Support waymo_info_test.pkl and ImageSets generating (#2180)

tools/create_data.py

@@ -199,6 +199,10 @@ def waymo_data_prep(root_path,
             test_mode=(split
                        in ['testing', 'testing_3d_camera_only_detection']))
         converter.convert()
+
+    from tools.dataset_converters.waymo_converter import \
+        create_ImageSets_img_ids
+    create_ImageSets_img_ids(osp.join(out_dir, 'kitti_format'), splits)
     # Generate waymo infos
     out_dir = osp.join(out_dir, 'kitti_format')
     kitti.create_waymo_info_file(
@@ -206,9 +210,11 @@ 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')
+    test_path = osp.join(out_dir, f'{info_prefix}_infos_test.pkl')
     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)
+    update_pkl_infos('waymo', out_dir=out_dir, pkl_path=test_path)
     GTDatabaseCreater(
         'WaymoDataset',
         out_dir,

tools/dataset_converters/kitti_data_utils.py

@@ -365,14 +365,14 @@ class WaymoInfoGatherer:
                 self.training,
                 self.relative_path,
                 use_prefix_id=True)
-            with open(
-                    get_timestamp_path(
-                        idx,
-                        self.path,
-                        self.training,
-                        relative_path=False,
-                        use_prefix_id=True)) as f:
-                info['timestamp'] = np.int64(f.read())
+        with open(
+                get_timestamp_path(
+                    idx,
+                    self.path,
+                    self.training,
+                    relative_path=False,
+                    use_prefix_id=True)) as f:
+            info['timestamp'] = np.int64(f.read())
         image_info['image_path'] = get_image_path(
             idx,
             self.path,

tools/dataset_converters/update_infos_to_v2.py

@@ -919,8 +919,9 @@ def update_waymo_infos(pkl_path, out_dir):
             'point_cloud']['num_features']
         temp_data_info['lidar_points']['timestamp'] = ori_info_dict[
             'timestamp']
-        temp_data_info['lidar_points']['lidar_path'] = Path(
-            ori_info_dict['point_cloud']['velodyne_path']).name
+        velo_path = ori_info_dict['point_cloud'].get('velodyne_path')
+        if velo_path is not None:
+            temp_data_info['lidar_points']['lidar_path'] = Path(velo_path).name
 
         # TODO discuss the usage of Tr_velo_to_cam in lidar
         Trv2c = ori_info_dict['calib']['Tr_velo_to_cam'].astype(np.float32)
@@ -953,93 +954,97 @@ def update_waymo_infos(pkl_path, out_dir):
             temp_data_info['lidar_sweeps'].append(lidar_sweep)
             temp_data_info['image_sweeps'].append(image_sweep)
 
-        anns = ori_info_dict['annos']
-        num_instances = len(anns['name'])
-
+        anns = ori_info_dict.get('annos')
         ignore_class_name = set()
-        instance_list = []
-        for instance_id in range(num_instances):
-            empty_instance = get_empty_instance()
-            empty_instance['bbox'] = anns['bbox'][instance_id].tolist()
+        if anns is not None:
+            num_instances = len(anns['name'])
 
-            if anns['name'][instance_id] in METAINFO['classes']:
-                empty_instance['bbox_label'] = METAINFO['classes'].index(
-                    anns['name'][instance_id])
-            else:
-                ignore_class_name.add(anns['name'][instance_id])
-                empty_instance['bbox_label'] = -1
+            instance_list = []
+            for instance_id in range(num_instances):
+                empty_instance = get_empty_instance()
+                empty_instance['bbox'] = anns['bbox'][instance_id].tolist()
 
-            empty_instance['bbox'] = anns['bbox'][instance_id].tolist()
+                if anns['name'][instance_id] in METAINFO['classes']:
+                    empty_instance['bbox_label'] = METAINFO['classes'].index(
+                        anns['name'][instance_id])
+                else:
+                    ignore_class_name.add(anns['name'][instance_id])
+                    empty_instance['bbox_label'] = -1
 
-            loc = anns['location'][instance_id]
-            dims = anns['dimensions'][instance_id]
-            rots = anns['rotation_y'][:, None][instance_id]
-            gt_bboxes_3d = np.concatenate([loc, dims,
-                                           rots]).astype(np.float32).tolist()
-            empty_instance['bbox_3d'] = gt_bboxes_3d
-            empty_instance['bbox_label_3d'] = copy.deepcopy(
-                empty_instance['bbox_label'])
-            empty_instance['bbox'] = anns['bbox'][instance_id].tolist()
-            empty_instance['truncated'] = int(
-                anns['truncated'][instance_id].tolist())
-            empty_instance['occluded'] = anns['occluded'][instance_id].tolist()
-            empty_instance['alpha'] = anns['alpha'][instance_id].tolist()
-            empty_instance['index'] = anns['index'][instance_id].tolist()
-            empty_instance['group_id'] = anns['group_ids'][instance_id].tolist(
-            )
-            empty_instance['difficulty'] = anns['difficulty'][
-                instance_id].tolist()
-            empty_instance['num_lidar_pts'] = anns['num_points_in_gt'][
-                instance_id].tolist()
-            empty_instance['camera_id'] = anns['camera_id'][
-                instance_id].tolist()
-            empty_instance = clear_instance_unused_keys(empty_instance)
-            instance_list.append(empty_instance)
-        temp_data_info['instances'] = instance_list
+                empty_instance['bbox'] = anns['bbox'][instance_id].tolist()
+
+                loc = anns['location'][instance_id]
+                dims = anns['dimensions'][instance_id]
+                rots = anns['rotation_y'][:, None][instance_id]
+                gt_bboxes_3d = np.concatenate([loc, dims, rots
+                                               ]).astype(np.float32).tolist()
+                empty_instance['bbox_3d'] = gt_bboxes_3d
+                empty_instance['bbox_label_3d'] = copy.deepcopy(
+                    empty_instance['bbox_label'])
+                empty_instance['bbox'] = anns['bbox'][instance_id].tolist()
+                empty_instance['truncated'] = int(
+                    anns['truncated'][instance_id].tolist())
+                empty_instance['occluded'] = anns['occluded'][
+                    instance_id].tolist()
+                empty_instance['alpha'] = anns['alpha'][instance_id].tolist()
+                empty_instance['index'] = anns['index'][instance_id].tolist()
+                empty_instance['group_id'] = anns['group_ids'][
+                    instance_id].tolist()
+                empty_instance['difficulty'] = anns['difficulty'][
+                    instance_id].tolist()
+                empty_instance['num_lidar_pts'] = anns['num_points_in_gt'][
+                    instance_id].tolist()
+                empty_instance['camera_id'] = anns['camera_id'][
+                    instance_id].tolist()
+                empty_instance = clear_instance_unused_keys(empty_instance)
+                instance_list.append(empty_instance)
+            temp_data_info['instances'] = instance_list
 
         # waymo provide the labels that sync with cam
-        anns = ori_info_dict['cam_sync_annos']
-        num_instances = len(anns['name'])
+        anns = ori_info_dict.get('cam_sync_annos')
         ignore_class_name = set()
-        instance_list = []
-        for instance_id in range(num_instances):
-            empty_instance = get_empty_instance()
-            empty_instance['bbox'] = anns['bbox'][instance_id].tolist()
+        if anns is not None:
+            num_instances = len(anns['name'])
+            instance_list = []
+            for instance_id in range(num_instances):
+                empty_instance = get_empty_instance()
+                empty_instance['bbox'] = anns['bbox'][instance_id].tolist()
 
-            if anns['name'][instance_id] in METAINFO['classes']:
-                empty_instance['bbox_label'] = METAINFO['classes'].index(
-                    anns['name'][instance_id])
-            else:
-                ignore_class_name.add(anns['name'][instance_id])
-                empty_instance['bbox_label'] = -1
+                if anns['name'][instance_id] in METAINFO['classes']:
+                    empty_instance['bbox_label'] = METAINFO['classes'].index(
+                        anns['name'][instance_id])
+                else:
+                    ignore_class_name.add(anns['name'][instance_id])
+                    empty_instance['bbox_label'] = -1
 
-            empty_instance['bbox'] = anns['bbox'][instance_id].tolist()
+                empty_instance['bbox'] = anns['bbox'][instance_id].tolist()
 
-            loc = anns['location'][instance_id]
-            dims = anns['dimensions'][instance_id]
-            rots = anns['rotation_y'][:, None][instance_id]
-            gt_bboxes_3d = np.concatenate([loc, dims,
-                                           rots]).astype(np.float32).tolist()
-            empty_instance['bbox_3d'] = gt_bboxes_3d
-            empty_instance['bbox_label_3d'] = copy.deepcopy(
-                empty_instance['bbox_label'])
-            empty_instance['bbox'] = anns['bbox'][instance_id].tolist()
-            empty_instance['truncated'] = int(
-                anns['truncated'][instance_id].tolist())
-            empty_instance['occluded'] = anns['occluded'][instance_id].tolist()
-            empty_instance['alpha'] = anns['alpha'][instance_id].tolist()
-            empty_instance['index'] = anns['index'][instance_id].tolist()
-            empty_instance['group_id'] = anns['group_ids'][instance_id].tolist(
-            )
-            empty_instance['camera_id'] = anns['camera_id'][
-                instance_id].tolist()
-            empty_instance = clear_instance_unused_keys(empty_instance)
-            instance_list.append(empty_instance)
-        temp_data_info['cam_sync_instances'] = instance_list
+                loc = anns['location'][instance_id]
+                dims = anns['dimensions'][instance_id]
+                rots = anns['rotation_y'][:, None][instance_id]
+                gt_bboxes_3d = np.concatenate([loc, dims, rots
+                                               ]).astype(np.float32).tolist()
+                empty_instance['bbox_3d'] = gt_bboxes_3d
+                empty_instance['bbox_label_3d'] = copy.deepcopy(
+                    empty_instance['bbox_label'])
+                empty_instance['bbox'] = anns['bbox'][instance_id].tolist()
+                empty_instance['truncated'] = int(
+                    anns['truncated'][instance_id].tolist())
+                empty_instance['occluded'] = anns['occluded'][
+                    instance_id].tolist()
+                empty_instance['alpha'] = anns['alpha'][instance_id].tolist()
+                empty_instance['index'] = anns['index'][instance_id].tolist()
+                empty_instance['group_id'] = anns['group_ids'][
+                    instance_id].tolist()
+                empty_instance['camera_id'] = anns['camera_id'][
+                    instance_id].tolist()
+                empty_instance = clear_instance_unused_keys(empty_instance)
+                instance_list.append(empty_instance)
+            temp_data_info['cam_sync_instances'] = instance_list
 
-        cam_instances = generate_waymo_camera_instances(
-            ori_info_dict, camera_types)
-        temp_data_info['cam_instances'] = cam_instances
+            cam_instances = generate_waymo_camera_instances(
+                ori_info_dict, camera_types)
+            temp_data_info['cam_instances'] = cam_instances
 
         temp_data_info, _ = clear_data_info_unused_keys(temp_data_info)
         converted_list.append(temp_data_info)

tools/dataset_converters/waymo_converter.py

@@ -9,8 +9,9 @@ except ImportError:
     raise ImportError('Please run "pip install waymo-open-dataset-tf-2-5-0" '
                       '>1.4.5 to install the official devkit first.')
 
+import os
 from glob import glob
-from os.path import join
+from os.path import exists, join
 
 import mmengine
 import numpy as np
@@ -132,9 +133,7 @@ class Waymo2KITTI(object):
 
             self.save_image(frame, file_idx, frame_idx)
             self.save_calib(frame, file_idx, frame_idx)
-            if 'testing_3d_camera_only_detection' not in self.load_dir:
-                # the camera only split doesn't contain lidar points.
-                self.save_lidar(frame, file_idx, frame_idx)
+            self.save_lidar(frame, file_idx, frame_idx)
             self.save_pose(frame, file_idx, frame_idx)
             self.save_timestamp(frame, file_idx, frame_idx)
 
@@ -230,6 +229,9 @@ class Waymo2KITTI(object):
         range_images, camera_projections, seg_labels, range_image_top_pose = \
             parse_range_image_and_camera_projection(frame)
 
+        if range_image_top_pose is None:
+            # the camera only split doesn't contain lidar points.
+            return
         # First return
         points_0, cp_points_0, intensity_0, elongation_0, mask_indices_0 = \
             self.convert_range_image_to_point_cloud(
@@ -601,3 +603,30 @@ class Waymo2KITTI(object):
         else:
             raise ValueError(mat.shape)
         return ret
+
+
+def create_ImageSets_img_ids(root_dir, splits):
+    save_dir = join(root_dir, 'ImageSets/')
+    if not exists(save_dir):
+        os.mkdir(save_dir)
+
+    idx_all = [[] for i in splits]
+    for i, split in enumerate(splits):
+        path = join(root_dir, splits[i], 'calib')
+        if not exists(path):
+            RawNames = []
+        else:
+            RawNames = os.listdir(path)
+
+        for name in RawNames:
+            if name.endswith('.txt'):
+                idx = name.replace('.txt', '\n')
+                idx_all[int(idx[0])].append(idx)
+        idx_all[i].sort()
+
+    open(save_dir + 'train.txt', 'w').writelines(idx_all[0])
+    open(save_dir + 'val.txt', 'w').writelines(idx_all[1])
+    open(save_dir + 'trainval.txt', 'w').writelines(idx_all[0] + idx_all[1])
+    open(save_dir + 'test.txt', 'w').writelines(idx_all[2])
+    # open(save_dir+'test_cam_only.txt','w').writelines(idx_all[3])
+    print('created txt files indicating what to collect in ', splits)