finish generate_sequence.py

configs/generate_config.yaml

@@ -1,23 +1,24 @@
 
-runners:
+runner:
   general:
     seed: 0
     device: cpu
     cuda_visible_devices: "0,1,2,3,4,5,6,7"
     
+  generate:
+    voxel_threshold: 0.005
+    overlap_threshold: 0.3
+    
   experiment:
     name: debug
     root_dir: "experiments"
   
-  generate:
+  dataset_list:
-    - name: OmniObject3d_train
+    - OmniObject3d
-      component: OmniObject3d
-      data_type: train
   
 datasets:
-  general:
-  components:
     OmniObject3d:
-      root_dir: "C:\\Document\\Local Project\\nbv_rec\\output"
+      root_dir: "C:\\Document\\Local Project\\nbv_rec\\sample_dataset"
+      output_dir: "C:\\Document\\Local Project\\nbv_rec\\sample_output"
 
     

core/dataset.py

@@ -0,0 +1,17 @@
+from PytorchBoot.dataset import BaseDataset
+import PytorchBoot.stereotype as stereotype
+
+@stereotype.dataset("nbv_reconstruction_dataset", comment="unfinished")
+class NBVReconstructionDataset(BaseDataset):
+    def __init__(self, config):
+        super(NBVReconstructionDataset, self).__init__(config)
+        self.config = config
+        
+    def get_datalist(self):
+        pass
+    
+    def load_view(path):
+        pass
+    
+    def load_data_item(self, idx):
+        pass

core/evaluation.py

@@ -0,0 +1,20 @@
+import torch
+import PytorchBoot.stereotype as stereotype
+
+@stereotype.evaluation_method("delta_pose_diff", comment="unfinished")
+class DeltaPoseDiff:
+    def __init__(self, config):
+        pass
+
+    def evaluate(self, output_list, data_list):
+        return
+    
+
+@stereotype.evaluation_method("coverage_rate_increase",comment="unfinished")
+class ConverageRateIncrease:
+    def __init__(self, config):
+        pass
+
+    def evaluate(self, output_list, data_list):
+        return
+    

core/pipeline.py

@@ -5,7 +5,7 @@ import PytorchBoot.stereotype as stereotype
 from PytorchBoot.factory.component_factory import ComponentFactory
 from PytorchBoot.utils import Log
 
-@stereotype.pipeline("nbv_reconstruction_pipeline")
+@stereotype.pipeline("nbv_reconstruction_pipeline", comment="should be tested")
 class NBVReconstructionPipeline(nn.Module):
     def __init__(self, config):
         super(NBVReconstructionPipeline, self).__init__()

modules/func_lib/__init__.py

@@ -1,5 +1,4 @@
 from modules.func_lib.samplers import (
-    cond_pc_sampler,
     cond_ode_sampler
 )
 from modules.func_lib.sde import (

runners/strategy_generator.py

@@ -1,16 +1,30 @@
 import os
+import json
 from PytorchBoot.runners.runner import Runner
 from PytorchBoot.config import ConfigManager
 import PytorchBoot.stereotype as stereotype
 
-@stereotype.runner("strategy_generator")
+from utils.data_load import DataLoadUtil
+from utils.reconstruction import ReconstructionUtil
+
+@stereotype.runner("strategy_generator", comment="unfinished")
 class StrategyGenerator(Runner):
     def __init__(self, config):
         super().__init__(config)
         self.load_experiment("generate")
         
     def run(self):
-        self.demo(seq=16,num=100)
+        dataset_name_list =  ConfigManager.get("runner", "dataset_list")
+        voxel_threshold, overlap_threshold = ConfigManager.get("runner","generate","voxel_threshold"), ConfigManager.get("runner","generate","overlap_threshold")
+        for dataset_name in dataset_name_list:
+            root_dir = ConfigManager.get("datasets", dataset_name, "root_dir")
+            output_dir = ConfigManager.get("datasets", dataset_name, "output_dir")
+            if not os.path.exists(output_dir):
+                os.makedirs(output_dir)
+                
+            scene_idx_list = DataLoadUtil.get_scene_idx_list(root_dir)
+            for scene_idx in scene_idx_list:
+                self.generate_sequence(root_dir, output_dir, scene_idx,voxel_threshold, overlap_threshold)
     
     def create_experiment(self, backup_name=None):
         super().create_experiment(backup_name)
@@ -20,54 +34,34 @@ class StrategyGenerator(Runner):
     def load_experiment(self, backup_name=None):
         super().load_experiment(backup_name)
     
-    def demo(self, seq, num=100):
+    def generate_sequence(self,root, output_dir, seq, voxel_threshold, overlap_threshold):
-        import os
+        frame_idx_list = DataLoadUtil.get_frame_idx_list(root, seq)
-        from utils.data_load import DataLoadUtil
+        model_pts = DataLoadUtil.load_model_points(root, seq)
-        from utils.reconstruction import ReconstructionUtil
-        import numpy as np
-        
-        component = self.config["generate"][0]["component"] #r"C:\Document\Local Project\nbv_rec\output"
-        data_dir = ConfigManager.get("datasets", "components", component, "root_dir")
-        model_path = os.path.join(data_dir, f"sequence.{seq}\\world_points.txt")
-        model_pts = np.loadtxt(model_path)
-        
-        output_dir = os.path.join(str(self.experiment_path), "output")
         pts_list = []
-        for idx in range(0,num):
+        for frame_idx in frame_idx_list:
-            path = DataLoadUtil.get_path(data_dir, seq, idx)
+            path = DataLoadUtil.get_path(root, seq, frame_idx)
             point_cloud = DataLoadUtil.get_point_cloud_world_from_path(path)
-            
+            sampled_point_cloud = ReconstructionUtil.downsample_point_cloud(point_cloud, voxel_threshold)
-            sampled_point_cloud = ReconstructionUtil.downsample_point_cloud(point_cloud, 0.005)
             pts_list.append(sampled_point_cloud)
-        
+        limited_useful_view, _ = ReconstructionUtil.compute_next_best_view_sequence_with_overlap(model_pts, pts_list, threshold=voxel_threshold, overlap_threshold=overlap_threshold)
-        sampled_model_pts = ReconstructionUtil.downsample_point_cloud(model_pts, 0.005)
+        data_pairs = self.generate_data_pairs(limited_useful_view)
-        np.savetxt(os.path.join(output_dir,"sampled_model_points.txt"), sampled_model_pts)
+        seq_save_data = {
-        thre = 0.005
+            "data_pairs": data_pairs,
-        
+            "best_sequence": limited_useful_view,
-        useful_view, useless_view = ReconstructionUtil.compute_next_best_view_sequence(model_pts, pts_list, threshold=thre)
+            "max_coverage_rate": limited_useful_view[-1][1]
-        print("useful:", useful_view)
+        }
-        print("useless:", useless_view)
+        output_label_path = DataLoadUtil.get_label_path(output_dir, seq)
-        
+        with open(output_label_path, 'w') as f:
-        selected_full_views = ReconstructionUtil.combine_point_with_view_sequence(pts_list, useful_view)
+            json.dump(seq_save_data, f)
-        downsampled_selected_full_views = ReconstructionUtil.downsample_point_cloud(selected_full_views, thre)
+            
-        np.savetxt(os.path.join(output_dir,"selected_full_views.txt"), downsampled_selected_full_views)
+    def generate_data_pairs(self, useful_view):
-        
+        data_pairs = []
-        limited_useful_view, limited_useless_view = ReconstructionUtil.compute_next_best_view_sequence_with_overlap(model_pts, pts_list, threshold=thre, overlap_threshold=0.3)
+        for next_view_idx in range(len(useful_view)):
-        print("limited_useful:", limited_useful_view)
+            scanned_views = useful_view[:next_view_idx]
-        print("limited_useless:", limited_useless_view)
+            next_view = useful_view[next_view_idx]
-        
+            data_pairs.append((scanned_views, next_view))
-        limited_selected_full_views = ReconstructionUtil.combine_point_with_view_sequence(pts_list, limited_useful_view)
+        return data_pairs
-        downsampled_limited_selected_full_views = ReconstructionUtil.downsample_point_cloud(limited_selected_full_views, thre)
+
-        np.savetxt(os.path.join(output_dir,"selected_full_views_limited.txt"), downsampled_limited_selected_full_views)
-        import json
-        for idx, score in limited_useful_view:
-            path = DataLoadUtil.get_path(data_dir, seq, idx)
-            point_cloud = DataLoadUtil.get_point_cloud_world_from_path(path)
-            print("saving useful view: ", idx, " | score: ", score)
-            np.savetxt(os.path.join(output_dir,f"useful_view_{idx}.txt"), point_cloud)
-        with open(os.path.join(output_dir,f"useful_view.json"), 'w') as f:
-            json.dump(limited_useful_view, f)
-        print("seq length: ", len(useful_view), "limited seq length: ", len(limited_useful_view))
     
         
     

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)