load pointwise laser origins

Author: David Josef Emmerichs

pcdet/datasets/dataset.py

@@ -234,7 +234,7 @@ def collate_batch(batch_list, _unused=False):
                         batch_size_ratio = len(val[0])
                         val = [i for item in val for i in item]
                     ret[key] = np.concatenate(val, axis=0)
-                elif key in ['points', 'voxel_coords']:
+                elif key in ['points', 'origins', 'voxel_coords']:
                     coors = []
                     if isinstance(val[0], list):
                         val =  [i for item in val for i in item]

pcdet/datasets/waymo/waymo_dataset.py

@@ -127,7 +127,7 @@ def load_data_to_shared_memory(self):
             if os.path.exists(f"/dev/shm/{sa_key}"):
                 continue
 
-            points = self.get_lidar(sequence_name, sample_idx)
+            points, _origins = self.get_lidar(sequence_name, sample_idx)
             common_utils.sa_create(f"shm://{sa_key}", points)
 
         dist.barrier()
@@ -193,15 +193,29 @@ def get_infos(self, raw_data_path, save_path, num_workers=multiprocessing.cpu_co
         all_sequences_infos = [item for infos in sequence_infos for item in infos]
         return all_sequences_infos
 
-    def get_lidar(self, sequence_name, sample_idx):
+    def get_lidar(self, sequence_name, sample_idx, load_origins=False):
         lidar_file = self.data_path / sequence_name / ('%04d.npy' % sample_idx)
         point_features = np.load(lidar_file)  # (N, 7): [x, y, z, intensity, elongation, NLZ_flag]
 
         points_all, NLZ_flag = point_features[:, 0:5], point_features[:, 5]
         if not self.dataset_cfg.get('DISABLE_NLZ_FLAG_ON_POINTS', False):
             points_all = points_all[NLZ_flag == -1]
         points_all[:, 3] = np.tanh(points_all[:, 3])
-        return points_all
+
+        if load_origins:
+            sample_info = self.seq_name_to_infos[sequence_name][sample_idx]
+            if 'extrinsics' not in sample_info:
+                raise ValueError('extrinsics not saved to database, use db version >= v0_6_0')
+            origins = [extr[:3, 3] for extr in sample_info['extrinsics']]
+            laser_counts = sample_info['num_points_of_each_lidar']
+            assert sum(laser_counts) == points_all.shape[0], (laser_counts, points_all.shape)
+            assert len(origins) == len(laser_counts), (origins, laser_counts)
+            origins = np.concatenate([np.tile(extr[None, :], (c, 1)) for c, extr in zip(laser_counts, origins)], axis=0)
+            assert origins.shape == points_all[:, :3].shape, (origins.shape, points_all.shape)
+        else:
+            origins = None
+
+        return points_all, origins
 
     @staticmethod
     def transform_prebox_to_current(pred_boxes3d, pose_pre, pose_cur):
@@ -243,7 +257,7 @@ def reorder_rois_for_refining(pred_bboxes):
             ordered_bboxes[bs_idx, :len(pred_bboxes[bs_idx])] = pred_bboxes[bs_idx]
         return ordered_bboxes
 
-    def get_sequence_data(self, info, points, sequence_name, sample_idx, sequence_cfg, load_pred_boxes=False):
+    def get_sequence_data(self, info, points, origins, sequence_name, sample_idx, sequence_cfg, load_pred_boxes=False):
         """
         Args:
             info:
@@ -256,7 +270,7 @@ def get_sequence_data(self, info, points, sequence_name, sample_idx, sequence_cf
 
         def remove_ego_points(points, center_radius=1.0):
             mask = ~((np.abs(points[:, 0]) < center_radius) & (np.abs(points[:, 1]) < center_radius))
-            return points[mask]
+            return points[mask], mask
 
         def load_pred_boxes_from_dict(sequence_name, sample_idx):
             """
@@ -268,6 +282,7 @@ def load_pred_boxes_from_dict(sequence_name, sample_idx):
             load_boxes[:, 7:9] = -0.1 * load_boxes[:, 7:9]  # transfer speed to negtive motion from t to t-1
             return load_boxes
 
+        load_origins = origins is not None
         pose_cur = info['pose'].reshape((4, 4))
         num_pts_cur = points.shape[0]
         sample_idx_pre_list = np.clip(sample_idx + np.arange(sequence_cfg.SAMPLE_OFFSET[0], sequence_cfg.SAMPLE_OFFSET[1]), 0, 0x7FFFFFFF)
@@ -281,6 +296,7 @@ def load_pred_boxes_from_dict(sequence_name, sample_idx):
             points = np.hstack([points, np.zeros((points.shape[0], 1)).astype(points.dtype)])
         points_pre_all = []
         num_points_pre = []
+        origins_pre_all = []
 
         pose_all = [pose_cur]
         pred_boxes_all = []
@@ -292,7 +308,7 @@ def load_pred_boxes_from_dict(sequence_name, sample_idx):
 
         for idx, sample_idx_pre in enumerate(sample_idx_pre_list):
 
-            points_pre = self.get_lidar(sequence_name, sample_idx_pre)
+            points_pre, origins_pre = self.get_lidar(sequence_name, sample_idx_pre, load_origins=load_origins)
             pose_pre = sequence_info[sample_idx_pre]['pose'].reshape((4, 4))
             expand_points_pre = np.concatenate([points_pre[:, :3], np.ones((points_pre.shape[0], 1))], axis=-1)
             points_pre_global = np.dot(expand_points_pre, pose_pre.T)[:, :3]
@@ -306,11 +322,19 @@ def load_pred_boxes_from_dict(sequence_name, sample_idx):
             else:
                 # add timestamp
                 points_pre = np.hstack([points_pre, 0.1 * (sample_idx - sample_idx_pre) * np.ones((points_pre.shape[0], 1)).astype(points_pre.dtype)])  # one frame 0.1s
-            points_pre = remove_ego_points(points_pre, 1.0)
+            points_pre, ego_mask = remove_ego_points(points_pre, 1.0)
             points_pre_all.append(points_pre)
             num_points_pre.append(points_pre.shape[0])
             pose_all.append(pose_pre)
 
+            if load_origins:
+                expand_origins_pre = np.concatenate([origins_pre[:, :3], np.ones((origins_pre.shape[0], 1))], axis=-1)
+                origins_pre_global = np.dot(expand_origins_pre, pose_pre.T)[:, :3]
+                expand_origins_pre_global = np.concatenate([origins_pre_global, np.ones((origins_pre_global.shape[0], 1))], axis=-1)
+                origins_pre = np.dot(expand_origins_pre_global, np.linalg.inv(pose_cur.T))[:, :3]
+                origins_pre = origins_pre[ego_mask]
+                origins_pre_all.append(origins_pre)
+
             if load_pred_boxes:
                 pose_pre = sequence_info[sample_idx_pre]['pose'].reshape((4, 4))
                 pred_boxes = load_pred_boxes_from_dict(sequence_name, sample_idx_pre)
@@ -321,6 +345,11 @@ def load_pred_boxes_from_dict(sequence_name, sample_idx):
         num_points_all = np.array([num_pts_cur] + num_points_pre).astype(np.int32)
         poses = np.concatenate(pose_all, axis=0).astype(np.float32)
 
+        if load_origins:
+            origins = np.concatenate([origins] + origins_pre_all, axis=0).astype(np.float32)
+        else:
+            origins = None
+
         if load_pred_boxes:
             temp_pred_boxes = self.reorder_rois_for_refining(pred_boxes_all)
             pred_boxes = temp_pred_boxes[:, :, 0:9]
@@ -329,7 +358,7 @@ def load_pred_boxes_from_dict(sequence_name, sample_idx):
         else:
             pred_boxes = pred_scores = pred_labels = None
 
-        return points, num_points_all, sample_idx_pre_list, poses, pred_boxes, pred_scores, pred_labels
+        return points, origins, num_points_all, sample_idx_pre_list, poses, pred_boxes, pred_scores, pred_labels
 
     def __len__(self):
         if self._merge_all_iters_to_one_epoch:
@@ -348,15 +377,15 @@ def __getitem__(self, index):
         input_dict = {
             'sample_idx': sample_idx
         }
-        if self.use_shared_memory and index < self.shared_memory_file_limit:
+        if self.use_shared_memory and index < self.shared_memory_file_limit and not self.dataset_cfg.get('USE_ORIGINS', False):
             sa_key = f'{sequence_name}___{sample_idx}'
             points = SharedArray.attach(f"shm://{sa_key}").copy()
         else:
-            points = self.get_lidar(sequence_name, sample_idx)
+            points, origins = self.get_lidar(sequence_name, sample_idx, load_origins=self.dataset_cfg.get('USE_ORIGINS', False))
 
         if self.dataset_cfg.get('SEQUENCE_CONFIG', None) is not None and self.dataset_cfg.SEQUENCE_CONFIG.ENABLED:
-            points, num_points_all, sample_idx_pre_list, poses, pred_boxes, pred_scores, pred_labels = self.get_sequence_data(
-                info, points, sequence_name, sample_idx, self.dataset_cfg.SEQUENCE_CONFIG,
+            points, origins, num_points_all, sample_idx_pre_list, poses, pred_boxes, pred_scores, pred_labels = self.get_sequence_data(
+                info, points, origins, sequence_name, sample_idx, self.dataset_cfg.SEQUENCE_CONFIG,
                 load_pred_boxes=self.dataset_cfg.get('USE_PREDBOX', False)
             )
             input_dict['poses'] = poses
@@ -369,6 +398,7 @@ def __getitem__(self, index):
 
         input_dict.update({
             'points': points,
+            'origins': origins,
             'frame_id': info['frame_id'],
         })
 
@@ -487,11 +517,11 @@ def create_groundtruth_database(self, info_path, save_path, used_classes=None, s
             pc_info = info['point_cloud']
             sequence_name = pc_info['lidar_sequence']
             sample_idx = pc_info['sample_idx']
-            points = self.get_lidar(sequence_name, sample_idx)
+            points, _origins = self.get_lidar(sequence_name, sample_idx)
 
             if use_sequence_data:
-                points, num_points_all, sample_idx_pre_list, _, _, _, _ = self.get_sequence_data(
-                    info, points, sequence_name, sample_idx, self.dataset_cfg.SEQUENCE_CONFIG
+                points, _origins, num_points_all, sample_idx_pre_list, _, _, _, _ = self.get_sequence_data(
+                    info, points, _origins, sequence_name, sample_idx, self.dataset_cfg.SEQUENCE_CONFIG
                 )
 
             annos = info['annos']
@@ -565,11 +595,11 @@ def create_gt_database_of_single_scene(self, info_with_idx, database_save_path=N
         pc_info = info['point_cloud']
         sequence_name = pc_info['lidar_sequence']
         sample_idx = pc_info['sample_idx']
-        points = self.get_lidar(sequence_name, sample_idx)
+        points, _origins = self.get_lidar(sequence_name, sample_idx)
 
         if use_sequence_data:
-            points, num_points_all, sample_idx_pre_list, _, _, _, _ = self.get_sequence_data(
-                info, points, sequence_name, sample_idx, self.dataset_cfg.SEQUENCE_CONFIG
+            points, _origins, num_points_all, sample_idx_pre_list, _, _, _, _ = self.get_sequence_data(
+                info, points, _origins, sequence_name, sample_idx, self.dataset_cfg.SEQUENCE_CONFIG
             )
 
         annos = info['annos']

pcdet/utils/common_utils.py

@@ -35,7 +35,7 @@ def drop_info_with_name(info, name):
 def apply_data_transform(data_dict, transforms):
     assert set(transforms.keys()).issubset({'point', 'box'})
     data_keys = {
-        'point': ['points'],
+        'point': ['points', 'origins'],
         'box': ['gt_boxes', 'roi_boxes']
     }
     for tf_type, tf in transforms.items():