add split dataset

core/dataset.py

@@ -1,10 +1,9 @@
-import os
 import numpy as np
 from PytorchBoot.dataset import BaseDataset
 import PytorchBoot.stereotype as stereotype
 
 import sys
-sys.path.append(r"/media/hofee/data/project/python/nbv_reconstruction/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
@@ -16,13 +15,22 @@ class NBVReconstructionDataset(BaseDataset):
         super(NBVReconstructionDataset, self).__init__(config)
         self.config = config
         self.root_dir = config["root_dir"]
+        self.split_file_path = config["split_file"]
+        self.scene_name_list = self.load_scene_name_list()
         self.datalist = self.get_datalist()
         self.pts_num = 1024
 
+    def load_scene_name_list(self):
+        scene_name_list = []
+        with open(self.split_file_path, "r") as f:
+            for line in f:
+                scene_name = line.strip()
+                scene_name_list.append(scene_name)
+        return scene_name_list
+    
     def get_datalist(self):
         datalist = []
-        scene_name_list = os.listdir(self.root_dir)
-        for scene_name in scene_name_list:
+        for scene_name in self.scene_name_list:
             label_path = DataLoadUtil.get_label_path(self.root_dir, scene_name)
             label_data = DataLoadUtil.load_label(label_path)
             for data_pair in label_data["data_pairs"]:
@@ -97,8 +105,12 @@ class NBVReconstructionDataset(BaseDataset):
 
 if __name__ == "__main__":
     import torch
+    seed = 0
+    torch.manual_seed(seed)
+    np.random.seed(seed)
     config = {
-        "root_dir": "/media/hofee/data/data/nbv_rec/sample",
+        "root_dir": "C:\\Document\\Local Project\\nbv_rec\\data\\sample",
+        "split_file": "C:\\Document\\Local Project\\nbv_rec\\data\\OmniObject3d_train.txt",
         "ratio": 0.05,
         "batch_size": 1,
         "num_workers": 0,