optimize preprocessor

utils/pts.py

@@ -1,6 +1,7 @@
 import numpy as np
 import open3d as o3d
 import torch
+from scipy.spatial import cKDTree
 
 class PtsUtil:
 
@@ -56,17 +57,36 @@ class PtsUtil:
         return overlapping_points
     
     @staticmethod
-    def filter_points(points, normals, cam_pose, theta=75, require_idx=False):
+    def new_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)
-        print(cos_theta, theta_rad)
         filtered_points= points[idx]
-        # ------  Debug Start ------
-        import ipdb;ipdb.set_trace()
-        # ------  Debug End ------
         if require_idx:
             return filtered_points, idx
-        return filtered_points
+        return filtered_points
+    
+    @staticmethod
+    def filter_points(points, points_normals, cam_pose,  voxel_size=0.002, theta=45, z_range=(0.2, 0.45)):
+        
+        """ filter with z range """
+        points_cam = PtsUtil.transform_point_cloud(points, np.linalg.inv(cam_pose))
+        idx = (points_cam[:, 2] > z_range[0]) & (points_cam[:, 2] < z_range[1])
+        z_filtered_points = points[idx]
+        
+        """ filter with normal """
+        sampled_points = PtsUtil.voxel_downsample_point_cloud(z_filtered_points, voxel_size)
+        kdtree = cKDTree(points_normals[:,:3])
+        _, indices = kdtree.query(sampled_points)
+        nearest_points = points_normals[indices]
+
+        normals = nearest_points[:, 3:]
+        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_sampled_points= sampled_points[idx] 
+        return filtered_sampled_points[:, :3]