hofee/nbv_reconstruction
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add scanned_pts_mask

Browse Source
This commit is contained in:
hofee 2024-10-06 12:01:10 +08:00
parent a84417ef62
commit 276f45dcc3
1 changed files with 17 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

25
core/nbv_dataset.py
View File

@ -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",