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

change to-1 to to_world

Browse Source
This commit is contained in:
hofee
2024-09-19 00:20:26 +08:00
parent 935069d68c
commit 8d5d6d5df4
3 changed files with 29 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files

38
core/dataset.py
View File

@@ -92,7 +92,7 @@ class NBVReconstructionDataset(BaseDataset):
nbv = data_item_info["next_best_view"]
max_coverage_rate = data_item_info["max_coverage_rate"]
scene_name = data_item_info["scene_name"]
scanned_views_pts, scanned_coverages_rate, scanned_n_to_1_pose = [], [], []
scanned_views_pts, scanned_coverages_rate, scanned_n_to_world_pose = [], [], []
first_frame_idx = scanned_views[0][0]
first_cam_info = DataLoadUtil.load_cam_info(DataLoadUtil.get_path(self.root_dir, scene_name, first_frame_idx), binocular=True)
first_frame_to_world = first_cam_info["cam_to_world"]
@@ -103,17 +103,15 @@ class NBVReconstructionDataset(BaseDataset):
cam_info = DataLoadUtil.load_cam_info(view_path, binocular=True)
n_to_world_pose = cam_info["cam_to_world"]
nR_to_world_pose = cam_info["cam_to_world_R"]
n_to_1_pose = np.dot(np.linalg.inv(first_frame_to_world), n_to_world_pose)
nR_to_1_pose = np.dot(np.linalg.inv(first_frame_to_world), nR_to_world_pose)
cached_data = None
if self.cache:
cached_data = self.load_from_cache(scene_name, first_frame_idx, frame_idx)
if cached_data is None:
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_1_pose)['points_world']
point_cloud_R = DataLoadUtil.get_point_cloud(depth_R, cam_info['cam_intrinsic'], nR_to_1_pose)['points_world']
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)
@@ -126,26 +124,26 @@ class NBVReconstructionDataset(BaseDataset):
scanned_views_pts.append(downsampled_target_point_cloud)
scanned_coverages_rate.append(coverage_rate)
n_to_1_6d = PoseUtil.matrix_to_rotation_6d_numpy(np.asarray(n_to_1_pose[:3,:3]))
n_to_1_trans = n_to_1_pose[:3,3]
n_to_1_9d = np.concatenate([n_to_1_6d, n_to_1_trans], axis=0)
scanned_n_to_1_pose.append(n_to_1_9d)
n_to_world_6d = PoseUtil.matrix_to_rotation_6d_numpy(np.asarray(n_to_world_pose[:3,:3]))
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)
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)
best_frame_to_world = cam_info["cam_to_world"]
best_to_1_pose = np.dot(np.linalg.inv(first_frame_to_world), best_frame_to_world)
best_to_1_6d = PoseUtil.matrix_to_rotation_6d_numpy(np.asarray(best_to_1_pose[:3,:3]))
best_to_1_trans = best_to_1_pose[:3,3]
best_to_1_9d = np.concatenate([best_to_1_6d, best_to_1_trans], axis=0)
best_to_world_6d = PoseUtil.matrix_to_rotation_6d_numpy(np.asarray(best_frame_to_world[:3,:3]))
best_to_world_trans = best_frame_to_world[:3,3]
best_to_world_9d = np.concatenate([best_to_world_6d, best_to_world_trans], axis=0)
data_item = {
"scanned_pts": np.asarray(scanned_views_pts,dtype=np.float32),
"scanned_coverage_rate": scanned_coverages_rate,
"scanned_n_to_1_pose_9d": np.asarray(scanned_n_to_1_pose,dtype=np.float32),
"scanned_n_to_world_pose_9d": np.asarray(scanned_n_to_world_pose,dtype=np.float32),
"best_coverage_rate": nbv_coverage_rate,
"best_to_1_pose_9d": np.asarray(best_to_1_9d,dtype=np.float32),
"best_to_world_pose_9d": np.asarray(best_to_world_9d,dtype=np.float32),
"max_coverage_rate": max_coverage_rate,
"scene_name": scene_name
}
@@ -180,12 +178,12 @@ class NBVReconstructionDataset(BaseDataset):
def collate_fn(batch):
collate_data = {}
collate_data["scanned_pts"] = [torch.tensor(item['scanned_pts']) for item in batch]
collate_data["scanned_n_to_1_pose_9d"] = [torch.tensor(item['scanned_n_to_1_pose_9d']) for item in batch]
collate_data["best_to_1_pose_9d"] = torch.stack([torch.tensor(item['best_to_1_pose_9d']) 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]
collate_data["best_to_world_pose_9d"] = torch.stack([torch.tensor(item['best_to_world_pose_9d']) for item in batch])
if "first_frame_to_world" in batch[0]:
collate_data["first_frame_to_world"] = torch.stack([torch.tensor(item["first_frame_to_world"]) for item in batch])
for key in batch[0].keys():
if key not in ["scanned_pts", "scanned_n_to_1_pose_9d", "best_to_1_pose_9d", "first_frame_to_world"]:
if key not in ["scanned_pts", "scanned_n_to_world_pose_9d", "best_to_world_pose_9d", "first_frame_to_world"]:
collate_data[key] = [item[key] for item in batch]
return collate_data
return collate_fn
@@ -233,7 +231,7 @@ if __name__ == "__main__":
# print(key, ":" ,value.shape)
# else:
# print(key, ":" ,len(value))
# if key == "scanned_n_to_1_pose_9d":
# if key == "scanned_n_to_world_pose_9d":
# for val in value:
# print(val.shape)
# if key == "scanned_pts":