upd

core/seq_dataset.py

@@ -8,7 +8,7 @@ import torch
 import os
 import sys
 
-sys.path.append(r"/data/hofee/project/nbv_rec/nbv_reconstruction")
+sys.path.append(r"/media/hofee/data/project/python/nbv_reconstruction/nbv_reconstruction")
 
 from utils.data_load import DataLoadUtil
 from utils.pose import PoseUtil
@@ -47,6 +47,8 @@ class SeqReconstructionDataset(BaseDataset):
         with open(self.split_file_path, "r") as f:
             for line in f:
                 scene_name = line.strip()
+                if not os.path.exists(os.path.join(self.root_dir, scene_name)):
+                    continue
                 scene_name_list.append(scene_name)
         return scene_name_list    
     
@@ -55,33 +57,23 @@ class SeqReconstructionDataset(BaseDataset):
 
 
     def get_datalist(self):
-        datalist = []        
+        datalist = []
         total = len(self.scene_name_list)
         for idx, scene_name in enumerate(self.scene_name_list):
             print(f"processing {scene_name} ({idx}/{total})")
-            seq_num = DataLoadUtil.get_label_num(self.root_dir, scene_name)
-            scene_max_coverage_rate = 0
-            max_coverage_rate_list = []
             scene_max_cr_idx = 0
-            for seq_idx in range(seq_num):
-                label_path = DataLoadUtil.get_label_path(
-                    self.root_dir, scene_name, seq_idx
-                )
-                label_data = DataLoadUtil.load_label(label_path)
-                max_coverage_rate = label_data["max_coverage_rate"]
-                if max_coverage_rate > scene_max_coverage_rate:
-                    scene_max_coverage_rate = max_coverage_rate
-                    scene_max_cr_idx = seq_idx
-                max_coverage_rate_list.append(max_coverage_rate)
-            best_label_path = DataLoadUtil.get_label_path(self.root_dir, scene_name, scene_max_cr_idx)
-            best_label_data = DataLoadUtil.load_label(best_label_path)
-            first_frame = best_label_data["best_sequence"][0]
-            best_seq_len = len(best_label_data["best_sequence"])
+            frame_len = DataLoadUtil.get_scene_seq_length(self.root_dir, scene_name)
+
+            for i in range(frame_len):
+                path = DataLoadUtil.get_path(self.root_dir, scene_name, i)
+                pts = DataLoadUtil.load_from_preprocessed_pts(path, "npy")
+                if pts.shape[0] == 0:
+                    continue
             datalist.append({
                 "scene_name": scene_name,
-                "first_frame": first_frame,
-                "best_seq_len": best_seq_len,
-                "max_coverage_rate": scene_max_coverage_rate,
+                "first_frame": i,
+                "best_seq_len": -1,
+                "max_coverage_rate": 1.0,
                 "label_idx": scene_max_cr_idx,
             })  
         return datalist
@@ -132,8 +124,7 @@ class SeqReconstructionDataset(BaseDataset):
             scanned_n_to_world_pose,
         ) = ([], [], [])
         view = data_item_info["first_frame"]
-        frame_idx = view[0]
-        coverage_rate = view[1]
+        frame_idx = view
         view_path = DataLoadUtil.get_path(self.root_dir, scene_name, frame_idx)
         cam_info = DataLoadUtil.load_cam_info(view_path, binocular=True)
         
@@ -145,7 +136,7 @@ class SeqReconstructionDataset(BaseDataset):
             target_point_cloud, self.pts_num
         )
         scanned_views_pts.append(downsampled_target_point_cloud)
-        scanned_coverages_rate.append(coverage_rate)
+        
         n_to_world_6d = PoseUtil.matrix_to_rotation_6d_numpy(
             np.asarray(n_to_world_pose[:3, :3])
         ) 
@@ -162,7 +153,6 @@ class SeqReconstructionDataset(BaseDataset):
         gt_pts = self.seq_combined_pts(scene_name, frame_list)
         data_item = {
             "first_scanned_pts": np.asarray(scanned_views_pts, dtype=np.float32), # Ndarray(S x Nv x 3)
-            "first_scanned_coverage_rate": scanned_coverages_rate, # List(S): Float, range(0, 1)
             "first_scanned_n_to_world_pose_9d": np.asarray(scanned_n_to_world_pose, dtype=np.float32), # Ndarray(S x 9)
             "seq_max_coverage_rate": max_coverage_rate, # Float, range(0, 1)
             "best_seq_len": best_seq_len, # Int
@@ -190,9 +180,9 @@ if __name__ == "__main__":
     np.random.seed(seed)
     
     config = {
-        "root_dir": "/data/hofee/data/new_full_data",
+        "root_dir": "/media/hofee/data/data/new_testset",
         "source": "seq_reconstruction_dataset",
-        "split_file": "/data/hofee/data/new_full_data_list/OmniObject3d_test.txt",
+        "split_file": "/media/hofee/data/data/OmniObject3d_test.txt",
         "load_from_preprocess": True,
         "filter_degree": 75,
         "num_workers": 0,
@@ -200,22 +190,16 @@ if __name__ == "__main__":
         "type": namespace.Mode.TEST,
     }
     
-    output_dir = "/data/hofee/trash_can/output_inference_test"
-    new_output_dir = "/data/hofee/inference_test"
+    output_dir = "/media/hofee/data/data/new_testset_output"
     os.makedirs(output_dir, exist_ok=True)
-    os.makedirs(new_output_dir, exist_ok=True)
 
     ds = SeqReconstructionDataset(config)
     for i in tqdm(range(len(ds)), desc="processing dataset"):
         output_path = os.path.join(output_dir, f"item_{i}.pkl")
-        if os.path.exists(output_path):
-            item = pickle.load(open(output_path, "rb"))
-        else:
-            item = ds.__getitem__(i)
+        item = ds.__getitem__(i)
         for key, value in item.items():
             if isinstance(value, np.ndarray):
                 item[key] = value.tolist()
-        new_output_path = os.path.join(new_output_dir, f"item_{i}.pkl")
-        with open(new_output_path, "wb") as f:
+        #import ipdb; ipdb.set_trace()
+        with open(output_path, "wb") as f:
             pickle.dump(item, f)
-

core/seq_dataset_preprocessed.py

@@ -0,0 +1,82 @@
+import numpy as np
+from PytorchBoot.dataset import BaseDataset
+import PytorchBoot.namespace as namespace
+import PytorchBoot.stereotype as stereotype
+from PytorchBoot.config import ConfigManager
+from PytorchBoot.utils.log_util import Log
+import pickle
+import torch
+import os
+import sys
+
+sys.path.append(r"C:\Document\Local Project\nbv_rec\nbv_reconstruction")
+
+from utils.data_load import DataLoadUtil
+from utils.pose import PoseUtil
+from utils.pts import PtsUtil
+
+@stereotype.dataset("seq_reconstruction_dataset_preprocessed")
+class SeqReconstructionDatasetPreprocessed(BaseDataset):
+    def __init__(self, config):
+        super(SeqReconstructionDatasetPreprocessed, self).__init__(config)
+        self.config = config
+        self.root_dir = config["root_dir"]
+        self.real_root_dir = r"/media/hofee/data/data/new_testset"
+        self.item_list = os.listdir(self.root_dir)
+
+    def __getitem__(self, index):
+        data = pickle.load(open(os.path.join(self.root_dir, self.item_list[index]), "rb"))
+        data_item = {
+            "first_scanned_pts": np.asarray(data["first_scanned_pts"], dtype=np.float32), # Ndarray(S x Nv x 3)
+            "first_scanned_n_to_world_pose_9d": np.asarray(data["first_scanned_n_to_world_pose_9d"], dtype=np.float32), # Ndarray(S x 9)
+            "seq_max_coverage_rate": data["seq_max_coverage_rate"], # Float, range(0, 1)
+            "best_seq_len": data["best_seq_len"], # Int
+            "scene_name": data["scene_name"], # String
+            "gt_pts": np.asarray(data["gt_pts"], dtype=np.float32), # Ndarray(N x 3)
+            "scene_path": os.path.join(self.real_root_dir, data["scene_name"]), # String
+            "O_to_L_pose": np.asarray(data["O_to_L_pose"], dtype=np.float32),
+        }
+        return data_item
+
+    def __len__(self):
+        return len(self.item_list)
+
+# -------------- Debug ---------------- #
+if __name__ == "__main__":
+    import torch
+
+    seed = 0
+    torch.manual_seed(seed)
+    np.random.seed(seed)
+    '''
+    OmniObject3d_test:
+    root_dir: "H:\\AI\\Datasets\\packed_test_data"
+    model_dir: "H:\\AI\\Datasets\\scaled_object_meshes"
+    source: seq_reconstruction_dataset
+    split_file: "H:\\AI\\Datasets\\data_list\\OmniObject3d_test.txt"
+    type: test
+    filter_degree: 75
+    eval_list:
+      - pose_diff
+      - coverage_rate_increase
+    ratio: 0.1
+    batch_size: 1
+    num_workers: 12
+    pts_num: 8192
+    load_from_preprocess: True
+    '''
+    config = {
+        "root_dir": "H:\\AI\\Datasets\\packed_test_data",
+        "source": "seq_reconstruction_dataset",
+        "split_file": "H:\\AI\\Datasets\\data_list\\OmniObject3d_test.txt",
+        "load_from_preprocess": True,
+        "ratio": 1,
+        "filter_degree": 75,
+        "num_workers": 0,
+        "pts_num": 8192,
+        "type": "test",
+    }
+    ds = SeqReconstructionDataset(config)
+    print(len(ds))
+    print(ds.__getitem__(10))
+