upd

app_heuristic.py

@@ -1,8 +0,0 @@
-from PytorchBoot.application import PytorchBootApplication
-from runners.heuristic import Heuristic
-
-@PytorchBootApplication("exp_heuristic")
-class ExpHeuristic:
-    @staticmethod
-    def start():
-        Heuristic("configs/local/heuristic_exp_config.yaml").run()

configs/local/heuristic_exp_config.yaml

@@ -1,71 +0,0 @@
-
-runner:
-  general:
-    seed: 0
-    device: cuda
-    cuda_visible_devices: "0,1,2,3,4,5,6,7"
-    
-  experiment:
-    name: exp_hemisphere_circle_trajectory
-    root_dir: "experiments"
-    epoch: -1 # -1 stands for last epoch
-
-  test:
-    dataset_list:
-      - OmniObject3d_test
-
-  blender_script_path: "/media/hofee/data/project/python/nbv_reconstruction/blender/data_renderer.py"
-  output_dir: "/media/hofee/data/results/nbv_rec_inference/hemisphere_random_241202"
-  voxel_size: 0.003
-  min_new_area: 1.0
-  heuristic_method: hemisphere_random
-
-dataset:
-  # OmniObject3d_train:
-  #   root_dir: "C:\\Document\\Datasets\\inference_test1"
-  #   model_dir: "C:\\Document\\Datasets\\scaled_object_meshes"
-  #   source: seq_reconstruction_dataset_preprocessed
-  #   split_file: "C:\\Document\\Datasets\\data_list\\sample.txt"
-  #   type: test
-  #   filter_degree: 75
-  #   ratio: 1
-  #   batch_size: 1
-  #   num_workers: 12
-  #   pts_num: 8192
-  #   load_from_preprocess: True
-
-  OmniObject3d_test:
-    root_dir: "/media/hofee/data/data/new_testset_output"
-    model_dir: "/media/hofee/data/data/scaled_object_meshes"
-    source: seq_reconstruction_dataset_preprocessed
-    # split_file: "C:\\Document\\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
-
-heuristic_methods:
-  hemisphere_random:
-    center: [0, 0, 0]
-    radius_fixed: True
-    fixed_radius: 0.6
-    min_radius: 0.4
-    max_radius: 0.8
-
-  hemisphere_circle_trajectory:
-    center: [0, 0, 0]
-    radius_fixed: False
-    fixed_radius: 0.6
-    min_radius: 0.4
-    max_radius: 0.8
-    phi_list: [15, 45, 75]
-    circle_times: 12
-
-
-

configs/local/inference_config.yaml

@@ -6,7 +6,7 @@ runner:
     cuda_visible_devices: "0,1,2,3,4,5,6,7"
     
   experiment:
-    name: train_ab_partial
+    name: train_ab_global_only
     root_dir: "experiments"
     epoch: -1 # -1 stands for last epoch
 
@@ -15,10 +15,10 @@ runner:
       - OmniObject3d_test
 
   blender_script_path: "/media/hofee/data/project/python/nbv_reconstruction/blender/data_renderer.py"
-  output_dir: "/media/hofee/data/results/nbv_rec_inference/partial_241202"
+  output_dir: "/media/hofee/data/data/new_inference_test_output"
   pipeline: nbv_reconstruction_pipeline
   voxel_size: 0.003
-  min_new_area: 1.0
+  
 dataset:
   # OmniObject3d_train:
   #   root_dir: "C:\\Document\\Datasets\\inference_test1"
@@ -66,7 +66,7 @@ module:
     global_feat: True
     feature_transform: False
   transformer_seq_encoder:
-    embed_dim: 320
+    embed_dim: 256
     num_heads: 4
     ffn_dim: 256
     num_layers: 3

configs/local/view_generate_config.yaml

@@ -7,17 +7,19 @@ runner:
     name: debug
     root_dir: experiments
   generate:
-    port: 5002
-    from: 1
-    to: 50 # -1 means all
-    object_dir: C:\\Document\\Datasets\\scaled_object_meshes
-    table_model_path: C:\\Document\\Datasets\\table.obj
-    output_dir: C:\\Document\\Datasets\\debug_generate_view
+    port: 5000
+    from: 0
+    to: -1 # -1 means all
+    object_dir: /media/hofee/data/data/scaled_object_meshes
+    table_model_path: "/media/hofee/data/data/others/table.obj"
+    output_dir: /media/hofee/data/data/new_testset
+    object_list_path: /media/hofee/data/data/OmniObject3d_test.txt
+    use_list: True
     binocular_vision: true
     plane_size: 10
     max_views: 512
     min_views: 128
-    random_view_ratio: 0.02
+    random_view_ratio: 0.01
     min_cam_table_included_degree: 20
     max_diag: 0.7
     min_diag: 0.01

core/pipeline.py

@@ -88,49 +88,26 @@ class NBVReconstructionPipeline(nn.Module):
         scanned_n_to_world_pose_9d_batch = data[
             "scanned_n_to_world_pose_9d"
         ]  # List(B): Tensor(S x 9)
-        scanned_pts_mask_batch = data["scanned_pts_mask"] # List(B): Tensor(S x N)
 
         device = next(self.parameters()).device
 
         embedding_list_batch = []
 
         combined_scanned_pts_batch = data["combined_scanned_pts"]  # Tensor(B x N x 3)
-        global_scanned_feat, per_point_feat_batch = self.pts_encoder.encode_points(
-            combined_scanned_pts_batch, require_per_point_feat=True
+        global_scanned_feat = self.pts_encoder.encode_points(
+            combined_scanned_pts_batch, require_per_point_feat=False
         )  # global_scanned_feat: Tensor(B x Dg)
-        batch_size = len(scanned_n_to_world_pose_9d_batch)
-        for i in range(batch_size):
-            seq_len = len(scanned_n_to_world_pose_9d_batch[i])
-            scanned_n_to_world_pose_9d = scanned_n_to_world_pose_9d_batch[i].to(device)  # Tensor(S x 9)
-            scanned_pts_mask = scanned_pts_mask_batch[i] # Tensor(S x N)
-            per_point_feat = per_point_feat_batch[i] # Tensor(N x Dp)
-            partial_point_feat_seq = []
-            for j in range(seq_len):
-                partial_per_point_feat = per_point_feat[scanned_pts_mask[j]]
-                if partial_per_point_feat.shape[0] == 0:
-                    partial_point_feat = torch.zeros(per_point_feat.shape[1], device=device)
-                else:
-                    partial_point_feat = torch.mean(partial_per_point_feat, dim=0) # Tensor(Dp)
-                partial_point_feat_seq.append(partial_point_feat)
-            partial_point_feat_seq = torch.stack(partial_point_feat_seq, dim=0) # Tensor(S x Dp)
 
+        for scanned_n_to_world_pose_9d in scanned_n_to_world_pose_9d_batch:
+            scanned_n_to_world_pose_9d = scanned_n_to_world_pose_9d.to(device)  # Tensor(S x 9)
             pose_feat_seq = self.pose_encoder.encode_pose(scanned_n_to_world_pose_9d)  # Tensor(S x Dp) 
-
-            seq_embedding = torch.cat([partial_point_feat_seq, pose_feat_seq], dim=-1)
-            
+            seq_embedding = pose_feat_seq
             embedding_list_batch.append(seq_embedding) # List(B): Tensor(S x (Dp))
         
         seq_feat = self.seq_encoder.encode_sequence(embedding_list_batch) # Tensor(B x Ds)
         main_feat = torch.cat([seq_feat, global_scanned_feat], dim=-1) # Tensor(B x (Ds+Dg))
 
         if torch.isnan(main_feat).any():
-            for i in range(len(main_feat)):
-                if torch.isnan(main_feat[i]).any():
-                    scanned_pts_mask = scanned_pts_mask_batch[i]
-                    Log.info(f"scanned_pts_mask shape: {scanned_pts_mask.shape}")
-                    Log.info(f"scanned_pts_mask sum: {scanned_pts_mask.sum()}")
-                    import ipdb
-                    ipdb.set_trace()
             Log.error("nan in main_feat", True)
 
         return main_feat

core/seq_dataset.py

@@ -47,8 +47,7 @@ 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
+                if os.path.exists(os.path.join(self.root_dir, scene_name)):
                     scene_name_list.append(scene_name)
         return scene_name_list
     
@@ -169,6 +168,7 @@ class SeqReconstructionDataset(BaseDataset):
 
 # -------------- Debug ---------------- #
 if __name__ == "__main__":
+    #import ipdb; ipdb.set_trace()
     import torch
     from tqdm import tqdm
     import pickle
@@ -199,6 +199,6 @@ if __name__ == "__main__":
         for key, value in item.items():
             if isinstance(value, np.ndarray):
                 item[key] = value.tolist()
-        #import ipdb; ipdb.set_trace()
+        import ipdb; ipdb.set_trace()
         with open(output_path, "wb") as f:
             pickle.dump(item, f)

core/seq_dataset_preprocessed.py

@@ -15,6 +15,7 @@ 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):
@@ -41,6 +42,7 @@ class SeqReconstructionDatasetPreprocessed(BaseDataset):
     def __len__(self):
         return len(self.item_list)
 
+
 # -------------- Debug ---------------- #
 if __name__ == "__main__":
     import torch

preprocess/pack_upload_data.py

@@ -29,8 +29,8 @@ def pack_all_scenes(root, scene_list, output_dir):
         pack_scene_data(root, scene, output_dir)
 
 if __name__ == "__main__":
-    root = r"H:\AI\Datasets\nbv_rec_part2"
-    output_dir = r"H:\AI\Datasets\upload_part2"
+    root = r"/media/hofee/repository/data_part_1"
+    output_dir = r"/media/hofee/repository/upload_part1"
     scene_list = os.listdir(root)
     from_idx = 0
     to_idx = len(scene_list)

preprocess/preprocessor.py

@@ -164,10 +164,10 @@ def save_scene_data(root, scene, scene_idx=0, scene_total=1,file_type="txt"):
 
 if __name__ == "__main__":
     #root = "/media/hofee/repository/new_data_with_normal"
-    root = r"H:\AI\Datasets\nbv_rec_part2"
+    root = "/media/hofee/data/data/new_testset"
     scene_list = os.listdir(root)
     from_idx = 0 # 1000
-    to_idx = 600 # 1500
+    to_idx = len(scene_list) # 1500
 
 
     cnt = 0
@@ -179,7 +179,11 @@ if __name__ == "__main__":
             print(f"Scene {scene} has been processed")
             cnt+=1
             continue
+        try:
             save_scene_data(root, scene, cnt, total, file_type="npy")
+        except Exception as e:
+            print(f"Error occurred when processing scene {scene}")
+            print(e)
         cnt+=1
         end = time.time()
         print(f"Time cost: {end-start}")

runners/heuristic.py

@@ -1,425 +0,0 @@
-import os
-import json
-from utils.render import RenderUtil
-from utils.pose import PoseUtil
-from utils.pts import PtsUtil
-from utils.reconstruction import ReconstructionUtil
-
-import torch
-from tqdm import tqdm
-import numpy as np
-import pickle
-
-from PytorchBoot.config import ConfigManager
-import PytorchBoot.namespace as namespace
-import PytorchBoot.stereotype as stereotype
-from PytorchBoot.factory import ComponentFactory
-
-from PytorchBoot.dataset import BaseDataset
-from PytorchBoot.runners.runner import Runner
-from PytorchBoot.utils import Log
-from PytorchBoot.status import status_manager
-from utils.data_load import DataLoadUtil
-
-@stereotype.runner("heuristic")
-class Heuristic(Runner):
-    def __init__(self, config_path):
-        
-        super().__init__(config_path)
-    
-        self.script_path = ConfigManager.get(namespace.Stereotype.RUNNER, "blender_script_path")
-        self.output_dir = ConfigManager.get(namespace.Stereotype.RUNNER, "output_dir")
-        self.voxel_size = ConfigManager.get(namespace.Stereotype.RUNNER, "voxel_size")
-        self.min_new_area = ConfigManager.get(namespace.Stereotype.RUNNER, "min_new_area")
-        self.heuristic_method = ConfigManager.get(namespace.Stereotype.RUNNER, "heuristic_method")
-        self.heuristic_method_config = ConfigManager.get("heuristic_methods", self.heuristic_method)
-        CM = 0.01
-        self.min_new_pts_num = self.min_new_area * (CM / self.voxel_size) **2
-
-        ''' Experiment '''
-        self.load_experiment("nbv_evaluator")
-        self.stat_result_path = os.path.join(self.output_dir, "stat.json")
-        if os.path.exists(self.stat_result_path):
-            with open(self.stat_result_path, "r") as f:
-                self.stat_result = json.load(f)
-        else:
-            self.stat_result = {}
-
-        ''' Test '''
-        self.test_config = ConfigManager.get(namespace.Stereotype.RUNNER, namespace.Mode.TEST)
-        self.test_dataset_name_list = self.test_config["dataset_list"]
-        self.test_set_list = []
-        self.test_writer_list = []
-        seen_name = set()
-        for test_dataset_name in self.test_dataset_name_list:
-            if test_dataset_name not in seen_name:
-                seen_name.add(test_dataset_name)
-            else:
-                raise ValueError("Duplicate test dataset name: {}".format(test_dataset_name))
-            test_set: BaseDataset = ComponentFactory.create(namespace.Stereotype.DATASET, test_dataset_name)
-            self.test_set_list.append(test_set)
-        self.print_info()
-        
-
-    def run(self):
-        Log.info("Loading from epoch {}.".format(self.current_epoch))
-        self.run_heuristic()
-        Log.success("Inference finished.")
-        
-
-    def run_heuristic(self):
- 
-        test_set: BaseDataset
-        for dataset_idx, test_set in enumerate(self.test_set_list):
-            status_manager.set_progress("heuristic", "heuristic", f"dataset", dataset_idx, len(self.test_set_list))
-            test_set_name = test_set.get_name()
-
-            total=int(len(test_set))
-            for i in tqdm(range(total), desc=f"Processing {test_set_name}", ncols=100):
-                try:
-                    data = test_set.__getitem__(i)
-                    scene_name = data["scene_name"]
-                    inference_result_path = os.path.join(self.output_dir, test_set_name, f"{scene_name}.pkl")
-                    if os.path.exists(inference_result_path):
-                        Log.info(f"Inference result already exists for scene: {scene_name}")
-                        continue
-                    
-                    status_manager.set_progress("heuristic", "heuristic", f"Batch[{test_set_name}]", i+1, total)
-                    output = self.predict_sequence(data)
-                    self.save_inference_result(test_set_name, data["scene_name"], output)
-                except Exception as e:
-                    print(e)
-                    Log.error(f"Error, {e}")
-                    continue
-                
-        status_manager.set_progress("heuristic", "heuristic", f"dataset", len(self.test_set_list), len(self.test_set_list))
-        
-    def predict_sequence(self, data, cr_increase_threshold=0, overlap_area_threshold=25, scan_points_threshold=10, max_iter=5000, max_retry=5000, max_success=5000):
-        scene_name = data["scene_name"]
-        Log.info(f"Processing scene: {scene_name}")
-        status_manager.set_status("heuristic", "heuristic", "scene", scene_name)
-
-        ''' data for rendering '''
-        scene_path = data["scene_path"]
-        O_to_L_pose = data["O_to_L_pose"]
-        voxel_threshold = self.voxel_size
-        filter_degree = 75
-        down_sampled_model_pts = data["gt_pts"]
-        
-        first_frame_to_world_9d = data["first_scanned_n_to_world_pose_9d"][0]
-        first_frame_to_world = np.eye(4)
-        first_frame_to_world[:3,:3] = PoseUtil.rotation_6d_to_matrix_numpy(first_frame_to_world_9d[:6])
-        first_frame_to_world[:3,3] = first_frame_to_world_9d[6:]
-        
-        # 获取扫描点
-        root = os.path.dirname(scene_path)
-        display_table_info = DataLoadUtil.get_display_table_info(root, scene_name)
-        radius = display_table_info["radius"]
-        scan_points = np.asarray(ReconstructionUtil.generate_scan_points(display_table_top=0,display_table_radius=radius))
-
-        # 生成位姿序列
-        if self.heuristic_method == "hemisphere_random":
-            pose_sequence = self.generate_hemisphere_random_sequence(
-                max_iter,
-                self.heuristic_method_config
-            )
-        elif self.heuristic_method == "hemisphere_circle_trajectory":
-            pose_sequence = self.generate_hemisphere_circle_sequence(
-                self.heuristic_method_config
-            )
-        else:
-            raise ValueError(f"Unknown heuristic method: {self.heuristic_method}")
-
-        # 执行第一帧
-        first_frame_target_pts, _, first_frame_scan_points_indices = RenderUtil.render_pts(
-            first_frame_to_world, scene_path, self.script_path, scan_points, 
-            voxel_threshold=voxel_threshold, filter_degree=filter_degree, nO_to_nL_pose=O_to_L_pose
-        )
-        
-        # 初始化结果存储
-        scanned_view_pts = [first_frame_target_pts]
-        history_indices = [first_frame_scan_points_indices]
-        pred_cr_seq = []
-        retry_duplication_pose = []
-        retry_no_pts_pose = []
-        retry_overlap_pose = []
-        pose_9d_seq = [first_frame_to_world_9d]
-        
-        last_pred_cr, _ = self.compute_coverage_rate(scanned_view_pts, None, down_sampled_model_pts, threshold=voxel_threshold)
-        pred_cr_seq.append(last_pred_cr)
-        last_pts_num = PtsUtil.voxel_downsample_point_cloud(first_frame_target_pts, voxel_threshold).shape[0]
-        
-        # 执行序列
-        retry = 0
-        success = 0
-        #import ipdb; ipdb.set_trace()
-        combined_scanned_pts_tensor = torch.tensor([0,0,0])
-        cnt = 0
-        for pred_pose in pose_sequence:
-            cnt += 1
-            if retry >= max_retry or success >= max_success:
-                break
-            
-            Log.green(f"迭代: {cnt}/{len(pose_sequence)}, 重试: {retry}/{max_retry}, 成功: {success}/{max_success}")
-            
-            try:
-                new_target_pts, _, new_scan_points_indices = RenderUtil.render_pts(
-                    pred_pose, scene_path, self.script_path, scan_points,
-                    voxel_threshold=voxel_threshold, filter_degree=filter_degree, nO_to_nL_pose=O_to_L_pose
-                )
-                
-                # 检查扫描点重叠
-                if not ReconstructionUtil.check_scan_points_overlap(history_indices, new_scan_points_indices, scan_points_threshold):
-                    curr_overlap_area_threshold = overlap_area_threshold
-                else:
-                    curr_overlap_area_threshold = overlap_area_threshold * 0.5
-
-                # 检查点云重叠
-                downsampled_new_target_pts = PtsUtil.voxel_downsample_point_cloud(new_target_pts, voxel_threshold)
-                overlap, _ = ReconstructionUtil.check_overlap(
-                    downsampled_new_target_pts, down_sampled_model_pts,
-                    overlap_area_threshold=curr_overlap_area_threshold,
-                    voxel_size=voxel_threshold,
-                    require_new_added_pts_num=True
-                )
-                
-                if not overlap:
-                    Log.yellow("no overlap!")
-                    retry += 1
-                    retry_overlap_pose.append(pred_pose.tolist())
-                    continue
-
-                if new_target_pts.shape[0] == 0:
-                    Log.red("新视角无点云")
-                    retry_no_pts_pose.append(pred_pose.tolist())
-                    retry += 1
-                    continue
-                    
-                history_indices.append(new_scan_points_indices)
-                
-                # 计算覆盖率
-                pred_cr, _ = self.compute_coverage_rate(scanned_view_pts, new_target_pts, down_sampled_model_pts, threshold=voxel_threshold)
-                Log.yellow(f"覆盖率: {pred_cr}, 上一次: {last_pred_cr}, 最大: {data['seq_max_coverage_rate']}")
-                
-                # 更新结果
-                pred_cr_seq.append(pred_cr)
-                scanned_view_pts.append(new_target_pts)
-                pose_6d = PoseUtil.matrix_to_rotation_6d_numpy(pred_pose[:3,:3])
-                pose_9d = np.concatenate([
-                    pose_6d,
-                    pred_pose[:3,3]
-                ])
-                pose_9d_seq.append(pose_9d)
-                # 处理点云数据用于combined_scanned_pts
-                combined_scanned_pts = np.vstack(scanned_view_pts)
-                voxel_downsampled_pts, _ = self.voxel_downsample_with_mapping(combined_scanned_pts, voxel_threshold)
-                random_downsampled_pts, _ = PtsUtil.random_downsample_point_cloud(voxel_downsampled_pts, 8192, require_idx=True)
-                combined_scanned_pts_tensor = torch.tensor(random_downsampled_pts, dtype=torch.float32)
-
-                
-                # 检查点数增量
-                pts_num = voxel_downsampled_pts.shape[0]
-                Log.info(f"点数增量: {pts_num - last_pts_num}, 当前: {pts_num}, 上一次: {last_pts_num}")
-                
-                if pts_num - last_pts_num < self.min_new_pts_num:
-                    if pred_cr <= data["seq_max_coverage_rate"] - 1e-2:
-                        retry += 1
-                        retry_duplication_pose.append(pred_pose.tolist())
-                        Log.red(f"点数增量过小 < {self.min_new_pts_num}")
-                    else:
-                        success += 1
-                        Log.success(f"达到目标覆盖率")
-                        
-                last_pts_num = pts_num
-                last_pred_cr = pred_cr
-                
-                if pred_cr >= data["seq_max_coverage_rate"] - 1e-3:
-                    Log.success(f"达到最大覆盖率: {pred_cr}")
-                    
-                    
-            except Exception as e:
-                import traceback
-                traceback.print_exc()
-                Log.error(f"场景 {scene_path} 处理出错: {e}")
-                retry_no_pts_pose.append(pred_pose.tolist())
-                retry += 1
-                continue
-
-        # 返回结果
-        result = {
-            "pred_pose_9d_seq": pose_9d_seq,
-            "combined_scanned_pts_tensor": combined_scanned_pts_tensor,
-            "target_pts_seq": scanned_view_pts,
-            "coverage_rate_seq": pred_cr_seq,
-            "max_coverage_rate": data["seq_max_coverage_rate"],
-            "pred_max_coverage_rate": max(pred_cr_seq),
-            "scene_name": scene_name,
-            "retry_no_pts_pose": retry_no_pts_pose,
-            "retry_duplication_pose": retry_duplication_pose,
-            "retry_overlap_pose": retry_overlap_pose,
-            "best_seq_len": data["best_seq_len"],
-        }
-        
-        self.stat_result[scene_name] = {
-            "coverage_rate_seq": pred_cr_seq,
-            "pred_max_coverage_rate": max(pred_cr_seq),
-            "pred_seq_len": len(pred_cr_seq),
-        }
-        print('success rate: ', max(pred_cr_seq))
-
-        return result
-
-    def voxel_downsample_with_mapping(self, point_cloud, voxel_size=0.003):
-        voxel_indices = np.floor(point_cloud / voxel_size).astype(np.int32)
-        unique_voxels, inverse, counts = np.unique(voxel_indices, axis=0, return_inverse=True, return_counts=True)
-        idx_sort = np.argsort(inverse)
-        idx_unique = idx_sort[np.cumsum(counts)-counts]
-        downsampled_points = point_cloud[idx_unique]
-        return downsampled_points, inverse
-    
-    def compute_coverage_rate(self, scanned_view_pts, new_pts, model_pts, threshold=0.005):
-        if new_pts is not None:
-            new_scanned_view_pts = scanned_view_pts + [new_pts]
-        else:
-            new_scanned_view_pts = scanned_view_pts
-        combined_point_cloud = np.vstack(new_scanned_view_pts)
-        down_sampled_combined_point_cloud = PtsUtil.voxel_downsample_point_cloud(combined_point_cloud,threshold)
-        return ReconstructionUtil.compute_coverage_rate(model_pts, down_sampled_combined_point_cloud, threshold)
-        
-
-    def save_inference_result(self, dataset_name, scene_name, output):
-        dataset_dir = os.path.join(self.output_dir, dataset_name)
-        if not os.path.exists(dataset_dir):
-            os.makedirs(dataset_dir)
-        output_path = os.path.join(dataset_dir, f"{scene_name}.pkl")
-        pickle.dump(output, open(output_path, "wb"))
-        with open(self.stat_result_path, "w") as f:
-            json.dump(self.stat_result, f)
-    
-
-    def get_checkpoint_path(self, is_last=False):
-        return os.path.join(self.experiment_path, namespace.Direcotry.CHECKPOINT_DIR_NAME,
-                            "Epoch_{}.pth".format(
-                                self.current_epoch if self.current_epoch != -1 and not is_last else "last"))
-
-    def load_checkpoint(self, is_last=False):
-        self.load(self.get_checkpoint_path(is_last))
-        Log.success(f"Loaded checkpoint from {self.get_checkpoint_path(is_last)}")
-        if is_last:
-            checkpoint_root = os.path.join(self.experiment_path, namespace.Direcotry.CHECKPOINT_DIR_NAME)
-            meta_path = os.path.join(checkpoint_root, "meta.json")
-            if not os.path.exists(meta_path):
-                raise FileNotFoundError(
-                    "No checkpoint meta.json file in the experiment {}".format(self.experiments_config["name"]))
-            file_path = os.path.join(checkpoint_root, "meta.json")
-            with open(file_path, "r") as f:
-                meta = json.load(f)
-            self.current_epoch = meta["last_epoch"]
-            self.current_iter = meta["last_iter"]
-
-    def load_experiment(self, backup_name=None):
-        super().load_experiment(backup_name)
-        self.current_epoch = self.experiments_config["epoch"]
-
-    def create_experiment(self, backup_name=None):
-        super().create_experiment(backup_name)
-       
-        
-    
-    def print_info(self):
-        def print_dataset(dataset: BaseDataset):
-            config = dataset.get_config()
-            name = dataset.get_name()
-            Log.blue(f"Dataset: {name}")
-            for k,v in config.items():
-                Log.blue(f"\t{k}: {v}")
-
-        super().print_info()
-        table_size = 70
-        Log.blue(f"{'+' + '-' * (table_size // 2)} Datasets {'-' * (table_size // 2)}" + '+')
-        for i, test_set in enumerate(self.test_set_list):
-            Log.blue(f"test dataset {i}: ")
-            print_dataset(test_set)
-
-        Log.blue(f"{'+' + '-' * (table_size // 2)}----------{'-' * (table_size // 2)}" + '+')
-
-    def generate_hemisphere_random_sequence(self, max_iter, config):
-        """Generate a random hemisphere sampling sequence"""
-        radius_fixed = config["radius_fixed"]
-        fixed_radius = config["fixed_radius"]
-        min_radius = config["min_radius"]
-        max_radius = config["max_radius"]
-        poses = []
-        center = np.array(config["center"])
-        
-        for _ in range(max_iter):
-            # 随机采样方向
-            direction = np.random.randn(3)
-            direction[2] = abs(direction[2])  # 确保在上半球
-            direction = direction / np.linalg.norm(direction)
-            
-            # 确定半径
-            if radius_fixed:
-                radius = fixed_radius
-            else:
-                radius = np.random.uniform(min_radius, max_radius)
-                
-            # 计算位置和朝向
-            position = center + direction * radius
-            z_axis = -direction
-            y_axis = np.array([0, 0, 1])
-            x_axis = np.cross(y_axis, z_axis)
-            x_axis = x_axis / np.linalg.norm(x_axis)
-            y_axis = np.cross(z_axis, x_axis)
-            
-            pose = np.eye(4)
-            pose[:3,:3] = np.stack([x_axis, y_axis, z_axis], axis=1)
-            pose[:3,3] = position
-            poses.append(pose)
-            
-        return poses
-
-    def generate_hemisphere_circle_sequence(self, config):
-        """Generate a circular trajectory sampling sequence"""
-        radius_fixed = config["radius_fixed"]
-        fixed_radius = config["fixed_radius"]
-        min_radius = config["min_radius"]
-        max_radius = config["max_radius"]
-        phi_list = config["phi_list"]
-        circle_times = config["circle_times"]
-        
-        poses = []
-        center = np.array(config["center"])
-        
-        for phi in phi_list:  # 仰角
-            phi_rad = np.deg2rad(phi)
-            for i in range(circle_times):  # 方位角
-                theta = i * (2 * np.pi / circle_times)
-                
-                # 确定半径
-                if radius_fixed:
-                    radius = fixed_radius
-                else:
-                    radius = np.random.uniform(min_radius, max_radius)
-                    
-                # 球坐标转笛卡尔坐标
-                x = radius * np.cos(theta) * np.sin(phi_rad)
-                y = radius * np.sin(theta) * np.sin(phi_rad)
-                z = radius * np.cos(phi_rad)
-                position = center + np.array([x, y, z])
-                
-                # 计算朝向
-                direction = (center - position) / np.linalg.norm(center - position)
-                z_axis = direction
-                y_axis = np.array([0, 0, 1])
-                x_axis = np.cross(y_axis, z_axis)
-                x_axis = x_axis / np.linalg.norm(x_axis)
-                y_axis = np.cross(z_axis, x_axis)
-                
-                pose = np.eye(4)
-                pose[:3,:3] = np.stack([x_axis, y_axis, z_axis], axis=1)
-                pose[:3,3] = position
-                poses.append(pose)
-                
-        return poses
-

runners/inference_server.py

@@ -25,7 +25,6 @@ class InferencerServer(Runner):
         self.pipeline:torch.nn.Module = ComponentFactory.create(namespace.Stereotype.PIPELINE, self.pipeline_name)
         self.pipeline = self.pipeline.to(self.device)
         self.pts_num = 8192
-        self.voxel_size = 0.002
             
         ''' Experiment '''
         self.load_experiment("inferencer_server")
@@ -35,14 +34,20 @@ class InferencerServer(Runner):
         scanned_pts = data["scanned_pts"]
         scanned_n_to_world_pose_9d = data["scanned_n_to_world_pose_9d"]
         combined_scanned_views_pts = np.concatenate(scanned_pts, axis=0)
-        voxel_downsampled_combined_scanned_pts = PtsUtil.voxel_downsample_point_cloud(
-            combined_scanned_views_pts, self.voxel_size
-        )
         fps_downsampled_combined_scanned_pts, fps_idx = PtsUtil.fps_downsample_point_cloud(
-            voxel_downsampled_combined_scanned_pts, self.pts_num, require_idx=True
+            combined_scanned_views_pts, self.pts_num, require_idx=True
         )
+        # combined_scanned_views_pts_mask = np.zeros(len(scanned_pts), dtype=np.uint8)
+        # start_idx = 0
+        # for i in range(len(scanned_pts)):
+        #     end_idx = start_idx + len(scanned_pts[i])
+        #     combined_scanned_views_pts_mask[start_idx:end_idx] = i
+        #     start_idx = end_idx
+
+        # fps_downsampled_combined_scanned_pts_mask = combined_scanned_views_pts_mask[fps_idx]
 
         input_data["scanned_pts"] = scanned_pts
+        # input_data["scanned_pts_mask"] = np.asarray(fps_downsampled_combined_scanned_pts_mask, dtype=np.uint8)
         input_data["scanned_n_to_world_pose_9d"] = np.asarray(scanned_n_to_world_pose_9d, dtype=np.float32)
         input_data["combined_scanned_pts"] = np.asarray(fps_downsampled_combined_scanned_pts, dtype=np.float32)
         return input_data

runners/inferencer.py

@@ -23,15 +23,11 @@ from utils.data_load import DataLoadUtil
 @stereotype.runner("inferencer")
 class Inferencer(Runner):
     def __init__(self, config_path):
-        
         super().__init__(config_path)
     
         self.script_path = ConfigManager.get(namespace.Stereotype.RUNNER, "blender_script_path")
         self.output_dir = ConfigManager.get(namespace.Stereotype.RUNNER, "output_dir")
         self.voxel_size = ConfigManager.get(namespace.Stereotype.RUNNER, "voxel_size")
-        self.min_new_area = ConfigManager.get(namespace.Stereotype.RUNNER, "min_new_area")
-        CM = 0.01
-        self.min_new_pts_num = self.min_new_area * (CM / self.voxel_size) **2
         ''' Pipeline '''
         self.pipeline_name = self.config[namespace.Stereotype.PIPELINE]
         self.pipeline:torch.nn.Module = ComponentFactory.create(namespace.Stereotype.PIPELINE, self.pipeline_name)
@@ -78,9 +74,10 @@ class Inferencer(Runner):
 
                 total=int(len(test_set))
                 for i in tqdm(range(total), desc=f"Processing {test_set_name}", ncols=100):
-                    try:
                     data = test_set.__getitem__(i)
                     scene_name = data["scene_name"]
+                    if scene_name != "omniobject3d-book_004":
+                        continue
                     inference_result_path = os.path.join(self.output_dir, test_set_name, f"{scene_name}.pkl")
                     if os.path.exists(inference_result_path):
                         Log.info(f"Inference result already exists for scene: {scene_name}")
@@ -89,14 +86,10 @@ class Inferencer(Runner):
                     status_manager.set_progress("inference", "inferencer", f"Batch[{test_set_name}]", i+1, total)
                     output = self.predict_sequence(data)
                     self.save_inference_result(test_set_name, data["scene_name"], output)
-                    except Exception as e:
-                        print(e)
-                        Log.error(f"Error, {e}")
-                        continue
                     
             status_manager.set_progress("inference", "inferencer", f"dataset", len(self.test_set_list), len(self.test_set_list))
         
-    def predict_sequence(self, data, cr_increase_threshold=0, overlap_area_threshold=25, scan_points_threshold=10, max_iter=50, max_retry = 10, max_success=3):
+    def predict_sequence(self, data, cr_increase_threshold=0, overlap_area_threshold=25, scan_points_threshold=10, max_iter=50, max_retry = 5):
         scene_name = data["scene_name"]
         Log.info(f"Processing scene: {scene_name}")
         status_manager.set_status("inference", "inferencer", "scene", scene_name)
@@ -115,14 +108,13 @@ class Inferencer(Runner):
         
         ''' data for inference '''
         input_data = {}
-        
         input_data["combined_scanned_pts"] = torch.tensor(data["first_scanned_pts"][0], dtype=torch.float32).to(self.device).unsqueeze(0)
-        input_data["scanned_pts_mask"] = [torch.zeros(input_data["combined_scanned_pts"].shape[1], dtype=torch.bool).to(self.device).unsqueeze(0)]
         input_data["scanned_n_to_world_pose_9d"] = [torch.tensor(data["first_scanned_n_to_world_pose_9d"], dtype=torch.float32).to(self.device)]
         input_data["mode"] = namespace.Mode.TEST
         input_pts_N = input_data["combined_scanned_pts"].shape[1]
         
         root = os.path.dirname(scene_path)
+        
         display_table_info = DataLoadUtil.get_display_table_info(root, scene_name)
         radius = display_table_info["radius"]
         scan_points = np.asarray(ReconstructionUtil.generate_scan_points(display_table_top=0,display_table_radius=radius))
@@ -137,13 +129,13 @@ class Inferencer(Runner):
         retry = 0
         pred_cr_seq = [last_pred_cr]
         success = 0
-        last_pts_num = PtsUtil.voxel_downsample_point_cloud(data["first_scanned_pts"][0], voxel_threshold).shape[0]
+        last_pts_num = PtsUtil.voxel_downsample_point_cloud(data["first_scanned_pts"][0], 0.002).shape[0]
         import time
-        while len(pred_cr_seq) < max_iter and retry < max_retry and success < max_success:
-            Log.green(f"iter: {len(pred_cr_seq)}, retry: {retry}/{max_retry}, success: {success}/{max_success}")
-            combined_scanned_pts = np.vstack(scanned_view_pts)
-            voxel_downsampled_combined_scanned_pts_np, inverse =  self.voxel_downsample_with_mapping(combined_scanned_pts, voxel_threshold)
+        while len(pred_cr_seq) < max_iter and retry < max_retry:
+            start_time = time.time()
             output = self.pipeline(input_data)
+            end_time = time.time()
+            print(f"Time taken for inference: {end_time - start_time} seconds")
             pred_pose_9d = output["pred_pose_9d"]
             pred_pose = torch.eye(4, device=pred_pose_9d.device)
 
@@ -151,6 +143,7 @@ class Inferencer(Runner):
             pred_pose[:3,3] = pred_pose_9d[0,6:]
             
             try:
+                start_time = time.time()
                 new_target_pts, new_target_normals, new_scan_points_indices = RenderUtil.render_pts(pred_pose, scene_path, self.script_path, scan_points, voxel_threshold=voxel_threshold, filter_degree=filter_degree, nO_to_nL_pose=O_to_L_pose)
                 #import ipdb; ipdb.set_trace()
                 if not ReconstructionUtil.check_scan_points_overlap(history_indices, new_scan_points_indices, scan_points_threshold):
@@ -159,16 +152,17 @@ class Inferencer(Runner):
                     curr_overlap_area_threshold = overlap_area_threshold * 0.5  
 
                 downsampled_new_target_pts = PtsUtil.voxel_downsample_point_cloud(new_target_pts, voxel_threshold)
-                overlap, _ = ReconstructionUtil.check_overlap(downsampled_new_target_pts, voxel_downsampled_combined_scanned_pts_np, overlap_area_threshold = curr_overlap_area_threshold, voxel_size=voxel_threshold, require_new_added_pts_num = True)
+                overlap, _ = ReconstructionUtil.check_overlap(downsampled_new_target_pts, down_sampled_model_pts, overlap_area_threshold = curr_overlap_area_threshold, voxel_size=voxel_threshold, require_new_added_pts_num = True)
                 if not overlap:
-                    Log.yellow("no overlap!")
                     retry += 1
                     retry_overlap_pose.append(pred_pose.cpu().numpy().tolist())
                     continue
                 
                 history_indices.append(new_scan_points_indices)
+                end_time = time.time()
+                print(f"Time taken for rendering: {end_time - start_time} seconds")
             except Exception as e:
-                Log.error(f"Error in scene {scene_path}, {e}")
+                Log.warning(f"Error in scene {scene_path}, {e}")
                 print("current pose: ", pred_pose)
                 print("curr_pred_cr: ", last_pred_cr)
                 retry_no_pts_pose.append(pred_pose.cpu().numpy().tolist())
@@ -176,94 +170,42 @@ class Inferencer(Runner):
                 continue
             
             if new_target_pts.shape[0] == 0:
-                Log.red("no pts in new target")
+                print("no pts in new target")
                 retry_no_pts_pose.append(pred_pose.cpu().numpy().tolist())
                 retry += 1
                 continue
             
+            start_time = time.time()
             pred_cr, _ = self.compute_coverage_rate(scanned_view_pts, new_target_pts, down_sampled_model_pts, threshold=voxel_threshold)
-            Log.yellow(f"{pred_cr}, {last_pred_cr}, max: , {data['seq_max_coverage_rate']}")
+            end_time = time.time()
+            print(f"Time taken for coverage rate computation: {end_time - start_time} seconds")
+            print(pred_cr, last_pred_cr, " max: ", data["seq_max_coverage_rate"])
             if pred_cr >= data["seq_max_coverage_rate"] - 1e-3:
                 print("max coverage rate reached!: ", pred_cr)
+                success += 1
         
 
-
+            retry = 0
             pred_cr_seq.append(pred_cr)
             scanned_view_pts.append(new_target_pts)
             
             input_data["scanned_n_to_world_pose_9d"] = [torch.cat([input_data["scanned_n_to_world_pose_9d"][0], pred_pose_9d], dim=0)]
-            start_indices = [0]
-            total_points = 0
-            for pts in scanned_view_pts:
-                total_points += pts.shape[0]
-                start_indices.append(total_points)
+            
             combined_scanned_pts = np.vstack(scanned_view_pts)
-            voxel_downsampled_combined_scanned_pts_np, inverse =  self.voxel_downsample_with_mapping(combined_scanned_pts, voxel_threshold)
-            random_downsampled_combined_scanned_pts_np, random_downsample_idx = PtsUtil.random_downsample_point_cloud(voxel_downsampled_combined_scanned_pts_np, input_pts_N, require_idx=True)
-            all_idx_unique = np.arange(len(voxel_downsampled_combined_scanned_pts_np))
-            all_random_downsample_idx = all_idx_unique[random_downsample_idx]
-            scanned_pts_mask = []
-            for idx, start_idx in enumerate(start_indices):
-                if idx == len(start_indices) - 1:
-                    break
-                end_idx = start_indices[idx+1]
-                view_inverse = inverse[start_idx:end_idx]
-                view_unique_downsampled_idx = np.unique(view_inverse)
-                view_unique_downsampled_idx_set = set(view_unique_downsampled_idx)
-                mask = np.array([idx in view_unique_downsampled_idx_set for idx in all_random_downsample_idx])
-                scanned_pts_mask.append(mask)
-
+            voxel_downsampled_combined_scanned_pts_np = PtsUtil.voxel_downsample_point_cloud(combined_scanned_pts, 0.002)
+            random_downsampled_combined_scanned_pts_np = PtsUtil.random_downsample_point_cloud(voxel_downsampled_combined_scanned_pts_np, input_pts_N)
             input_data["combined_scanned_pts"] = torch.tensor(random_downsampled_combined_scanned_pts_np, dtype=torch.float32).unsqueeze(0).to(self.device)
-            #import ipdb; ipdb.set_trace()
-            input_data["scanned_pts_mask"] = [torch.tensor(scanned_pts_mask, dtype=torch.bool)]
-
 
+            if success > 3:
+                break
             last_pred_cr = pred_cr
             pts_num = voxel_downsampled_combined_scanned_pts_np.shape[0]
-            Log.info(f"delta pts num:,{pts_num - last_pts_num },{pts_num}, {last_pts_num}")
-
-            if pts_num - last_pts_num < self.min_new_pts_num and pred_cr <= data["seq_max_coverage_rate"] - 1e-2:
+            if pts_num - last_pts_num < 10 and pred_cr < data["seq_max_coverage_rate"] - 1e-3:
                 retry += 1
                 retry_duplication_pose.append(pred_pose.cpu().numpy().tolist())
-                Log.red(f"delta pts num < {self.min_new_pts_num}:, {pts_num}, {last_pts_num}")
-            elif pts_num - last_pts_num < self.min_new_pts_num and pred_cr > data["seq_max_coverage_rate"] - 1e-2:
-                success += 1
-                Log.success(f"delta pts num < {self.min_new_pts_num}:, {pts_num}, {last_pts_num}")
-
+                print("delta pts num < 10:", pts_num, last_pts_num)
             last_pts_num = pts_num
 
-
-        input_data["scanned_n_to_world_pose_9d"] = input_data["scanned_n_to_world_pose_9d"][0].cpu().numpy().tolist()
-        result = {
-            "pred_pose_9d_seq": input_data["scanned_n_to_world_pose_9d"],
-            "combined_scanned_pts": input_data["combined_scanned_pts"],
-            "target_pts_seq": scanned_view_pts,
-            "coverage_rate_seq": pred_cr_seq,
-            "max_coverage_rate": data["seq_max_coverage_rate"],
-            "pred_max_coverage_rate": max(pred_cr_seq),
-            "scene_name": scene_name,
-            "retry_no_pts_pose": retry_no_pts_pose,
-            "retry_duplication_pose": retry_duplication_pose,
-            "retry_overlap_pose": retry_overlap_pose,
-            "best_seq_len": data["best_seq_len"],
-        }
-        self.stat_result[scene_name] = {
-            "coverage_rate_seq": pred_cr_seq,
-            "pred_max_coverage_rate": max(pred_cr_seq),
-            "pred_seq_len": len(pred_cr_seq),
-        }
-        print('success rate: ', max(pred_cr_seq))
-
-        return result
-
-    def voxel_downsample_with_mapping(self, point_cloud, voxel_size=0.003):
-        voxel_indices = np.floor(point_cloud / voxel_size).astype(np.int32)
-        unique_voxels, inverse, counts = np.unique(voxel_indices, axis=0, return_inverse=True, return_counts=True)
-        idx_sort = np.argsort(inverse)
-        idx_unique = idx_sort[np.cumsum(counts)-counts]
-        downsampled_points = point_cloud[idx_unique]
-        return downsampled_points, inverse
-    
     def compute_coverage_rate(self, scanned_view_pts, new_pts, model_pts, threshold=0.005):
         if new_pts is not None:
             new_scanned_view_pts = scanned_view_pts + [new_pts]

runners/view_generator.py

@@ -9,7 +9,7 @@ class ViewGenerator(Runner):
         self.config_path = config_path
         
     def run(self):
-        result = subprocess.run(['blender', '-b', '-P', '../blender/run_blender.py', '--', self.config_path])
+        result = subprocess.run(['/home/hofee/blender-4.0.2-linux-x64/blender', '-b', '-P', '../blender/run_blender.py', '--', self.config_path])
         print()
 
     def create_experiment(self, backup_name=None):

utils/render.py

@@ -84,10 +84,13 @@ class RenderUtil:
             params_data_path = os.path.join(temp_dir, "params.json")
             with open(params_data_path, 'w') as f:
                 json.dump(params, f)
+            start_time = time.time()
             result = subprocess.run([
                 '/home/hofee/blender-4.0.2-linux-x64/blender', '-b', '-P', script_path, '--', temp_dir
             ], capture_output=True, text=True)
-            # print(result)
+            end_time = time.time()
+
+            print(f"-- Time taken for blender: {end_time - start_time} seconds")
             path = os.path.join(temp_dir, "tmp")
             cam_info = DataLoadUtil.load_cam_info(path, binocular=True)
             depth_L, depth_R = DataLoadUtil.load_depth(
@@ -95,6 +98,7 @@ class RenderUtil:
                     cam_info["far_plane"], 
                     binocular=True
                 )
+            start_time = time.time()
             mask_L, mask_R = DataLoadUtil.load_seg(path, binocular=True)
             normal_L = DataLoadUtil.load_normal(path, binocular=True, left_only=True)
             ''' target points '''
@@ -131,5 +135,7 @@ class RenderUtil:
             if not has_points:
                 target_points = np.zeros((0, 3))
                 target_normals = np.zeros((0, 3))
+            end_time = time.time()
+            print(f"-- Time taken for processing: {end_time - start_time} seconds")
             #import ipdb; ipdb.set_trace()
             return target_points, target_normals, scan_points_indices

utils/vis.py

@@ -175,9 +175,6 @@ class visualizeUtil:
         np.savetxt(os.path.join(output_dir, "nrm.txt"), visualized_nrm)
         np.savetxt(os.path.join(output_dir, "pts.txt"), pts_world)
         
-    # @staticmethod
-    # def save_
-        
 # ------ Debug ------
 
 if __name__ == "__main__":