update strategy and overlap rate compute method

utils/data_load.py

@@ -17,11 +17,50 @@ class DataLoadUtil:
         return path
     
     @staticmethod
-    def load_model_points(root, scene_name):
-        model_path = os.path.join(root, scene_name, "sampled_model_points.txt")
+    def get_sampled_model_points_path(root, scene_name):
+        path = os.path.join(root,scene_name, f"sampled_model_points.txt")
+        return path
+    
+    @staticmethod
+    def get_scene_seq_length(root, scene_name):
+        camera_params_path = os.path.join(root, scene_name, "camera_params")
+        return len(os.listdir(camera_params_path))
+    
+    @staticmethod
+    def load_downsampled_world_model_points(root, scene_name):
+        model_path = DataLoadUtil.get_sampled_model_points_path(root, scene_name)
+        model_points = np.loadtxt(model_path)
+        return model_points
+
+    @staticmethod
+    def save_downsampled_world_model_points(root, scene_name, model_points):
+        model_path = DataLoadUtil.get_sampled_model_points_path(root, scene_name)
+        np.savetxt(model_path, model_points)
+    
+    @staticmethod
+    def load_original_model_points(model_dir, object_name):
+        model_path = os.path.join(model_dir, object_name, "mesh.obj")
         mesh = trimesh.load(model_path)
         return mesh.vertices
-
+    
+    @staticmethod
+    def load_scene_info(root, scene_name):
+        scene_info_path = os.path.join(root, scene_name, "scene_info.json")
+        with open(scene_info_path, "r") as f:
+            scene_info = json.load(f)
+        return scene_info
+    
+    @staticmethod
+    def load_target_object_pose(root, scene_name):
+        scene_info = DataLoadUtil.load_scene_info(root, scene_name)
+        target_name = scene_info["target_name"]
+        transformation = scene_info[target_name]
+        location = transformation["location"]
+        rotation_euler = transformation["rotation_euler"]
+        pose_mat = trimesh.transformations.euler_matrix(*rotation_euler)
+        pose_mat[:3, 3] = location
+        return pose_mat
+    
     @staticmethod
     def load_depth(path):
         depth_path = os.path.join(os.path.dirname(path), "depth", os.path.basename(path) + ".png")
@@ -83,8 +122,9 @@ 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, 3)
-        target_mask = np.all(mask == target_mask_label)
+        mask = mask.reshape(-1)
+
+        target_mask = mask == target_mask_label
         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)
         
@@ -104,10 +144,10 @@ class DataLoadUtil:
         return point_cloud['points_world']
     
     @staticmethod
-    def get_point_cloud_list_from_seq(root, seq_idx, num_frames):
+    def get_point_cloud_list_from_seq(root, scene_name, num_frames):
         point_cloud_list = []
-        for idx in range(num_frames):
-            path = DataLoadUtil.get_path(root, seq_idx, idx)
+        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