diff --git a/core/dataset.py b/core/dataset.py index 679c4a9..9abea77 100644 --- a/core/dataset.py +++ b/core/dataset.py @@ -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": diff --git a/core/evaluation.py b/core/evaluation.py index 97793fe..4ad1c8e 100644 --- a/core/evaluation.py +++ b/core/evaluation.py @@ -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] diff --git a/core/pipeline.py b/core/pipeline.py index 5066dcd..294ed50 100644 --- a/core/pipeline.py +++ b/core/pipeline.py @@ -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():