local_only: dataset

This commit is contained in:
hofee 2024-10-29 12:40:06 +00:00
parent 234c8bccc3
commit 8d7299b482

View File

@ -165,13 +165,8 @@ class NBVReconstructionDataset(BaseDataset):
[best_to_world_6d, best_to_world_trans], axis=0 [best_to_world_6d, best_to_world_trans], axis=0
) )
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)
data_item = { data_item = {
"scanned_pts": np.asarray(scanned_views_pts, dtype=np.float32), # Ndarray(S x Nv x 3) "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_coverage_rate": scanned_coverages_rate, # List(S): Float, range(0, 1) "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) "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) "best_coverage_rate": nbv_coverage_rate, # Float, range(0, 1)
@ -203,9 +198,6 @@ class NBVReconstructionDataset(BaseDataset):
collate_data["best_to_world_pose_9d"] = torch.stack( collate_data["best_to_world_pose_9d"] = torch.stack(
[torch.tensor(item["best_to_world_pose_9d"]) for item in batch] [torch.tensor(item["best_to_world_pose_9d"]) for item in batch]
) )
collate_data["combined_scanned_pts"] = torch.stack(
[torch.tensor(item["combined_scanned_pts"]) for item in batch]
)
collate_data["scanned_pts_mask"] = torch.stack( collate_data["scanned_pts_mask"] = torch.stack(
[torch.tensor(item["scanned_pts_mask"]) for item in batch] [torch.tensor(item["scanned_pts_mask"]) for item in batch]
) )
@ -216,7 +208,6 @@ class NBVReconstructionDataset(BaseDataset):
"scanned_pts_mask", "scanned_pts_mask",
"scanned_n_to_world_pose_9d", "scanned_n_to_world_pose_9d",
"best_to_world_pose_9d", "best_to_world_pose_9d",
"combined_scanned_pts",
]: ]:
collate_data[key] = [item[key] for item in batch] collate_data[key] = [item[key] for item in batch]
return collate_data return collate_data