update basic framework

utils/data_load.py

@@ -0,0 +1,135 @@
+import os
+import OpenEXR
+import Imath
+import numpy as np
+import json
+import cv2
+
+class DataLoadUtil:
+    
+    @staticmethod
+    def get_path(root, scene_idx, frame_idx):
+        path = os.path.join(root, f"sequence.{scene_idx}", f"step{frame_idx}")
+        return path
+    
+    @staticmethod
+    def read_exr_depth(depth_path):
+        file = OpenEXR.InputFile(depth_path)
+
+        dw = file.header()['dataWindow']
+        width = dw.max.x - dw.min.x + 1
+        height = dw.max.y - dw.min.y + 1
+
+        pix_type = Imath.PixelType(Imath.PixelType.FLOAT)
+        depth_map = np.frombuffer(file.channel('R', pix_type), dtype=np.float32)
+
+        depth_map.shape = (height, width)
+        
+        return depth_map
+
+    @staticmethod
+    def load_depth(path):
+        depth_path = path + ".camera.Depth.exr"
+        depth_map = DataLoadUtil.read_exr_depth(depth_path)
+        return depth_map
+    
+    @staticmethod
+    def load_rgb(path):
+        rgb_path = path + ".camera.png"
+        rgb_image = cv2.imread(rgb_path, cv2.IMREAD_COLOR)
+        return rgb_image
+    
+    @staticmethod
+    def load_seg(path):
+        seg_path = path + ".camera.semantic segmentation.png"
+        seg_image = cv2.imread(seg_path, cv2.IMREAD_COLOR)
+        return seg_image
+    
+    @staticmethod
+    def load_cam_info(path):
+        label_path = path + ".camera_params.json"
+        with open(label_path, 'r') as f:
+            label_data = json.load(f)
+        cam_transform = np.asarray(label_data['cam_to_world']).reshape(
+            (4, 4)
+        ).T
+    
+        offset = np.asarray([
+            [1, 0, 0, 0],
+            [0, -1, 0, 0],
+            [0, 0, 1, 0],
+            [0, 0, 0, 1]])  
+
+        cam_to_world = cam_transform @ offset
+        
+        
+        
+        f_x = label_data['f_x']
+        f_y = label_data['f_y']
+        c_x = label_data['c_x']
+        c_y = label_data['c_y']
+        cam_intrinsic = np.array([[f_x, 0, c_x], [0, f_y, c_y], [0, 0, 1]])
+        
+        return {
+            "cam_to_world": cam_to_world,
+            "cam_intrinsic": cam_intrinsic
+        }
+
+        
+    @staticmethod
+    def get_target_point_cloud(depth, cam_intrinsic, cam_extrinsic, mask, target_mask_label=(255,255,255)):
+        h, w = depth.shape
+        i, j = np.meshgrid(np.arange(w), np.arange(h), indexing='xy')
+        
+        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_world = np.dot(cam_extrinsic, points_camera_aug.T).T[:, :3]
+        mask = mask.reshape(-1, 3)
+        target_mask = np.all(mask == target_mask_label, axis=-1)
+        return {
+            "points_world": points_world[target_mask],
+            "points_camera": points_camera[target_mask]
+        }
+        
+    @staticmethod
+    def get_target_point_cloud(depth, cam_intrinsic, cam_extrinsic, mask, target_mask_label=(255,255,255)):
+        h, w = depth.shape
+        i, j = np.meshgrid(np.arange(w), np.arange(h), indexing='xy')
+        
+        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_world = np.dot(cam_extrinsic, points_camera_aug.T).T[:, :3]
+        mask = mask.reshape(-1, 3)
+        target_mask = np.all(mask == target_mask_label, axis=-1)
+        return {
+            "points_world": points_world[target_mask],
+            "points_camera": points_camera[target_mask]
+        }
+        
+    @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']
+    
+    @staticmethod
+    def get_point_cloud_list_from_seq(root, seq_idx, num_frames):
+        point_cloud_list = []
+        for idx in range(num_frames):
+            path = DataLoadUtil.get_path(root, seq_idx, idx)
+            point_cloud = DataLoadUtil.get_point_cloud_world_from_path(path)
+            point_cloud_list.append(point_cloud)
+        return point_cloud_list
+