sync

utils/data_load.py

@@ -62,14 +62,46 @@ class DataLoadUtil:
         return pose_mat
     
     @staticmethod
-    def load_depth(path):
-        depth_path = os.path.join(os.path.dirname(path), "depth", os.path.basename(path) + ".png")
-        depth = cv2.imread(depth_path, cv2.IMREAD_UNCHANGED)
-        depth = depth.astype(np.float32) / 65535.0
-        min_depth = 0.01
-        max_depth = 5.0
-        depth_meters = min_depth + (max_depth - min_depth) * depth
-        return depth_meters
+    def load_depth(path, min_depth=0.01,max_depth=5.0,binocular=False):
+        
+        def load_depth_from_real_path(real_path, min_depth, max_depth):
+            depth = cv2.imread(real_path, cv2.IMREAD_UNCHANGED)
+            depth = depth.astype(np.float32) / 65535.0
+            min_depth = min_depth
+            max_depth = max_depth
+            depth_meters = min_depth + (max_depth - min_depth) * depth
+            return depth_meters
+        
+        if binocular:
+            depth_path_L = os.path.join(os.path.dirname(path), "depth", os.path.basename(path) + "_L.png")
+            depth_path_R = os.path.join(os.path.dirname(path), "depth", os.path.basename(path) + "_R.png")
+            depth_meters_L = load_depth_from_real_path(depth_path_L, min_depth, max_depth)
+            depth_meters_R = load_depth_from_real_path(depth_path_R, min_depth, max_depth)
+            return depth_meters_L, depth_meters_R
+        else:
+            depth_path = os.path.join(os.path.dirname(path), "depth", os.path.basename(path) + ".png")
+            depth_meters = load_depth_from_real_path(depth_path, min_depth, max_depth)
+            return depth_meters
+        
+    @staticmethod
+    def load_seg(path, binocular=False):
+        if binocular:
+            def clean_mask(mask_image):
+                green = [0, 255, 0, 255]
+                red = [255, 0, 0, 255]
+                threshold = 2
+                mask_image = np.where(np.abs(mask_image - green) <= threshold, green, mask_image)
+                mask_image = np.where(np.abs(mask_image - red) <= threshold, red, mask_image)
+                return mask_image
+            mask_path_L = os.path.join(os.path.dirname(path), "mask", os.path.basename(path) + "_L.png")
+            mask_image_L = clean_mask(cv2.imread(mask_path_L, cv2.IMREAD_UNCHANGED))
+            mask_path_R = os.path.join(os.path.dirname(path), "mask", os.path.basename(path) + "_R.png")
+            mask_image_R = clean_mask(cv2.imread(mask_path_R, cv2.IMREAD_UNCHANGED))
+            return mask_image_L, mask_image_R
+        else:
+            mask_path = os.path.join(os.path.dirname(path), "mask", os.path.basename(path) + ".png")
+            mask_image = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
+            return mask_image
     
     @staticmethod
     def load_label(path):
@@ -83,12 +115,6 @@ class DataLoadUtil:
         rgb_image = cv2.imread(rgb_path, cv2.IMREAD_COLOR)
         return rgb_image
     
-    @staticmethod
-    def load_seg(path):
-        mask_path = os.path.join(os.path.dirname(path), "mask", os.path.basename(path) + ".png")
-        mask_image = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
-        return mask_image
-    
     @staticmethod
     def cam_pose_transformation(cam_pose_before):
         offset = np.asarray([
@@ -100,20 +126,27 @@ class DataLoadUtil:
         return cam_pose_after
     
     @staticmethod
-    def load_cam_info(path):
+    def load_cam_info(path, binocular=False):
         camera_params_path = os.path.join(os.path.dirname(path), "camera_params", os.path.basename(path) + ".json")
         with open(camera_params_path, 'r') as f:
             label_data = json.load(f)
         cam_to_world = np.asarray(label_data["extrinsic"])
         cam_to_world = DataLoadUtil.cam_pose_transformation(cam_to_world)
         cam_intrinsic = np.asarray(label_data["intrinsic"])
-        return {
+        cam_info = {
             "cam_to_world": cam_to_world,
-            "cam_intrinsic": cam_intrinsic
+            "cam_intrinsic": cam_intrinsic,
+            "far_plane": label_data["far_plane"],
+            "near_plane": label_data["near_plane"]
         }
-
+        if binocular:
+            cam_to_world_R = np.asarray(label_data["extrinsic_R"])
+            cam_to_world_R = DataLoadUtil.cam_pose_transformation(cam_to_world_R)
+            cam_info["cam_to_world_R"] = cam_to_world_R
+        return cam_info
+    
     @staticmethod
-    def get_target_point_cloud(depth, cam_intrinsic, cam_extrinsic, mask, target_mask_label=255):
+    def get_target_point_cloud(depth, cam_intrinsic, cam_extrinsic, mask, target_mask_label=(0,255,0,255)):
         h, w = depth.shape
         i, j = np.meshgrid(np.arange(w), np.arange(h), indexing='xy')
         
@@ -122,9 +155,10 @@ class DataLoadUtil:
         y = (j - cam_intrinsic[1, 2]) * z / cam_intrinsic[1, 1]
         
         points_camera = np.stack((x, y, z), axis=-1).reshape(-1, 3)
-        mask = mask.reshape(-1)
+        mask = mask.reshape(-1,4)
 
-        target_mask = mask == target_mask_label
+        target_mask = (mask == target_mask_label).all(axis=-1)  
+        
         target_points_camera = points_camera[target_mask]
         target_points_camera_aug = np.concatenate([target_points_camera, np.ones((target_points_camera.shape[0], 1))], axis=-1)
         
@@ -133,22 +167,45 @@ class DataLoadUtil:
             "points_world": target_points_world,
             "points_camera": target_points_camera
         }
-    
         
     @staticmethod
-    def get_point_cloud_world_from_path(path):
-        cam_info = DataLoadUtil.load_cam_info(path)
-        depth = DataLoadUtil.load_depth(path)
-        mask = DataLoadUtil.load_seg(path)
-        point_cloud = DataLoadUtil.get_target_point_cloud(depth, cam_info['cam_intrinsic'], cam_info['cam_to_world'], mask)
-        return point_cloud['points_world']
+    def get_target_point_cloud_world_from_path(path, binocular=False):
+        cam_info = DataLoadUtil.load_cam_info(path, binocular=binocular)
+        if binocular:
+            voxel_size = 0.0005
+            depth_L, depth_R = DataLoadUtil.load_depth(path, cam_info['near_plane'], cam_info['far_plane'], binocular=True)
+            mask_L, mask_R = DataLoadUtil.load_seg(path, binocular=True)
+            point_cloud_L = DataLoadUtil.get_target_point_cloud(depth_L, cam_info['cam_intrinsic'], cam_info['cam_to_world'], mask_L)['points_world']
+            point_cloud_R = DataLoadUtil.get_target_point_cloud(depth_R, cam_info['cam_intrinsic'], cam_info['cam_to_world_R'], mask_R)['points_world']
+            overlap_points = DataLoadUtil.get_overlapping_points(point_cloud_L, point_cloud_R, voxel_size)
+            return overlap_points
+        else:
+            depth = DataLoadUtil.load_depth(path, cam_info['near_plane'], cam_info['far_plane'])
+            mask = DataLoadUtil.load_seg(path)
+            point_cloud = DataLoadUtil.get_target_point_cloud(depth, cam_info['cam_intrinsic'], cam_info['cam_to_world'], mask)['points_world']
+            return point_cloud
     
     @staticmethod
-    def get_point_cloud_list_from_seq(root, scene_name, num_frames):
-        point_cloud_list = []
-        for frame_idx in range(num_frames):
-            path = DataLoadUtil.get_path(root, scene_name, frame_idx)
-            point_cloud = DataLoadUtil.get_point_cloud_world_from_path(path)
-            point_cloud_list.append(point_cloud)
-        return point_cloud_list
-    
+    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):
+        voxels_L, indices_L = DataLoadUtil.voxelize_points(point_cloud_L, voxel_size)
+        voxels_R, _ = DataLoadUtil.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]
+        return overlapping_points
+    
+    @staticmethod
+    def load_points_normals(root, scene_name):
+        points_path = os.path.join(root, scene_name, "points_and_normals.txt")
+        points_normals = np.loadtxt(points_path)
+        return points_normals

utils/reconstruction.py

@@ -101,5 +101,20 @@ class ReconstructionUtil:
             runner_name = status_info["runner_name"]
             sm.set_progress(app_name, runner_name, "processed view", len(point_cloud_list), len(point_cloud_list))
         return view_sequence, remaining_views
+    
+    @staticmethod
+    def filter_points(points, points_normals, cam_pose,  voxel_size=0.005, theta=45):
+        sampled_points = PtsUtil.voxel_downsample_point_cloud(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)
+        filtered_sampled_points= sampled_points[cos_theta > np.cos(theta_rad)]
+        
+        return filtered_sampled_points[:, :3]