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
Show Stats 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":

10
core/evaluation.py
View File

@ -20,7 +20,7 @@ class PoseDiff:
rot_angle_list = []
trans_dist_list = []
for output, data in zip(output_list, data_list):
gt_pose_9d = data['best_to_1_pose_9d']
gt_pose_9d = data['best_to_world_pose_9d']
pred_pose_9d = output['pred_pose_9d']
gt_rot_6d = gt_pose_9d[:, :6]
gt_trans = gt_pose_9d[:, 6:]
@ -62,10 +62,10 @@ class ConverageRateIncrease:
voxel_threshold_list = data["voxel_threshold"]
filter_degree_list = data["filter_degree"]
first_frame_to_world = data["first_frame_to_world"]
pred_n_to_1_pose_mats = torch.eye(4, device=pred_pose_9ds.device).unsqueeze(0).repeat(pred_pose_9ds.shape[0], 1, 1)
pred_n_to_1_pose_mats[:,:3,:3] = PoseUtil.rotation_6d_to_matrix_tensor_batch(pred_pose_9ds[:, :6])
pred_n_to_1_pose_mats[:,:3,3] = pred_pose_9ds[:, 6:]
pred_n_to_world_pose_mats = torch.matmul(first_frame_to_world, pred_n_to_1_pose_mats)
pred_n_to_world_pose_mats = torch.eye(4, device=pred_pose_9ds.device).unsqueeze(0).repeat(pred_pose_9ds.shape[0], 1, 1)
pred_n_to_world_pose_mats[:,:3,:3] = PoseUtil.rotation_6d_to_matrix_tensor_batch(pred_pose_9ds[:, :6])
pred_n_to_world_pose_mats[:,:3,3] = pred_pose_9ds[:, 6:]
pred_n_to_world_pose_mats = torch.matmul(first_frame_to_world, pred_n_to_world_pose_mats)
for idx in range(len(scanned_cr)):
model_points_normals = model_points_normals_list[idx]
scanned_view_pts = scanned_view_pts_list[idx]

12
core/pipeline.py
View File

@ -40,8 +40,8 @@ class NBVReconstructionPipeline(nn.Module):
def forward_train(self, data):
seq_feat = self.get_seq_feat(data)
''' get std '''
best_to_1_pose_9d_batch = data["best_to_1_pose_9d"]
perturbed_x, random_t, target_score, std = self.pertube_data(best_to_1_pose_9d_batch)
best_to_world_pose_9d_batch = data["best_to_world_pose_9d"]
perturbed_x, random_t, target_score, std = self.pertube_data(best_to_world_pose_9d_batch)
input_data = {
"sampled_pose": perturbed_x,
"t": random_t,
@ -66,16 +66,16 @@ class NBVReconstructionPipeline(nn.Module):
def get_seq_feat(self, data):
scanned_pts_batch = data['scanned_pts']
scanned_n_to_1_pose_9d_batch = data['scanned_n_to_1_pose_9d']
scanned_n_to_world_pose_9d_batch = data['scanned_n_to_world_pose_9d']
pts_feat_seq_list = []
pose_feat_seq_list = []
device = next(self.parameters()).device
for scanned_pts,scanned_n_to_1_pose_9d in zip(scanned_pts_batch,scanned_n_to_1_pose_9d_batch):
for scanned_pts,scanned_n_to_world_pose_9d in zip(scanned_pts_batch,scanned_n_to_world_pose_9d_batch):
scanned_pts = scanned_pts.to(device)
scanned_n_to_1_pose_9d = scanned_n_to_1_pose_9d.to(device)
scanned_n_to_world_pose_9d = scanned_n_to_world_pose_9d.to(device)
pts_feat_seq_list.append(self.pts_encoder.encode_points(scanned_pts))
pose_feat_seq_list.append(self.pose_encoder.encode_pose(scanned_n_to_1_pose_9d))
pose_feat_seq_list.append(self.pose_encoder.encode_pose(scanned_n_to_world_pose_9d))
seq_feat = self.seq_encoder.encode_sequence(pts_feat_seq_list, pose_feat_seq_list)
if torch.isnan(seq_feat).any():