update preprocessor

configs/local/view_generate_config.yaml

@@ -7,16 +7,16 @@ runner:
     name: debug
     root_dir: experiments
   generate:
-    port: 5003
+    port: 5002
-    from: 3000
+    from: 2200
-    to: -1 # -1 means all
+    to: 2300 # -1 means all
     object_dir: /media/hofee/data/data/scaled_object_meshes
     table_model_path: /media/hofee/data/data/others/table.obj
-    output_dir: /media/hofee/repository/new_full_data
+    output_dir: /media/hofee/repository/new_data_with_normal
     binocular_vision: true
     plane_size: 10
     max_views: 512
-    min_views: 64
+    min_views: 128
     random_view_ratio: 0.2
     min_cam_table_included_degree: 20
     max_diag: 0.7

preprocess/preprocessor.py

@@ -1,6 +1,13 @@
 import os
 import json
 import numpy as np
+
+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
 from utils.pts import PtsUtil
@@ -55,8 +62,7 @@ def get_world_points(depth, cam_intrinsic, cam_extrinsic):
     points_camera_world = np.dot(cam_extrinsic, points_camera_aug.T).T[:, :3]
     return points_camera_world
 
-def get_world_normals(normal_image, cam_extrinsic):
+def get_world_normals(normals, cam_extrinsic):
-    normals = normal_image.reshape(-1, 3)
     normals = normals / np.linalg.norm(normals, axis=1, keepdims=True)
     normals_world = np.dot(cam_extrinsic[:3, :3], normals.T).T  
     return normals_world
@@ -75,7 +81,7 @@ def get_scan_points_indices(scan_points, mask, display_table_mask_label, cam_int
     return selected_points_indices
     
 
-def save_scene_data(root, scene):
+def save_scene_data(root, scene, scene_idx=0, scene_total=1):
     
     ''' configuration '''
     target_mask_label = (0, 255, 0, 255)
@@ -88,11 +94,12 @@ def save_scene_data(root, scene):
     ''' scan points '''
     display_table_info = DataLoadUtil.get_display_table_info(root, scene)
     radius = display_table_info["radius"]
-    scan_points = ReconstructionUtil.generate_scan_points(display_table_top=0,display_table_radius=radius)
+    scan_points = np.asarray(ReconstructionUtil.generate_scan_points(display_table_top=0,display_table_radius=radius))
     
     ''' read frame data(depth|mask|normal) '''
     frame_num = DataLoadUtil.get_scene_seq_length(root, scene)
     for frame_id in range(frame_num):
+        print(f"[scene({scene_idx}/{scene_total})|frame({frame_id}/{frame_num})]Processing {scene} frame {frame_id}")
         path = DataLoadUtil.get_path(root, scene, frame_id)
         cam_info = DataLoadUtil.load_cam_info(path, binocular=True)
         depth_L, depth_R = DataLoadUtil.load_depth(
@@ -104,30 +111,36 @@ def save_scene_data(root, scene):
         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_L"], cam_info["cam_extrinsic_L"])
+        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_R"], cam_info["cam_extrinsic_R"])
+        scene_points_R = get_world_points(depth_R, cam_info["cam_intrinsic"], cam_info["cam_to_world_R"])
-        scene_points_L, random_sample_idx_L = PtsUtil.random_downsample_point_cloud(
+        sampled_scene_points_L, random_sample_idx_L = PtsUtil.random_downsample_point_cloud(
                 scene_points_L, random_downsample_N, require_idx=True
             )
-        scene_points_R = PtsUtil.random_downsample_point_cloud(
+        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(
-                scene_points_L, scene_points_R, voxel_size, require_idx=True
+                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_L = mask_L.reshape(-1, 4)
         mask_L = (mask_L == target_mask_label).all(axis=-1)
         mask_overlap = mask_L[random_sample_idx_L][overlap_idx_L]
-        scene_normals_L = get_world_normals(normal_L, cam_info["cam_extrinsic_L"])
+        scene_normals_L = normal_L.reshape(-1, 3)
+        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_normals = scene_normals_L[mask_overlap]
         filtered_target_points, filtered_idx = PtsUtil.filter_points(
-            target_points, target_normals, cam_info["cam_extrinsic_L"], filter_degree, require_idx=True
+            target_points, target_normals, cam_info["cam_to_world"], filter_degree, require_idx=True
             )
         
         ''' train_input_mask '''
@@ -135,7 +148,7 @@ def save_scene_data(root, scene):
         
         
         ''' scan points indices '''
-        scan_points_indices = get_scan_points_indices(scan_points, mask_L, display_table_mask_label, cam_info["cam_intrinsic_L"], cam_info["cam_extrinsic_L"]) 
+        scan_points_indices = get_scan_points_indices(scan_points, mask_img, display_table_mask_label, cam_info["cam_intrinsic"], cam_info["cam_to_world"]) 
         
         save_full_points(root, scene, frame_id, train_input_points)
         save_target_points(root, scene, frame_id, filtered_target_points)
@@ -146,6 +159,19 @@ def save_scene_data(root, scene):
 
 
 if __name__ == "__main__":
-    root = ""
+    #root = "/media/hofee/repository/new_data_with_normal"
-    for scene in os.listdir(root):
+    root = "/media/hofee/repository/test_sample"
-        save_scene_data(root, scene)
+    list_path = "/media/hofee/repository/test_sample/test_sample_list.txt"
+    scene_list = []
+    
+    with open(list_path, "r") as f:
+        for line in f:
+            scene_list.append(line.strip())
+
+    from_idx = 0
+    to_idx = len(scene_list)
+    cnt = 0
+    total = to_idx - from_idx
+    for scene in scene_list[from_idx:to_idx]:
+        save_scene_data(root, scene, cnt, total)
+        cnt+=1

utils/data_load.py

@@ -190,7 +190,7 @@ class DataLoadUtil:
                 mask_path = os.path.join(
                     os.path.dirname(path), "mask", os.path.basename(path) + ".png"
                 )
-            mask_image = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
+            mask_image = cv2.imread(mask_path, cv2.IMREAD_UNCHANGED)
             return mask_image
         
     @staticmethod
@@ -199,23 +199,22 @@ class DataLoadUtil:
             normal_path_L = os.path.join(
                 os.path.dirname(path), "normal", os.path.basename(path) + "_L.png"
             )
-            normal_image_L = cv2.imread(normal_path_L, cv2.IMREAD_UNCHANGED)
+            normal_image_L = cv2.imread(normal_path_L, cv2.IMREAD_COLOR)
             normal_path_R = os.path.join(
                 os.path.dirname(path), "normal", os.path.basename(path) + "_R.png"
             )
-            normal_image_R = cv2.imread(normal_path_R, cv2.IMREAD_UNCHANGED)
+            normal_image_R = cv2.imread(normal_path_R, cv2.IMREAD_COLOR)
-            return normal_image_L, normal_image_R
+            return normal_image_L[:3,:3], normal_image_R[:3,:3]
         else:
             if binocular and left_only:
                 normal_path = os.path.join(
                     os.path.dirname(path), "normal", os.path.basename(path) + "_L.png"
                 )
-                
             else:
                 normal_path = os.path.join(
                     os.path.dirname(path), "normal", os.path.basename(path) + ".png"
                 )
-            normal_image = cv2.imread(normal_path, cv2.IMREAD_UNCHANGED)
+            normal_image = cv2.imread(normal_path, cv2.IMREAD_COLOR)
             return normal_image
 
     @staticmethod

utils/pts.py

@@ -20,6 +20,8 @@ class PtsUtil:
     @staticmethod
     def random_downsample_point_cloud(point_cloud, num_points, require_idx=False):
         if point_cloud.shape[0] == 0:
+            if require_idx:
+                return point_cloud, np.array([])
             return point_cloud
         idx = np.random.choice(len(point_cloud), num_points, replace=True)
         if require_idx:
@@ -60,7 +62,11 @@ class PtsUtil:
         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