From 276f45dcc369e06e10bc957bec0eb075fdfe0172 Mon Sep 17 00:00:00 2001
From: hofee <64160135+GitHofee@users.noreply.github.com>
Date: Sun, 6 Oct 2024 12:01:10 +0800
Subject: [PATCH] add scanned_pts_mask

---
 core/nbv_dataset.py | 25 +++++++++++++++++--------
 1 file changed, 17 insertions(+), 8 deletions(-)

diff --git a/core/nbv_dataset.py b/core/nbv_dataset.py
index 702134e..3880ed5 100644
--- a/core/nbv_dataset.py
+++ b/core/nbv_dataset.py
@@ -162,18 +162,23 @@ class NBVReconstructionDataset(BaseDataset):
         )
 
         combined_scanned_views_pts = np.concatenate(scanned_views_pts, axis=0)
-        voxel_downsampled_combined_scanned_pts_np, _ = (
-            PtsUtil.voxelize_points(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
-            )
+        fps_downsampled_combined_scanned_pts, fps_mask = PtsUtil.fps_downsample_point_cloud(
+            combined_scanned_views_pts, self.pts_num, require_mask=True
         )
+        
+        view_start_indices = np.cumsum([0] + [pts.shape[0] for pts in scanned_views_pts[:-1]])
+        scanned_pts_mask = []
+
+        for i, start_idx in enumerate(view_start_indices[:-1]):
+            end_idx = view_start_indices[i + 1]
+            view_mask = fps_mask[start_idx:end_idx]
+            scanned_pts_mask.append(view_mask)
+        
         data_item = {
             "scanned_pts": np.asarray(scanned_views_pts, dtype=np.float32),
+            "scanned_pts_mask": np.asarray(scanned_pts_mask, dtype=np.uint8),
             "combined_scanned_pts": np.asarray(
-                random_downsampled_combined_scanned_pts_np, dtype=np.float32
+                fps_downsampled_combined_scanned_pts, dtype=np.float32
             ),
             "scanned_coverage_rate": scanned_coverages_rate,
             "scanned_n_to_world_pose_9d": np.asarray(
@@ -199,6 +204,9 @@ class NBVReconstructionDataset(BaseDataset):
             collate_data["scanned_pts"] = [
                 torch.tensor(item["scanned_pts"]) for item in batch
             ]
+            collate_data["scanned_pts_mask"] = [
+                torch.tensor(item["scanned_pts_mask"]) for item in batch
+            ]
             collate_data["scanned_n_to_world_pose_9d"] = [
                 torch.tensor(item["scanned_n_to_world_pose_9d"]) for item in batch
             ]
@@ -218,6 +226,7 @@ class NBVReconstructionDataset(BaseDataset):
             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",
                     "first_frame_to_world",