global_and_partial_global: all

core/nbv_dataset.py

@@ -7,6 +7,7 @@ from PytorchBoot.utils.log_util import Log
 import torch
 import os
 import sys
+import time
 
 sys.path.append(r"/data/hofee/project/nbv_rec/nbv_reconstruction")
 
@@ -34,7 +35,7 @@ class NBVReconstructionDataset(BaseDataset):
             #self.model_dir = config["model_dir"]
             self.filter_degree = config["filter_degree"]
         if self.type == namespace.Mode.TRAIN:
-            scale_ratio = 100
+            scale_ratio = 50
             self.datalist = self.datalist*scale_ratio
         if self.cache:
             expr_root = ConfigManager.get("runner", "experiment", "root_dir")
@@ -114,8 +115,13 @@ class NBVReconstructionDataset(BaseDataset):
         except Exception as e:
             Log.error(f"Save cache failed: {e}")
 
-    def voxel_downsample_with_mask(self, pts, voxel_size):
-        pass
+    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 __getitem__(self, index):
@@ -129,6 +135,9 @@ class NBVReconstructionDataset(BaseDataset):
             scanned_coverages_rate,
             scanned_n_to_world_pose,
         ) = ([], [], [])
+        start_time = time.time()
+        start_indices = [0]
+        total_points = 0
         for view in scanned_views:
             frame_idx = view[0]
             coverage_rate = view[1]
@@ -140,7 +149,7 @@ class NBVReconstructionDataset(BaseDataset):
                 DataLoadUtil.load_from_preprocessed_pts(view_path)
             )
             downsampled_target_point_cloud = PtsUtil.random_downsample_point_cloud(
-                target_point_cloud, self.pts_num
+                target_point_cloud, self.pts_num, replace=False
             )
             scanned_views_pts.append(downsampled_target_point_cloud)
             scanned_coverages_rate.append(coverage_rate)
@@ -150,8 +159,12 @@ class NBVReconstructionDataset(BaseDataset):
             n_to_world_trans = n_to_world_pose[:3, 3]
             n_to_world_9d = np.concatenate([n_to_world_6d, n_to_world_trans], axis=0)
             scanned_n_to_world_pose.append(n_to_world_9d)
+            total_points += len(downsampled_target_point_cloud)
+            start_indices.append(total_points)
 
 
+        end_time = time.time()
+        #Log.info(f"load data time: {end_time - start_time}")
         nbv_idx, nbv_coverage_rate = nbv[0], nbv[1]
         nbv_path = DataLoadUtil.get_path(self.root_dir, scene_name, nbv_idx)
         cam_info = DataLoadUtil.load_cam_info(nbv_path)
@@ -165,13 +178,33 @@ class NBVReconstructionDataset(BaseDataset):
             [best_to_world_6d, best_to_world_trans], axis=0
         )
 
+        start_time = time.time()
+        
         combined_scanned_views_pts = np.concatenate(scanned_views_pts, axis=0)
-        voxel_downsampled_combined_scanned_pts_np = PtsUtil.voxel_downsample_point_cloud(combined_scanned_views_pts, 0.002)
-        random_downsampled_combined_scanned_pts_np = PtsUtil.random_downsample_point_cloud(voxel_downsampled_combined_scanned_pts_np, self.pts_num)
+        #Log.info(f"combined_scanned_views_pts shape: {combined_scanned_views_pts.shape}")
+        voxel_downsampled_combined_scanned_pts_np, inverse = self.voxel_downsample_with_mapping(combined_scanned_views_pts, 0.003)
+        random_downsampled_combined_scanned_pts_np, random_downsample_idx = PtsUtil.random_downsample_point_cloud(voxel_downsampled_combined_scanned_pts_np, self.pts_num, 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)
+        #Log.info(f"random_downsampled_combined_scanned_pts_np shape: {random_downsampled_combined_scanned_pts_np.shape}")
+        end_time = time.time()
+        #Log.info(f"downsample time: {end_time - start_time}")
         
         data_item = {
             "scanned_pts": np.asarray(scanned_views_pts, dtype=np.float32), # Ndarray(S x Nv x 3)
             "combined_scanned_pts": np.asarray(random_downsampled_combined_scanned_pts_np, dtype=np.float32), # Ndarray(N x 3)
+            "scanned_pts_mask": np.asarray(scanned_pts_mask, dtype=np.bool), # Ndarray(N)
             "scanned_coverage_rate": scanned_coverages_rate, # List(S): Float, range(0, 1)
             "scanned_n_to_world_pose_9d": np.asarray(scanned_n_to_world_pose, dtype=np.float32), # Ndarray(S x 9)
             "best_coverage_rate": nbv_coverage_rate, # Float, range(0, 1)
@@ -197,7 +230,9 @@ class NBVReconstructionDataset(BaseDataset):
             collate_data["scanned_n_to_world_pose_9d"] = [
                 torch.tensor(item["scanned_n_to_world_pose_9d"]) for item in batch
             ]
-            
+            collate_data["scanned_pts_mask"] = [
+                torch.tensor(item["scanned_pts_mask"]) for item in batch
+            ]
             ''' ------ Fixed Length ------ '''
             
             collate_data["best_to_world_pose_9d"] = torch.stack(
@@ -206,17 +241,14 @@ class NBVReconstructionDataset(BaseDataset):
             collate_data["combined_scanned_pts"] = torch.stack(
                 [torch.tensor(item["combined_scanned_pts"]) for item in batch]
             )
-            collate_data["scanned_pts_mask"] = torch.stack(
-                [torch.tensor(item["scanned_pts_mask"]) for item in batch]
-            )
-                
+            
             for key in batch[0].keys():
                 if key not in [
                     "scanned_pts",
-                    "scanned_pts_mask",
                     "scanned_n_to_world_pose_9d",
                     "best_to_world_pose_9d",
                     "combined_scanned_pts",
+                    "scanned_pts_mask",
                 ]:
                     collate_data[key] = [item[key] for item in batch]
             return collate_data
@@ -232,9 +264,9 @@ if __name__ == "__main__":
     torch.manual_seed(seed)
     np.random.seed(seed)
     config = {
-        "root_dir": "/data/hofee/data/packed_preprocessed_data",
+        "root_dir": "/data/hofee/nbv_rec_part2_preprocessed",
         "source": "nbv_reconstruction_dataset",
-        "split_file": "/data/hofee/data/OmniObject3d_train.txt",
+        "split_file": "/data/hofee/data/sample.txt",
         "load_from_preprocess": True,
         "ratio": 0.5,
         "batch_size": 2,