finish generate_sequence.py

utils/data_load.py

@@ -4,6 +4,7 @@ import Imath
 import numpy as np
 import json
 import cv2
+import re
 
 class DataLoadUtil:
     
@@ -12,6 +13,38 @@ class DataLoadUtil:
         path = os.path.join(root, f"sequence.{scene_idx}", f"step{frame_idx}")
         return path
     
+    @staticmethod
+    def get_label_path(root, scene_idx):
+        path = os.path.join(root, f"sequence.{scene_idx}_label.json")
+        return path
+    
+    @staticmethod
+    def get_scene_idx_list(root):
+        scene_dir = os.listdir(root)
+        scene_idx_list = []
+        for scene in scene_dir:
+            if "sequence" in scene:
+                scene_idx = int(re.search(r'\d+', scene).group())
+                scene_idx_list.append(scene_idx)
+        return scene_idx_list
+    
+    @staticmethod
+    def get_frame_idx_list(root, scene_idx):
+        scene_path = os.path.join(root, f"sequence.{scene_idx}")
+        view_dir = os.listdir(scene_path)
+        seen_frame_idx = set() 
+        for view in view_dir:
+            if "step" in view:
+                frame_idx = int(re.search(r'\d+', view).group())
+                seen_frame_idx.add(frame_idx)
+        return list(seen_frame_idx)
+    
+    @staticmethod
+    def load_model_points(root,scene_idx):
+        model_path = os.path.join(root, f"sequence.{scene_idx}", "world_points.txt")
+        model_pts = np.loadtxt(model_path)
+        return model_pts
+    
     @staticmethod
     def read_exr_depth(depth_path):
         file = OpenEXR.InputFile(depth_path)