core/nbv_dataset.py

@@ -124,20 +124,63 @@ class NBVReconstructionDataset(BaseDataset):
             scanned_n_to_world_pose,
             scanned_target_pts_num,
         ) = ([], [], [], [])
+        target_pts_num_dict = DataLoadUtil.load_target_pts_num_dict(
+            self.root_dir, scene_name
+        )
         for view in scanned_views:
             frame_idx = view[0]
             coverage_rate = view[1]
             view_path = DataLoadUtil.get_path(self.root_dir, scene_name, frame_idx)
             cam_info = DataLoadUtil.load_cam_info(view_path, binocular=True)
-            
-            n_to_world_pose = cam_info["cam_to_world"]  
-            target_point_cloud = (
-                DataLoadUtil.load_from_preprocessed_pts(view_path)
-            )
-            target_pts_num = target_point_cloud.shape[0]
-            downsampled_target_point_cloud = PtsUtil.random_downsample_point_cloud(
-                target_point_cloud, self.pts_num
-            )
+            target_pts_num = target_pts_num_dict[frame_idx]
+            n_to_world_pose = cam_info["cam_to_world"]
+            nR_to_world_pose = cam_info["cam_to_world_R"]
+
+            if self.load_from_preprocess:
+                downsampled_target_point_cloud = (
+                    DataLoadUtil.load_from_preprocessed_pts(view_path)
+                )
+            else:
+                cached_data = None
+                if self.cache:
+                    cached_data = self.load_from_cache(scene_name, frame_idx)
+
+                if cached_data is None:
+                    print("load depth")
+                    depth_L, depth_R = DataLoadUtil.load_depth(
+                        view_path,
+                        cam_info["near_plane"],
+                        cam_info["far_plane"],
+                        binocular=True,
+                    )
+                    point_cloud_L = DataLoadUtil.get_point_cloud(
+                        depth_L, cam_info["cam_intrinsic"], n_to_world_pose
+                    )["points_world"]
+                    point_cloud_R = DataLoadUtil.get_point_cloud(
+                        depth_R, cam_info["cam_intrinsic"], nR_to_world_pose
+                    )["points_world"]
+
+                    point_cloud_L = PtsUtil.random_downsample_point_cloud(
+                        point_cloud_L, 65536
+                    )
+                    point_cloud_R = PtsUtil.random_downsample_point_cloud(
+                        point_cloud_R, 65536
+                    )
+                    overlap_points = PtsUtil.get_overlapping_points(
+                        point_cloud_L, point_cloud_R
+                    )
+                    downsampled_target_point_cloud = (
+                        PtsUtil.random_downsample_point_cloud(
+                            overlap_points, self.pts_num
+                        )
+                    )
+                    if self.cache:
+                        self.save_to_cache(
+                            scene_name, frame_idx, downsampled_target_point_cloud
+                        )
+                else:
+                    downsampled_target_point_cloud = cached_data
+
             scanned_views_pts.append(downsampled_target_point_cloud)
             scanned_coverages_rate.append(coverage_rate)
             n_to_world_6d = PoseUtil.matrix_to_rotation_6d_numpy(

core/seq_dataset.py

@@ -74,12 +74,24 @@ class SeqNBVReconstructionDataset(BaseDataset):
         max_coverage_rate = data_item_info["max_coverage_rate"]
         scene_name = data_item_info["scene_name"]
         first_cam_info = DataLoadUtil.load_cam_info(DataLoadUtil.get_path(self.root_dir, scene_name, first_frame_idx), binocular=True)
+       
         first_view_path = DataLoadUtil.get_path(self.root_dir, scene_name, first_frame_idx)
         first_left_cam_pose = first_cam_info["cam_to_world"]
+        first_right_cam_pose = first_cam_info["cam_to_world_R"]
         first_center_cam_pose = first_cam_info["cam_to_world_O"]
-        first_target_point_cloud = DataLoadUtil.load_from_preprocessed_pts(first_view_path)
-        first_pts_num = first_target_point_cloud.shape[0]
-        first_downsampled_target_point_cloud = PtsUtil.random_downsample_point_cloud(first_target_point_cloud, self.pts_num)
+        if self.load_from_preprocess:
+            first_downsampled_target_point_cloud = DataLoadUtil.load_from_preprocessed_pts(first_view_path)
+        else:    
+            first_depth_L, first_depth_R = DataLoadUtil.load_depth(first_view_path, first_cam_info['near_plane'], first_cam_info['far_plane'], binocular=True)
+            
+            first_point_cloud_L = DataLoadUtil.get_point_cloud(first_depth_L, first_cam_info['cam_intrinsic'], first_left_cam_pose)['points_world']
+            first_point_cloud_R = DataLoadUtil.get_point_cloud(first_depth_R, first_cam_info['cam_intrinsic'], first_right_cam_pose)['points_world']
+            
+            first_point_cloud_L = PtsUtil.random_downsample_point_cloud(first_point_cloud_L, 65536)
+            first_point_cloud_R = PtsUtil.random_downsample_point_cloud(first_point_cloud_R, 65536)
+            first_overlap_points = PtsUtil.get_overlapping_points(first_point_cloud_L, first_point_cloud_R)
+            first_downsampled_target_point_cloud = PtsUtil.random_downsample_point_cloud(first_overlap_points, self.pts_num)
+     
         first_to_world_rot_6d = PoseUtil.matrix_to_rotation_6d_numpy(np.asarray(first_left_cam_pose[:3,:3]))
         first_to_world_trans = first_left_cam_pose[:3,3]
         first_to_world_9d = np.concatenate([first_to_world_rot_6d, first_to_world_trans], axis=0)
@@ -90,9 +102,6 @@ class SeqNBVReconstructionDataset(BaseDataset):
         model_points_normals = DataLoadUtil.load_points_normals(self.root_dir, scene_name)
 
         data_item = {
-            "first_pts_num": np.asarray(
-                first_pts_num, dtype=np.int32
-            ),
             "first_pts": np.asarray([first_downsampled_target_point_cloud],dtype=np.float32),
             "combined_scanned_pts": np.asarray(first_downsampled_target_point_cloud,dtype=np.float32),
             "first_to_world_9d": np.asarray([first_to_world_9d],dtype=np.float32),

utils/data_load.py

@@ -237,7 +237,7 @@ class DataLoadUtil:
     @staticmethod
     def load_from_preprocessed_pts(path):
         npy_path = os.path.join(
-            os.path.dirname(path), "pts", os.path.basename(path) + ".npy"
+            os.path.dirname(path), "points", os.path.basename(path) + ".npy"
         )
         pts = np.load(npy_path)
         return pts

utils/reconstruction.py

@@ -8,9 +8,9 @@ class ReconstructionUtil:
     def compute_coverage_rate(target_point_cloud, combined_point_cloud, threshold=0.01):
         kdtree = cKDTree(combined_point_cloud)
         distances, _ = kdtree.query(target_point_cloud)
-        covered_points_num = np.sum(distances < threshold)
-        coverage_rate = covered_points_num / target_point_cloud.shape[0]
-        return coverage_rate, covered_points_num
+        covered_points = np.sum(distances < threshold*2)
+        coverage_rate = covered_points / target_point_cloud.shape[0]
+        return coverage_rate
     
     @staticmethod
     def compute_overlap_rate(new_point_cloud, combined_point_cloud, threshold=0.01):
@@ -46,12 +46,10 @@ class ReconstructionUtil:
         downsampled_max_rec_pts = PtsUtil.voxel_downsample_point_cloud(max_rec_pts, threshold)
         
         max_rec_pts_num = downsampled_max_rec_pts.shape[0]
-        max_real_rec_pts_coverage, _ = ReconstructionUtil.compute_coverage_rate(target_point_cloud, downsampled_max_rec_pts, threshold)
+        max_rec_pts_coverage = ReconstructionUtil.compute_coverage_rate(target_point_cloud, downsampled_max_rec_pts, threshold)
         
-        new_coverage, new_covered_num = ReconstructionUtil.compute_coverage_rate(downsampled_max_rec_pts, combined_point_cloud, threshold)
+        new_coverage = ReconstructionUtil.compute_coverage_rate(downsampled_max_rec_pts, combined_point_cloud, threshold)
         current_coverage = new_coverage
-        current_covered_num = new_covered_num
-        
         remaining_views = list(range(len(point_cloud_list)))
         view_sequence = [(init_view, current_coverage)]
         cnt_processed_view = 0
@@ -63,7 +61,6 @@ class ReconstructionUtil:
             best_view = None
             best_coverage_increase = -1
             best_combined_point_cloud = None
-            best_covered_num = 0
         
             for view_index in remaining_views:
                 if point_cloud_list[view_index].shape[0] == 0:
@@ -84,26 +81,25 @@ class ReconstructionUtil:
                 start = time.time()
                 new_combined_point_cloud = np.vstack([combined_point_cloud, point_cloud_list[view_index]])
                 new_downsampled_combined_point_cloud = PtsUtil.voxel_downsample_point_cloud(new_combined_point_cloud,threshold)
-                new_coverage, new_covered_num = ReconstructionUtil.compute_coverage_rate(downsampled_max_rec_pts, new_downsampled_combined_point_cloud, threshold)
+                new_coverage = ReconstructionUtil.compute_coverage_rate(downsampled_max_rec_pts, new_downsampled_combined_point_cloud, threshold)
                 end = time.time()   
                 #print(f"compute_coverage_rate Time: {end-start}")
                 coverage_increase = new_coverage - current_coverage
                 if coverage_increase > best_coverage_increase:
                     best_coverage_increase = coverage_increase
                     best_view = view_index
-                    best_covered_num = new_covered_num
                     best_combined_point_cloud = new_downsampled_combined_point_cloud
                     
             
             if best_view is not None:
-                if best_coverage_increase <=1e-3 or best_covered_num - current_covered_num <= 5:
+                if best_coverage_increase <=1e-3:
                     break
                 
                 selected_views.append(best_view)
                 best_rec_pts_num = best_combined_point_cloud.shape[0]
-                print(f"Current rec pts num: {curr_rec_pts_num}, Best rec pts num: {best_rec_pts_num}, Best cover pts: {best_covered_num}, Max rec pts num: {max_rec_pts_num}")
-                print(f"Current coverage: {current_coverage}, Best coverage increase: {best_coverage_increase}, Max Real coverage: {max_real_rec_pts_coverage}")
-                current_covered_num = best_covered_num
+                print(f"Current rec pts num: {curr_rec_pts_num}, Best rec pts num: {best_rec_pts_num}, Max rec pts num: {max_rec_pts_num}")
+                print(f"Current coverage: {current_coverage}, Best coverage increase: {best_coverage_increase}, Max coverage: {max_rec_pts_coverage}")
+                
                 curr_rec_pts_num = best_rec_pts_num
                 combined_point_cloud = best_combined_point_cloud
                 remaining_views.remove(best_view)