Merge branch 'ab_global_only' of http://git.hofee.top/hofee/nbv_reconstruction into ab_global_only

configs/local/inference_config.yaml

@@ -6,16 +6,16 @@ runner:
     cuda_visible_devices: "0,1,2,3,4,5,6,7"
     
   experiment:
-    name: train_ab_global_only
+    name: train_ab_global_only_p++_wp
     root_dir: "experiments"
-    epoch: -1 # -1 stands for last epoch
+    epoch: 922 # -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/global_only_ycb_241204"
+  output_dir: "/media/hofee/data/data/p++_wp_temp_cluster"
   pipeline: nbv_reconstruction_pipeline
   voxel_size: 0.003
   min_new_area: 1.0
@@ -34,8 +34,8 @@ dataset:
   #   load_from_preprocess: True
 
   OmniObject3d_test:
-    root_dir: "/media/hofee/data/results/ycb_preprocessed_dataset"
-    model_dir: "/media/hofee/data/data/ycb_obj"
+    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
@@ -52,7 +52,7 @@ dataset:
 pipeline:
   nbv_reconstruction_pipeline:
     modules:
-      pts_encoder: pointnet_encoder
+      pts_encoder: pointnet++_encoder
       seq_encoder: transformer_seq_encoder
       pose_encoder: pose_encoder
       view_finder: gf_view_finder
@@ -60,6 +60,9 @@ pipeline:
     global_scanned_feat: True
 
 module:
+  pointnet++_encoder:
+    in_dim: 3
+
   pointnet_encoder:
     in_dim: 3
     out_dim: 1024

configs/local/strategy_generate_config.yaml

@@ -15,13 +15,13 @@ runner:
     overlap_area_threshold: 30
     compute_with_normal: False
     scan_points_threshold: 10
-    overwrite: False
+    overwrite: False  
     seq_num: 10
     dataset_list:
       - OmniObject3d
   
 datasets:
     OmniObject3d:
-      root_dir: /media/hofee/data/results/ycb_view_data
+      root_dir: /media/hofee/data/data/test_bottle/view 
       from: 0
       to: -1 # ..-1 means end

configs/local/view_generate_config.yaml

@@ -8,11 +8,11 @@ runner:
     root_dir: experiments
   generate:
     port: 5002
-    from: 1
+    from: 0
     to: 50 # -1 means all
-    object_dir: /media/hofee/data/data/ycb_obj
+    object_dir: /media/hofee/data/data/test_bottle/bottle_mesh
     table_model_path: /media/hofee/data/data/others/table.obj
-    output_dir: /media/hofee/data/results/ycb_view_data
+    output_dir: /media/hofee/data/data/test_bottle/view
     binocular_vision: true
     plane_size: 10
     max_views: 512
@@ -34,7 +34,7 @@ runner:
         max_y: 0.05
         min_z: 0.01
         max_z: 0.01
-        random_rotation_ratio: 0.3
+        random_rotation_ratio: 0.0
       random_objects:
         num: 4
         cluster: 0.9

core/pipeline.py

@@ -75,6 +75,8 @@ class NBVReconstructionPipeline(nn.Module):
 
     def forward_test(self, data):
         main_feat = self.get_main_feat(data)
+        repeat_num = data.get("repeat_num", 100)
+        main_feat = main_feat.repeat(repeat_num, 1)
         estimated_delta_rot_9d, in_process_sample = self.view_finder.next_best_view(
             main_feat
         )

core/seq_dataset.py

@@ -64,11 +64,15 @@ class SeqReconstructionDataset(BaseDataset):
             scene_max_cr_idx = 0
             frame_len = DataLoadUtil.get_scene_seq_length(self.root_dir, scene_name)
 
-            for i in range(frame_len):
+            for i in range(10,frame_len):
                 path = DataLoadUtil.get_path(self.root_dir, scene_name, i)
                 pts = DataLoadUtil.load_from_preprocessed_pts(path, "npy")
+                print(pts.shape)
                 if pts.shape[0] == 0:
                     continue
+                else:
+                    break
+            print(i)
             datalist.append({
                 "scene_name": scene_name,
                 "first_frame": i,
@@ -180,9 +184,9 @@ if __name__ == "__main__":
     np.random.seed(seed)
     
     config = {
-        "root_dir": "/media/hofee/data/results/ycb_view_data",
+        "root_dir": "/media/hofee/data/data/test_bottle/view",
         "source": "seq_reconstruction_dataset",
-        "split_file": "/media/hofee/data/results/ycb_test.txt",
+        "split_file": "/media/hofee/data/data/test_bottle/test_bottle.txt",
         "load_from_preprocess": True,
         "filter_degree": 75,
         "num_workers": 0,
@@ -190,7 +194,7 @@ if __name__ == "__main__":
         "type": namespace.Mode.TEST,
     }
     
-    output_dir = "/media/hofee/data/results/ycb_preprocessed_dataset"
+    output_dir = "/media/hofee/data/data/test_bottle/preprocessed_dataset"
     os.makedirs(output_dir, exist_ok=True)
 
     ds = SeqReconstructionDataset(config)

core/seq_dataset_preprocessed.py

@@ -21,7 +21,7 @@ class SeqReconstructionDatasetPreprocessed(BaseDataset):
         super(SeqReconstructionDatasetPreprocessed, self).__init__(config)
         self.config = config
         self.root_dir = config["root_dir"]
-        self.real_root_dir = r"/media/hofee/data/results/ycb_view_data"
+        self.real_root_dir = r"/media/hofee/data/data/new_testset"
         self.item_list = os.listdir(self.root_dir)
 
     def __getitem__(self, index):
@@ -66,7 +66,7 @@ if __name__ == "__main__":
     load_from_preprocess: True
     '''
     config = {
-        "root_dir": "H:\\AI\\Datasets\\packed_test_data",
+        "root_dir": "/media/hofee/data/data/test_bottle/preprocessed_dataset",
         "source": "seq_reconstruction_dataset",
         "split_file": "H:\\AI\\Datasets\\data_list\\OmniObject3d_test.txt",
         "load_from_preprocess": True,

preprocess/preprocessor.py

@@ -164,7 +164,7 @@ 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"/media/hofee/data/results/ycb_view_data"
+    root = r"/media/hofee/data/data/test_bottle/view"
     scene_list = os.listdir(root)
     from_idx = 0 # 1000
     to_idx = len(scene_list) # 1500

runners/inferencer.py

@@ -83,6 +83,7 @@ class Inferencer(Runner):
                         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
@@ -143,74 +144,87 @@ class Inferencer(Runner):
             voxel_downsampled_combined_scanned_pts_np, inverse =  self.voxel_downsample_with_mapping(combined_scanned_pts, voxel_threshold)
             output = self.pipeline(input_data)
             pred_pose_9d = output["pred_pose_9d"]
-            import ipdb; ipdb.set_trace()
             pred_pose = torch.eye(4, device=pred_pose_9d.device)
-
-            pred_pose[:3,:3] = PoseUtil.rotation_6d_to_matrix_tensor_batch(pred_pose_9d[:,:6])[0]
-            pred_pose[:3,3] = pred_pose_9d[0,6:]
-            
-            try:
-                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)
+            # # save pred_pose_9d ------
+            # root = "/media/hofee/data/project/python/nbv_reconstruction/nbv_reconstruction/temp_output_result"
+            # scene_dir = os.path.join(root, scene_name)
+            # if not os.path.exists(scene_dir):
+            #     os.makedirs(scene_dir)
+            # pred_9d_path = os.path.join(scene_dir,f"pred_pose_9d_{len(pred_cr_seq)}.npy")
+            # pts_path = os.path.join(scene_dir,f"combined_scanned_pts_{len(pred_cr_seq)}.txt")
+            # np_combined_scanned_pts = input_data["combined_scanned_pts"][0].cpu().numpy()
+            # np.save(pred_9d_path, pred_pose_9d.cpu().numpy())
+            # np.savetxt(pts_path, np_combined_scanned_pts)
+            # # ----- ----- -----
+            pred_pose_9d_candidates = PredictResult(pred_pose_9d.cpu().numpy(), input_pts=input_data["combined_scanned_pts"][0].cpu().numpy(), cluster_params=dict(eps=0.25, min_samples=3)).candidate_9d_poses
+            for pred_pose_9d in pred_pose_9d_candidates:
                 #import ipdb; ipdb.set_trace()
-                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  
+                pred_pose_9d = torch.tensor(pred_pose_9d, dtype=torch.float32).to(self.device).unsqueeze(0)
+                pred_pose[:3,:3] = PoseUtil.rotation_6d_to_matrix_tensor_batch(pred_pose_9d[:,:6])[0]
+                pred_pose[:3,3] = pred_pose_9d[0,6:]
+                try:
+                    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):
+                        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, voxel_downsampled_combined_scanned_pts_np, overlap_area_threshold = curr_overlap_area_threshold, voxel_size=voxel_threshold, require_new_added_pts_num = True)
-                if not overlap:
-                    Log.yellow("no overlap!")
+                    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)
+                    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)
+                except Exception as e:
+                    Log.error(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())
                     retry += 1
-                    retry_overlap_pose.append(pred_pose.cpu().numpy().tolist())
                     continue
                 
-                history_indices.append(new_scan_points_indices)
-            except Exception as e:
-                Log.error(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())
-                retry += 1
-                continue
-            
-            if new_target_pts.shape[0] == 0:
-                Log.red("no pts in new target")
-                retry_no_pts_pose.append(pred_pose.cpu().numpy().tolist())
-                retry += 1
-                continue
-            
-            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']}")
-            if pred_cr >= data["seq_max_coverage_rate"] - 1e-3:
-                print("max coverage rate reached!: ", pred_cr)
+                if new_target_pts.shape[0] == 0:
+                    Log.red("no pts in new target")
+                    retry_no_pts_pose.append(pred_pose.cpu().numpy().tolist())
+                    retry += 1
+                    continue
                 
-        
-
-            pred_cr_seq.append(pred_cr)
-            scanned_view_pts.append(new_target_pts)
+                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']}")
+                if pred_cr >= data["seq_max_coverage_rate"] - 1e-3:
+                    print("max coverage rate reached!: ", pred_cr)
+                    
             
-            input_data["scanned_n_to_world_pose_9d"] = [torch.cat([input_data["scanned_n_to_world_pose_9d"][0], pred_pose_9d], dim=0)]
-            
-            combined_scanned_pts = np.vstack(scanned_view_pts)
-            voxel_downsampled_combined_scanned_pts_np = PtsUtil.voxel_downsample_point_cloud(combined_scanned_pts, voxel_threshold)
-            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)
 
-            
-            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}")
+                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)]
+                
+                combined_scanned_pts = np.vstack(scanned_view_pts)
+                voxel_downsampled_combined_scanned_pts_np = PtsUtil.voxel_downsample_point_cloud(combined_scanned_pts, voxel_threshold)
+                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)
 
-            if pts_num - last_pts_num < self.min_new_pts_num and pred_cr <= data["seq_max_coverage_rate"] - 1e-2:
-                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}")
+                
+                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}")
 
-            last_pts_num = pts_num
+                if pts_num - last_pts_num < self.min_new_pts_num and pred_cr <= data["seq_max_coverage_rate"] - 1e-2:
+                    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}")
+
+                last_pts_num = pts_num
+                break
             
 
         input_data["scanned_n_to_world_pose_9d"] = input_data["scanned_n_to_world_pose_9d"][0].cpu().numpy().tolist()

utils/render.py

@@ -88,6 +88,7 @@ class RenderUtil:
                 '/home/hofee/blender-4.0.2-linux-x64/blender', '-b', '-P', script_path, '--', temp_dir
             ], capture_output=True, text=True)
             #print(result)
+            #import ipdb; ipdb.set_trace()
             path = os.path.join(temp_dir, "tmp")
             cam_info = DataLoadUtil.load_cam_info(path, binocular=True)
             depth_L, depth_R = DataLoadUtil.load_depth(

utils/vis.py

@@ -7,6 +7,7 @@ import trimesh
 sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 from utils.data_load import DataLoadUtil
 from utils.pts import PtsUtil
+from utils.pose import PoseUtil
 
 class visualizeUtil:
     
@@ -33,7 +34,22 @@ class visualizeUtil:
         all_cam_axis = np.array(all_cam_axis).reshape(-1, 3)
         np.savetxt(os.path.join(output_dir, "all_cam_pos.txt"), all_cam_pos)
         np.savetxt(os.path.join(output_dir, "all_cam_axis.txt"), all_cam_axis)
-        
+    
+    @staticmethod
+    def get_cam_pose_and_cam_axis(cam_pose, is_6d_pose):
+        if is_6d_pose:
+            matrix_cam_pose = np.eye(4)
+            matrix_cam_pose[:3,:3] = PoseUtil.rotation_6d_to_matrix_numpy(cam_pose[:6])
+            matrix_cam_pose[:3, 3] = cam_pose[6:]
+        else:
+            matrix_cam_pose = cam_pose
+        cam_pos = matrix_cam_pose[:3, 3]
+        cam_axis = matrix_cam_pose[:3, 2]
+        num_samples = 10
+        sample_points = [cam_pos + 0.02*t * cam_axis for t in range(num_samples)]
+        sample_points = np.array(sample_points)
+        return cam_pos, sample_points
+    
     @staticmethod
     def save_all_combined_pts(root, scene, output_dir):
         length = DataLoadUtil.get_scene_seq_length(root, scene)