add preprocess

utils/pts.py

@@ -18,13 +18,49 @@ class PtsUtil:
         return points_h[:, :3]
     
     @staticmethod
-    def random_downsample_point_cloud(point_cloud, num_points):
+    def random_downsample_point_cloud(point_cloud, num_points, require_idx=False):
         if point_cloud.shape[0] == 0:
             return point_cloud
         idx = np.random.choice(len(point_cloud), num_points, replace=True)
+        if require_idx:
+            return point_cloud[idx], idx
         return point_cloud[idx]
     
     @staticmethod
     def random_downsample_point_cloud_tensor(point_cloud, num_points):
         idx = torch.randint(0, len(point_cloud), (num_points,))
-        return point_cloud[idx]
+        return point_cloud[idx]
+    
+    @staticmethod
+    def voxelize_points(points, voxel_size):
+        voxel_indices = np.floor(points / voxel_size).astype(np.int32)
+        unique_voxels = np.unique(voxel_indices, axis=0, return_inverse=True)
+        return unique_voxels
+
+    @staticmethod
+    def get_overlapping_points(point_cloud_L, point_cloud_R, voxel_size=0.005, require_idx=False):
+        voxels_L, indices_L = PtsUtil.voxelize_points(point_cloud_L, voxel_size)
+        voxels_R, _ = PtsUtil.voxelize_points(point_cloud_R, voxel_size)
+
+        voxel_indices_L = voxels_L.view([("", voxels_L.dtype)] * 3)
+        voxel_indices_R = voxels_R.view([("", voxels_R.dtype)] * 3)
+        overlapping_voxels = np.intersect1d(voxel_indices_L, voxel_indices_R)
+        mask_L = np.isin(
+            indices_L, np.where(np.isin(voxel_indices_L, overlapping_voxels))[0]
+        )
+        overlapping_points = point_cloud_L[mask_L]
+        if require_idx:
+            return overlapping_points, mask_L
+        return overlapping_points
+    
+    @staticmethod
+    def filter_points(points, normals, cam_pose, theta=75, require_idx=False):
+        camera_axis = -cam_pose[:3, 2] 
+        normals_normalized = normals / np.linalg.norm(normals, axis=1, keepdims=True)
+        cos_theta = np.dot(normals_normalized, camera_axis)
+        theta_rad = np.deg2rad(theta)
+        idx = cos_theta > np.cos(theta_rad)
+        filtered_points= points[idx]
+        if require_idx:
+            return filtered_points, idx
+        return filtered_points