upd infernce

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"C:\Document\Local Project\nbv_rec\nbv_reconstruction")
 
 from utils.data_load import DataLoadUtil
 from utils.pose import PoseUtil
@@ -55,7 +55,9 @@ class SeqReconstructionDataset(BaseDataset):
 
     def get_datalist(self):
         datalist = []
-        for scene_name in self.scene_name_list:
+        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 = []
@@ -182,17 +184,33 @@ if __name__ == "__main__":
     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": "/data/hofee/data/new_full_data",
+        "root_dir": "H:\\AI\\Datasets\\packed_test_data",
         "source": "seq_reconstruction_dataset",
-        "split_file": "/data/hofee/data/sample.txt",
+        "split_file": "H:\\AI\\Datasets\\data_list\\OmniObject3d_test.txt",
         "load_from_preprocess": True,
-        "ratio": 0.5,
-        "batch_size": 2,
+        "ratio": 1,
         "filter_degree": 75,
         "num_workers": 0,
-        "pts_num": 4096,
-        "type": namespace.Mode.TRAIN,
+        "pts_num": 8192,
+        "type": "test",
     }
     ds = SeqReconstructionDataset(config)
     print(len(ds))

core/seq_dataset_preprocessed.py

@@ -0,0 +1,85 @@
+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"H:\AI\Datasets\packed_test_data"
+        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_coverage_rate": data["first_scanned_coverage_rate"], # List(S): Float, range(0, 1)
+            "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))
+