optimize preprocessor

preprocess/preprocessor.py

@@ -1,12 +1,9 @@
 import os
-import json
 import numpy as np
-
+import time
 import sys
 np.random.seed(0)
-# append parent directory to sys.path
 sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-print(sys.path)
 
 from utils.reconstruction import ReconstructionUtil
 from utils.data_load import DataLoadUtil
@@ -18,26 +15,12 @@ def save_np_pts(path, pts: np.ndarray, file_type="txt"):
     else:
         np.save(path, pts)
 
-def save_full_points(root, scene, frame_idx, full_points: np.ndarray, file_type="txt"):
-    pts_path = os.path.join(root,scene, "scene_pts", f"{frame_idx}.{file_type}")
-    if not os.path.exists(os.path.join(root,scene, "scene_pts")):
-        os.makedirs(os.path.join(root,scene, "scene_pts"))
-    save_np_pts(pts_path, full_points, file_type)
-
 def save_target_points(root, scene, frame_idx, target_points: np.ndarray, file_type="txt"):
-    pts_path = os.path.join(root,scene, "target_pts", f"{frame_idx}.{file_type}")
+    pts_path = os.path.join(root,scene, "pts", f"{frame_idx}.{file_type}")
-    if not os.path.exists(os.path.join(root,scene, "target_pts")):
+    if not os.path.exists(os.path.join(root,scene, "pts")):
-        os.makedirs(os.path.join(root,scene, "target_pts"))
+        os.makedirs(os.path.join(root,scene, "pts"))
     save_np_pts(pts_path, target_points, file_type)
     
-def save_mask_idx(root, scene, frame_idx, mask_idx: np.ndarray,filtered_idx, file_type="txt"):
-    indices_path = os.path.join(root,scene, "mask_idx", f"{frame_idx}.{file_type}")
-    if not os.path.exists(os.path.join(root,scene, "mask_idx")):
-        os.makedirs(os.path.join(root,scene, "mask_idx"))
-    save_np_pts(indices_path, mask_idx, file_type)
-    filtered_path = os.path.join(root,scene, "mask_idx", f"{frame_idx}_filtered.{file_type}")
-    save_np_pts(filtered_path, filtered_idx, file_type)
-    
 def save_scan_points_indices(root, scene, frame_idx, scan_points_indices: np.ndarray, file_type="txt"):
     indices_path = os.path.join(root,scene, "scan_points_indices", f"{frame_idx}.{file_type}")
     if not os.path.exists(os.path.join(root,scene, "scan_points_indices")):
@@ -49,27 +32,36 @@ def save_scan_points(root, scene, scan_points: np.ndarray):
     save_np_pts(scan_points_path, scan_points)
     
 
-def get_world_points(depth, cam_intrinsic, cam_extrinsic):
+def old_get_world_points(depth, cam_intrinsic, cam_extrinsic):
     h, w = depth.shape
     i, j = np.meshgrid(np.arange(w), np.arange(h), indexing="xy")
-
+    # ----- Debug Trace ----- #
+    import ipdb; ipdb.set_trace()
+    # ------------------------ #
     z = depth
     x = (i - cam_intrinsic[0, 2]) * z / cam_intrinsic[0, 0]
     y = (j - cam_intrinsic[1, 2]) * z / cam_intrinsic[1, 1]
+    points_camera = np.stack((x, y, z), axis=-1).reshape(-1, 3)
+    points_camera_aug = np.concatenate((points_camera, np.ones((points_camera.shape[0], 1))), axis=-1)
+    points_camera_world = np.dot(cam_extrinsic, points_camera_aug.T).T[:, :3]
+    
+    return points_camera_world
+
+def get_world_points(depth, mask, cam_intrinsic, cam_extrinsic):
+    z = depth[mask]
+    i, j = np.nonzero(mask)
+    x = (j - cam_intrinsic[0, 2]) * z / cam_intrinsic[0, 0]
+    y = (i - cam_intrinsic[1, 2]) * z / cam_intrinsic[1, 1]
 
     points_camera = np.stack((x, y, z), axis=-1).reshape(-1, 3)
     points_camera_aug = np.concatenate((points_camera, np.ones((points_camera.shape[0], 1))), axis=-1)
     points_camera_world = np.dot(cam_extrinsic, points_camera_aug.T).T[:, :3]
-    return points_camera_world
     
-def get_world_normals(normals, cam_extrinsic):
+    return points_camera_world
-    normals = normals / np.linalg.norm(normals, axis=1, keepdims=True)
-    normals_world = np.dot(cam_extrinsic[:3, :3], normals.T).T  
-    return normals_world
 
 def get_scan_points_indices(scan_points, mask, display_table_mask_label, cam_intrinsic, cam_extrinsic):
     scan_points_homogeneous = np.hstack((scan_points, np.ones((scan_points.shape[0], 1))))
-    points_camera = np.dot(cam_extrinsic, scan_points_homogeneous.T).T[:, :3]
+    points_camera = np.dot(np.linalg.inv(cam_extrinsic), scan_points_homogeneous.T).T[:, :3]
     points_image_homogeneous = np.dot(cam_intrinsic, points_camera.T).T
     points_image_homogeneous /= points_image_homogeneous[:, 2:]
     pixel_x = points_image_homogeneous[:, 0].astype(int)
@@ -77,7 +69,8 @@ def get_scan_points_indices(scan_points, mask, display_table_mask_label, cam_int
     h, w = mask.shape[:2]
     valid_indices = (pixel_x >= 0) & (pixel_x < w) & (pixel_y >= 0) & (pixel_y < h)
     mask_colors = mask[pixel_y[valid_indices], pixel_x[valid_indices]]
-    selected_points_indices = mask_colors == display_table_mask_label
+    selected_points_indices = np.where((mask_colors == display_table_mask_label).all(axis=-1))[0]
+    selected_points_indices = np.where(valid_indices)[0][selected_points_indices]
     return selected_points_indices
     
 
@@ -86,14 +79,16 @@ def save_scene_data(root, scene, scene_idx=0, scene_total=1):
     ''' configuration '''
     target_mask_label = (0, 255, 0, 255)
     display_table_mask_label=(0, 0, 255, 255)
-    random_downsample_N = 65536
+    random_downsample_N = 32768
-    train_input_pts_num = 8192
     voxel_size=0.002
     filter_degree = 75
+    min_z = 0.2
+    max_z = 0.45
     
     ''' scan points '''
     display_table_info = DataLoadUtil.get_display_table_info(root, scene)
     radius = display_table_info["radius"]
+    
     scan_points = np.asarray(ReconstructionUtil.generate_scan_points(display_table_top=0,display_table_radius=radius))
     
     ''' read frame data(depth|mask|normal) '''
@@ -107,52 +102,50 @@ def save_scene_data(root, scene, scene_idx=0, scene_total=1):
                 cam_info["far_plane"], 
                 binocular=True
             )
-        mask_L = DataLoadUtil.load_seg(path, binocular=True, left_only=True)
+        mask_L, mask_R = DataLoadUtil.load_seg(path, binocular=True)
-        normal_L = DataLoadUtil.load_normal(path, binocular=True, left_only=True)
-        
-        ''' scene points '''
-        scene_points_L = get_world_points(depth_L, cam_info["cam_intrinsic"], cam_info["cam_to_world"])
-        scene_points_R = get_world_points(depth_R, cam_info["cam_intrinsic"], cam_info["cam_to_world_R"])
-        sampled_scene_points_L, random_sample_idx_L = PtsUtil.random_downsample_point_cloud(
-                scene_points_L, random_downsample_N, require_idx=True
-            )
-        sampled_scene_points_R = PtsUtil.random_downsample_point_cloud(
-                scene_points_R, random_downsample_N
-            )
-        scene_overlap_points, overlap_idx_L = PtsUtil.get_overlapping_points(
-                sampled_scene_points_L, sampled_scene_points_R, voxel_size, require_idx=True
-            )
-        
-        if scene_overlap_points.shape[0] < 1024:
-            scene_overlap_points = sampled_scene_points_L
-            overlap_idx_L = np.arange(sampled_scene_points_L.shape[0])
-        train_input_points, train_input_idx = PtsUtil.random_downsample_point_cloud(
-                scene_overlap_points, train_input_pts_num, require_idx=True
-            )
         
         ''' target points '''
-        mask_img = mask_L
+        mask_img_L = mask_L
-        mask_L = mask_L.reshape(-1, 4)
+        mask_img_R = mask_R
-        mask_L = (mask_L == target_mask_label).all(axis=-1)
+
-        mask_overlap = mask_L[random_sample_idx_L][overlap_idx_L]
+        target_mask_img_L = (mask_L == target_mask_label).all(axis=-1)
-        scene_normals_L = normal_L.reshape(-1, 3)
+        target_mask_img_R = (mask_R == target_mask_label).all(axis=-1)
-        target_overlap_normals = scene_normals_L[random_sample_idx_L][overlap_idx_L][mask_overlap]
+
-        target_normals = get_world_normals(target_overlap_normals, cam_info["cam_to_world"])
+        
-        target_points = scene_overlap_points[mask_overlap]
+        target_points_L = get_world_points(depth_L, target_mask_img_L, cam_info["cam_intrinsic"], cam_info["cam_to_world"])
-        filtered_target_points, filtered_idx = PtsUtil.filter_points(
+        target_points_R = get_world_points(depth_R, target_mask_img_R, cam_info["cam_intrinsic"], cam_info["cam_to_world_R"])
-            target_points, target_normals, cam_info["cam_to_world"], filter_degree, require_idx=True
+
+        sampled_target_points_L = PtsUtil.random_downsample_point_cloud(
+                target_points_L, random_downsample_N
+            )
+        sampled_target_points_R = PtsUtil.random_downsample_point_cloud(
+                target_points_R, random_downsample_N
             )
 
-        ''' train_input_mask '''
+        has_points = sampled_target_points_L.shape[0] > 0 and sampled_target_points_R.shape[0] > 0
-        mask_train_input = mask_overlap[train_input_idx]
+        if has_points:
+            target_points = PtsUtil.get_overlapping_points(
+                    sampled_target_points_L, sampled_target_points_R, voxel_size
+                )
+
+        
+        has_points = target_points.shape[0] > 0
+        if has_points:
+            points_normals = DataLoadUtil.load_points_normals(root, scene, display_table_as_world_space_origin=True)
+            target_points = PtsUtil.filter_points(
+                target_points, points_normals, cam_info["cam_to_world"],voxel_size=0.002, theta = filter_degree, z_range=(min_z, max_z)
+                )
 
 
         ''' scan points indices '''
-        scan_points_indices = get_scan_points_indices(scan_points, mask_img, display_table_mask_label, cam_info["cam_intrinsic"], cam_info["cam_to_world"]) 
+        scan_points_indices_L = get_scan_points_indices(scan_points, mask_img_L, display_table_mask_label, cam_info["cam_intrinsic"], cam_info["cam_to_world"]) 
+        scan_points_indices_R = get_scan_points_indices(scan_points, mask_img_R, display_table_mask_label, cam_info["cam_intrinsic"], cam_info["cam_to_world_R"])
+        scan_points_indices = np.intersect1d(scan_points_indices_L, scan_points_indices_R)
         
-        save_full_points(root, scene, frame_id, train_input_points)
+        if not has_points:
-        save_target_points(root, scene, frame_id, filtered_target_points)
+            target_points = np.zeros((0, 3))
-        save_mask_idx(root, scene, frame_id, mask_train_input, filtered_idx=filtered_idx)
+            
+        save_target_points(root, scene, frame_id, target_points)
         save_scan_points_indices(root, scene, frame_id, scan_points_indices)
     
     save_scan_points(root, scene, scan_points) # The "done" flag of scene preprocess
@@ -160,8 +153,8 @@ def save_scene_data(root, scene, scene_idx=0, scene_total=1):
 
 if __name__ == "__main__":
     #root = "/media/hofee/repository/new_data_with_normal"
-    root = "/media/hofee/repository/test_sample"
+    root = r"C:\Document\Local Project\nbv_rec\nbv_reconstruction\test\test_sample"
-    list_path = "/media/hofee/repository/test_sample/test_sample_list.txt"
+    list_path = r"C:\Document\Local Project\nbv_rec\nbv_reconstruction\test\test_sample/test_sample_list.txt"
     scene_list = []
     
     with open(list_path, "r") as f:
@@ -171,7 +164,11 @@ if __name__ == "__main__":
     from_idx = 0
     to_idx = len(scene_list)
     cnt = 0
+    import time
     total = to_idx - from_idx
     for scene in scene_list[from_idx:to_idx]:
+        start = time.time()
         save_scene_data(root, scene, cnt, total)
         cnt+=1
+        end = time.time()
+        print(f"Time cost: {end-start}")

utils/data_load.py

@@ -204,7 +204,9 @@ class DataLoadUtil:
                 os.path.dirname(path), "normal", os.path.basename(path) + "_R.png"
             )
             normal_image_R = cv2.imread(normal_path_R, cv2.IMREAD_COLOR)
-            return normal_image_L[:3,:3], normal_image_R[:3,:3]
+            normalized_normal_image_L = normal_image_L / 255.0 * 2.0 - 1.0
+            normalized_normal_image_R = normal_image_R / 255.0 * 2.0 - 1.0
+            return normalized_normal_image_L, normalized_normal_image_R
         else:
             if binocular and left_only:
                 normal_path = os.path.join(
@@ -215,7 +217,8 @@ class DataLoadUtil:
                     os.path.dirname(path), "normal", os.path.basename(path) + ".png"
                 )
             normal_image = cv2.imread(normal_path, cv2.IMREAD_COLOR)
-            return normal_image
+            normalized_normal_image = normal_image / 255.0 * 2.0 - 1.0
+            return normalized_normal_image
 
     @staticmethod
     def load_label(path):

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
+    
+    @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]

utils/reconstruction.py

@@ -132,7 +132,7 @@ class ReconstructionUtil:
 
     
     @staticmethod
-    def generate_scan_points(display_table_top, display_table_radius, min_distance=0.03, max_points_num = 100, max_attempts = 1000):
+    def generate_scan_points(display_table_top, display_table_radius, min_distance=0.03, max_points_num = 500, max_attempts = 1000):
         points = []
         attempts = 0
         while len(points) < max_points_num and attempts < max_attempts: