add data_splitor, modify dataset and encoder

core/dataset.py

@@ -1,6 +1,7 @@
 import numpy as np
 from PytorchBoot.dataset import BaseDataset
 import PytorchBoot.stereotype as stereotype
+from torch.nn.utils.rnn import pad_sequence
 
 import sys
 sys.path.append(r"C:\Document\Local Project\nbv_rec\nbv_reconstruction")
@@ -18,7 +19,7 @@ class NBVReconstructionDataset(BaseDataset):
         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
+        self.pts_num = config["pts_num"]
 
     def load_scene_name_list(self):
         scene_name_list = []
@@ -76,13 +77,9 @@ class NBVReconstructionDataset(BaseDataset):
 
         nbv_idx, nbv_coverage_rate = nbv[0], nbv[1]
         nbv_path = DataLoadUtil.get_path(self.root_dir, scene_name, nbv_idx)
-        nbv_depth = DataLoadUtil.load_depth(nbv_path)
         cam_info = DataLoadUtil.load_cam_info(nbv_path)
-        nbv_mask = DataLoadUtil.load_seg(nbv_path)
         best_frame_to_world = cam_info["cam_to_world"]
         best_to_1_pose = np.dot(np.linalg.inv(first_frame_to_world), best_frame_to_world)
-        best_target_point_cloud = DataLoadUtil.get_target_point_cloud(nbv_depth, cam_info["cam_intrinsic"], best_to_1_pose, nbv_mask)["points_world"]
-        downsampled_best_target_point_cloud = PtsUtil.random_downsample_point_cloud(best_target_point_cloud, self.pts_num)
         best_to_1_6d = PoseUtil.matrix_to_rotation_6d_numpy(np.asarray(best_to_1_pose[:3,:3]))
         best_to_1_trans = best_to_1_pose[:3,3]
         best_to_1_9d = np.concatenate([best_to_1_6d, best_to_1_trans], axis=0)
@@ -91,7 +88,6 @@ class NBVReconstructionDataset(BaseDataset):
             "scanned_pts": np.asarray(scanned_views_pts,dtype=np.float32),
             "scanned_coverage_rate": np.asarray(scanned_coverages_rate,dtype=np.float32),
             "scanned_n_to_1_pose_9d": np.asarray(scanned_n_to_1_pose,dtype=np.float32),
-            "best_pts": np.asarray(downsampled_best_target_point_cloud,dtype=np.float32),
             "best_coverage_rate": nbv_coverage_rate,
             "best_to_1_pose_9d": best_to_1_9d,
             "max_coverage_rate": max_coverage_rate,
@@ -102,6 +98,27 @@ class NBVReconstructionDataset(BaseDataset):
 
     def __len__(self):
         return len(self.datalist)
+    
+    def get_collate_fn(self):
+        def collate_fn(batch):
+            scanned_pts = [item['scanned_pts'] for item in batch]
+            scanned_n_to_1_pose_9d = [item['scanned_n_to_1_pose_9d'] for item in batch]
+            rest = {}
+            for key in batch[0].keys():
+                if key in ['scanned_pts', 'scanned_n_to_1_pose_9d']:
+                    continue
+                if isinstance(batch[0][key], torch.Tensor):
+                    rest[key] = torch.stack([item[key] for item in batch])
+                elif isinstance(batch[0][key], str):
+                    rest[key] = [item[key] for item in batch]
+                else:
+                    rest[key] = [item[key] for item in batch] 
+            return {
+                'scanned_pts': scanned_pts,
+                'scanned_n_to_1_pose_9d': scanned_n_to_1_pose_9d,
+                **rest
+            }
+        return collate_fn
 
 if __name__ == "__main__":
     import torch
@@ -111,9 +128,10 @@ if __name__ == "__main__":
     config = {
         "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,
+        "ratio": 0.5,
+        "batch_size": 2,
         "num_workers": 0,
+        "pts_num": 2048
     }
     ds = NBVReconstructionDataset(config)
     print(len(ds))
@@ -126,11 +144,18 @@ if __name__ == "__main__":
         for pts in data["scanned_pts"][0]:
             #np.savetxt(f"pts_{cnt}.txt", pts)
             cnt+=1
-        best_pts = data["best_pts"][0]
         #np.savetxt("best_pts.txt", best_pts)
         for key, value in data.items():
             if isinstance(value, torch.Tensor):
                 print(key, ":" ,value.shape)
+            else:
+                print(key, ":" ,len(value))
+            if key == "scanned_n_to_1_pose_9d":
+                for val in value:
+                    print(val.shape)
+            if key == "scanned_pts":
+                for val in value:
+                    print(val.shape)
             
                 
         print()

core/pipeline.py

@@ -17,6 +17,9 @@ class NBVReconstructionPipeline(nn.Module):
         
     def forward(self, data):
         mode = data["mode"]
+        # ----- Debug Trace ----- #
+        import ipdb; ipdb.set_trace()
+        # ------------------------ #
         if mode == namespace.Mode.TRAIN:
             return self.forward_train(data)
         elif mode == namespace.Mode.TEST: