Improve WOD performance with two LiDAR returns, speed up WOD data...

Improve WOD performance with two LiDAR returns, speed up WOD data pre-processing with multiprocesssing.Pool

pcdet/datasets/waymo/waymo_dataset.py

@@ -77,7 +77,6 @@ class WaymoDataset(DatasetTemplate):
         return sequence_file
 
     def get_infos(self, raw_data_path, save_path, num_workers=multiprocessing.cpu_count(), has_label=True, sampled_interval=1):
-        import concurrent.futures as futures
         from functools import partial
         from . import waymo_utils
         print('---------------The waymo sample interval is %d, total sequecnes is %d-----------------'
@@ -92,10 +91,10 @@ class WaymoDataset(DatasetTemplate):
             for sequence_file in self.sample_sequence_list
         ]
 
-        # process_single_sequence(sample_sequence_file_list[0])
-        with futures.ThreadPoolExecutor(num_workers) as executor:
-            sequence_infos = list(tqdm(executor.map(process_single_sequence, sample_sequence_file_list),
+        with multiprocessing.Pool(num_workers) as p:
+            sequence_infos = list(tqdm(p.imap(process_single_sequence, sample_sequence_file_list),
                                        total=len(sample_sequence_file_list)))
+
         all_sequences_infos = [item for infos in sequence_infos for item in infos]
         return all_sequences_infos
 
@@ -251,8 +250,8 @@ class WaymoDataset(DatasetTemplate):
 
     def create_groundtruth_database(self, info_path, save_path, used_classes=None, split='train', sampled_interval=10,
                                     processed_data_tag=None):
-        database_save_path = save_path / ('pcdet_gt_database_%s_sampled_%d' % (split, sampled_interval))
-        db_info_save_path = save_path / ('pcdet_waymo_dbinfos_%s_sampled_%d.pkl' % (split, sampled_interval))
+        database_save_path = save_path / ('%s_gt_database_%s_sampled_%d' % (processed_data_tag, split, sampled_interval))
+        db_info_save_path = save_path / ('%s_waymo_dbinfos_%s_sampled_%d.pkl' % (processed_data_tag, split, sampled_interval))
 
         database_save_path.mkdir(parents=True, exist_ok=True)
         all_db_infos = {}
@@ -274,7 +273,21 @@ class WaymoDataset(DatasetTemplate):
             difficulty = annos['difficulty']
             gt_boxes = annos['gt_boxes_lidar']
 
+            if k % 4 != 0 and len(names) > 0:
+                mask = (names == 'Vehicle')
+                names = names[~mask]
+                difficulty = difficulty[~mask]
+                gt_boxes = gt_boxes[~mask]
+
+            if k % 2 != 0 and len(names) > 0:
+                mask = (names == 'Pedestrian')
+                names = names[~mask]
+                difficulty = difficulty[~mask]
+                gt_boxes = gt_boxes[~mask]
+
             num_obj = gt_boxes.shape[0]
+            if num_obj == 0:
+                continue
 
             box_idxs_of_pts = roiaware_pool3d_utils.points_in_boxes_gpu(
                 torch.from_numpy(points[:, 0:3]).unsqueeze(dim=0).float().cuda(),
@@ -308,15 +321,15 @@ class WaymoDataset(DatasetTemplate):
 
 def create_waymo_infos(dataset_cfg, class_names, data_path, save_path,
                        raw_data_tag='raw_data', processed_data_tag='waymo_processed_data',
-                       workers=multiprocessing.cpu_count()):
+                       workers=min(16, multiprocessing.cpu_count())):
     dataset = WaymoDataset(
         dataset_cfg=dataset_cfg, class_names=class_names, root_path=data_path,
         training=False, logger=common_utils.create_logger()
     )
     train_split, val_split = 'train', 'val'
 
-    train_filename = save_path / ('waymo_infos_%s.pkl' % train_split)
-    val_filename = save_path / ('waymo_infos_%s.pkl' % val_split)
+    train_filename = save_path / ('%s_infos_%s.pkl' % (processed_data_tag, train_split))
+    val_filename = save_path / ('%s_infos_%s.pkl' % (processed_data_tag, val_split))
 
     print('---------------Start to generate data infos---------------')
 
@@ -343,8 +356,8 @@ def create_waymo_infos(dataset_cfg, class_names, data_path, save_path,
     print('---------------Start create groundtruth database for data augmentation---------------')
     dataset.set_split(train_split)
     dataset.create_groundtruth_database(
-        info_path=train_filename, save_path=save_path, split='train', sampled_interval=10,
-        used_classes=['Vehicle', 'Pedestrian', 'Cyclist']
+        info_path=train_filename, save_path=save_path, split='train', sampled_interval=1,
+        used_classes=['Vehicle', 'Pedestrian', 'Cyclist'], processed_data_tag=processed_data_tag
     )
     print('---------------Data preparation Done---------------')
 
@@ -355,6 +368,7 @@ if __name__ == '__main__':
     parser = argparse.ArgumentParser(description='arg parser')
     parser.add_argument('--cfg_file', type=str, default=None, help='specify the config of dataset')
     parser.add_argument('--func', type=str, default='create_waymo_infos', help='')
+    parser.add_argument('--processed_data_tag', type=str, default='waymo_processed_data_v0_3_1', help='')
     args = parser.parse_args()
 
     if args.func == 'create_waymo_infos':
@@ -362,11 +376,12 @@ if __name__ == '__main__':
         from easydict import EasyDict
         dataset_cfg = EasyDict(yaml.load(open(args.cfg_file)))
         ROOT_DIR = (Path(__file__).resolve().parent / '../../../').resolve()
+        dataset_cfg.PROCESSED_DATA_TAG = args.processed_data_tag
         create_waymo_infos(
             dataset_cfg=dataset_cfg,
             class_names=['Vehicle', 'Pedestrian', 'Cyclist'],
             data_path=ROOT_DIR / 'data' / 'waymo',
             save_path=ROOT_DIR / 'data' / 'waymo',
             raw_data_tag='raw_data',
-            processed_data_tag=dataset_cfg.PROCESSED_DATA_TAG
+            processed_data_tag=args.processed_data_tag
         )

pcdet/datasets/waymo/waymo_utils.py

@@ -61,7 +61,7 @@ def generate_labels(frame):
     return annotations
 
 
-def convert_range_image_to_point_cloud(frame, range_images, camera_projections, range_image_top_pose, ri_index=0):
+def convert_range_image_to_point_cloud(frame, range_images, camera_projections, range_image_top_pose, ri_index=(0, 1)):
     """
     Modified from the codes of Waymo Open Dataset.
     Convert range images to point cloud.
@@ -97,61 +97,73 @@ def convert_range_image_to_point_cloud(frame, range_images, camera_projections,
     range_image_top_pose_tensor = transform_utils.get_transform(
         range_image_top_pose_tensor_rotation,
         range_image_top_pose_tensor_translation)
+
     for c in calibrations:
-        range_image = range_images[c.name][ri_index]
-        if len(c.beam_inclinations) == 0:  # pylint: disable=g-explicit-length-test
-            beam_inclinations = range_image_utils.compute_inclination(
-                tf.constant([c.beam_inclination_min, c.beam_inclination_max]),
-                height=range_image.shape.dims[0])
-        else:
-            beam_inclinations = tf.constant(c.beam_inclinations)
-
-        beam_inclinations = tf.reverse(beam_inclinations, axis=[-1])
-        extrinsic = np.reshape(np.array(c.extrinsic.transform), [4, 4])
-
-        range_image_tensor = tf.reshape(
-            tf.convert_to_tensor(range_image.data), range_image.shape.dims)
-        pixel_pose_local = None
-        frame_pose_local = None
-        if c.name == dataset_pb2.LaserName.TOP:
-            pixel_pose_local = range_image_top_pose_tensor
-            pixel_pose_local = tf.expand_dims(pixel_pose_local, axis=0)
-            frame_pose_local = tf.expand_dims(frame_pose, axis=0)
-        range_image_mask = range_image_tensor[..., 0] > 0
-        range_image_NLZ = range_image_tensor[..., 3]
-        range_image_intensity = range_image_tensor[..., 1]
-        range_image_elongation = range_image_tensor[..., 2]
-        range_image_cartesian = range_image_utils.extract_point_cloud_from_range_image(
-            tf.expand_dims(range_image_tensor[..., 0], axis=0),
-            tf.expand_dims(extrinsic, axis=0),
-            tf.expand_dims(tf.convert_to_tensor(beam_inclinations), axis=0),
-            pixel_pose=pixel_pose_local,
-            frame_pose=frame_pose_local)
-
-        range_image_cartesian = tf.squeeze(range_image_cartesian, axis=0)
-        points_tensor = tf.gather_nd(range_image_cartesian,
-                                     tf.where(range_image_mask))
-        points_NLZ_tensor = tf.gather_nd(range_image_NLZ, tf.compat.v1.where(range_image_mask))
-        points_intensity_tensor = tf.gather_nd(range_image_intensity, tf.compat.v1.where(range_image_mask))
-        points_elongation_tensor = tf.gather_nd(range_image_elongation, tf.compat.v1.where(range_image_mask))
-        cp = camera_projections[c.name][0]
-        cp_tensor = tf.reshape(tf.convert_to_tensor(cp.data), cp.shape.dims)
-        cp_points_tensor = tf.gather_nd(cp_tensor, tf.where(range_image_mask))
-        points.append(points_tensor.numpy())
-        cp_points.append(cp_points_tensor.numpy())
-        points_NLZ.append(points_NLZ_tensor.numpy())
-        points_intensity.append(points_intensity_tensor.numpy())
-        points_elongation.append(points_elongation_tensor.numpy())
+        points_single, cp_points_single, points_NLZ_single, points_intensity_single, points_elongation_single \
+            = [], [], [], [], []
+        for cur_ri_index in ri_index:
+            range_image = range_images[c.name][cur_ri_index]
+            if len(c.beam_inclinations) == 0:  # pylint: disable=g-explicit-length-test
+                beam_inclinations = range_image_utils.compute_inclination(
+                    tf.constant([c.beam_inclination_min, c.beam_inclination_max]),
+                    height=range_image.shape.dims[0])
+            else:
+                beam_inclinations = tf.constant(c.beam_inclinations)
+
+            beam_inclinations = tf.reverse(beam_inclinations, axis=[-1])
+            extrinsic = np.reshape(np.array(c.extrinsic.transform), [4, 4])
+
+            range_image_tensor = tf.reshape(
+                tf.convert_to_tensor(range_image.data), range_image.shape.dims)
+            pixel_pose_local = None
+            frame_pose_local = None
+            if c.name == dataset_pb2.LaserName.TOP:
+                pixel_pose_local = range_image_top_pose_tensor
+                pixel_pose_local = tf.expand_dims(pixel_pose_local, axis=0)
+                frame_pose_local = tf.expand_dims(frame_pose, axis=0)
+            range_image_mask = range_image_tensor[..., 0] > 0
+            range_image_NLZ = range_image_tensor[..., 3]
+            range_image_intensity = range_image_tensor[..., 1]
+            range_image_elongation = range_image_tensor[..., 2]
+            range_image_cartesian = range_image_utils.extract_point_cloud_from_range_image(
+                tf.expand_dims(range_image_tensor[..., 0], axis=0),
+                tf.expand_dims(extrinsic, axis=0),
+                tf.expand_dims(tf.convert_to_tensor(beam_inclinations), axis=0),
+                pixel_pose=pixel_pose_local,
+                frame_pose=frame_pose_local)
+
+            range_image_cartesian = tf.squeeze(range_image_cartesian, axis=0)
+            points_tensor = tf.gather_nd(range_image_cartesian,
+                                         tf.where(range_image_mask))
+            points_NLZ_tensor = tf.gather_nd(range_image_NLZ, tf.compat.v1.where(range_image_mask))
+            points_intensity_tensor = tf.gather_nd(range_image_intensity, tf.compat.v1.where(range_image_mask))
+            points_elongation_tensor = tf.gather_nd(range_image_elongation, tf.compat.v1.where(range_image_mask))
+            cp = camera_projections[c.name][0]
+            cp_tensor = tf.reshape(tf.convert_to_tensor(cp.data), cp.shape.dims)
+            cp_points_tensor = tf.gather_nd(cp_tensor, tf.where(range_image_mask))
+
+            points_single.append(points_tensor.numpy())
+            cp_points_single.append(cp_points_tensor.numpy())
+            points_NLZ_single.append(points_NLZ_tensor.numpy())
+            points_intensity_single.append(points_intensity_tensor.numpy())
+            points_elongation_single.append(points_elongation_tensor.numpy())
+
+        points.append(np.concatenate(points_single, axis=0))
+        cp_points.append(np.concatenate(cp_points_single, axis=0))
+        points_NLZ.append(np.concatenate(points_NLZ_single, axis=0))
+        points_intensity.append(np.concatenate(points_intensity_single, axis=0))
+        points_elongation.append(np.concatenate(points_elongation_single, axis=0))
 
     return points, cp_points, points_NLZ, points_intensity, points_elongation
 
 
-def save_lidar_points(frame, cur_save_path):
+def save_lidar_points(frame, cur_save_path, use_two_returns=True):
     range_images, camera_projections, range_image_top_pose = \
         frame_utils.parse_range_image_and_camera_projection(frame)
 
-    points, cp_points, points_in_NLZ_flag, points_intensity, points_elongation = \
-        convert_range_image_to_point_cloud(frame, range_images, camera_projections, range_image_top_pose)
+    points, cp_points, points_in_NLZ_flag, points_intensity, points_elongation = convert_range_image_to_point_cloud(
+        frame, range_images, camera_projections, range_image_top_pose, ri_index=(0, 1) if use_two_returns else (0,)
+    )
 
     # 3d points in vehicle frame.
     points_all = np.concatenate(points, axis=0)
@@ -169,7 +181,7 @@ def save_lidar_points(frame, cur_save_path):
     return num_points_of_each_lidar
 
 
-def process_single_sequence(sequence_file, save_path, sampled_interval, has_label=True):
+def process_single_sequence(sequence_file, save_path, sampled_interval, has_label=True, use_two_returns=True):
     sequence_name = os.path.splitext(os.path.basename(sequence_file))[0]
 
     # print('Load record (sampled_interval=%d): %s' % (sampled_interval, sequence_name))
@@ -214,7 +226,9 @@ def process_single_sequence(sequence_file, save_path, sampled_interval, has_labe
             annotations = generate_labels(frame)
             info['annos'] = annotations
 
-        num_points_of_each_lidar = save_lidar_points(frame, cur_save_dir / ('%04d.npy' % cnt))
+        num_points_of_each_lidar = save_lidar_points(
+            frame, cur_save_dir / ('%04d.npy' % cnt), use_two_returns=use_two_returns
+        )
         info['num_points_of_each_lidar'] = num_points_of_each_lidar
 
         sequence_infos.append(info)

setup.py

@@ -40,7 +40,7 @@ class PostInstallation(install):
 
 
 if __name__ == '__main__':
-    version = '0.3.0+%s' % get_git_commit_number()
+    version = '0.3.1+%s' % get_git_commit_number()
     write_version_to_file(version, 'pcdet/version.py')
 
     setup(

tools/cfgs/dataset_configs/waymo_dataset.yaml

 DATASET: 'WaymoDataset'
 DATA_PATH: '../data/waymo'
-PROCESSED_DATA_TAG: 'waymo_processed_data'
+PROCESSED_DATA_TAG: 'waymo_processed_data_v0_3_1'
 
 POINT_CLOUD_RANGE: [-75.2, -75.2, -2, 75.2, 75.2, 4]
 
@@ -20,7 +20,7 @@ DATA_AUGMENTOR:
         - NAME: gt_sampling
           USE_ROAD_PLANE: False
           DB_INFO_PATH:
-              - pcdet_waymo_dbinfos_train_sampled_10.pkl
+              - waymo_processed_data_v0_3_1_waymo_dbinfos_train_sampled_1.pkl
           PREPARE: {
              filter_by_min_points: ['Vehicle:5', 'Pedestrian:5', 'Cyclist:5'],
              filter_by_difficulty: [-1],
@@ -62,6 +62,6 @@ DATA_PROCESSOR:
       VOXEL_SIZE: [0.1, 0.1, 0.15]
       MAX_POINTS_PER_VOXEL: 5
       MAX_NUMBER_OF_VOXELS: {
-        'train': 80000,
-        'test': 90000
+        'train': 150000,
+        'test': 150000
       }