diff --git a/app_generate.py b/app_generate.py index 9dada25..d56a6d1 100644 --- a/app_generate.py +++ b/app_generate.py @@ -6,4 +6,4 @@ class GenerateApp: @staticmethod def start(): StrategyGenerator("configs/strategy_generate_config.yaml").run() - \ No newline at end of file + \ No newline at end of file diff --git a/core/dataset.py b/core/dataset.py index 7d39e67..4dcc5de 100644 --- a/core/dataset.py +++ b/core/dataset.py @@ -8,14 +8,16 @@ sys.path.append(r"/media/hofee/data/project/python/nbv_reconstruction/nbv_recons from utils.data_load import DataLoadUtil from utils.pose import PoseUtil +from utils.pts import PtsUtil -@stereotype.dataset("nbv_reconstruction_dataset", comment="not tested") +@stereotype.dataset("nbv_reconstruction_dataset") class NBVReconstructionDataset(BaseDataset): def __init__(self, config): super(NBVReconstructionDataset, self).__init__(config) self.config = config self.root_dir = config["root_dir"] self.datalist = self.get_datalist() + self.pts_num = 1024 def get_datalist(self): datalist = [] @@ -43,9 +45,9 @@ 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_cam_pose = [], [], [] + scanned_views_pts, scanned_coverages_rate, scanned_n_to_1_pose = [], [], [] first_frame_idx = scanned_views[0][0] - first_frame_pose = DataLoadUtil.load_cam_info(DataLoadUtil.get_path(self.root_dir, scene_name, first_frame_idx))["cam_to_world"] + first_frame_to_world = DataLoadUtil.load_cam_info(DataLoadUtil.get_path(self.root_dir, scene_name, first_frame_idx))["cam_to_world"] for view in scanned_views: frame_idx = view[0] coverage_rate = view[1] @@ -53,35 +55,41 @@ class NBVReconstructionDataset(BaseDataset): depth = DataLoadUtil.load_depth(view_path) cam_info = DataLoadUtil.load_cam_info(view_path) mask = DataLoadUtil.load_seg(view_path) - target_point_cloud = DataLoadUtil.get_target_point_cloud(depth, cam_info["cam_intrinsic"], cam_info["cam_to_world"], mask) - scanned_views_pts.append(target_point_cloud) - scanned_coverages_rate.append(coverage_rate) - cam_pose = DataLoadUtil.load_cam_info(view_path)["cam_to_world"] - - cam_pose_6d = PoseUtil.matrix_to_rotation_6d_numpy(np.asarray(cam_pose[:3,:3])) - translation = cam_pose[:3,3] - cam_pose_9d = np.concatenate([cam_pose_6d, translation], axis=0) - scanned_cam_pose.append(cam_pose_9d) + frame_curr_to_world = cam_info["cam_to_world"] + n_to_1_pose = np.dot(np.linalg.inv(first_frame_to_world), frame_curr_to_world) + target_point_cloud = DataLoadUtil.get_target_point_cloud(depth, cam_info["cam_intrinsic"], n_to_1_pose, mask)["points_world"] + downsampled_target_point_cloud = PtsUtil.random_downsample_point_cloud(target_point_cloud, self.pts_num) + 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) nbv_idx, nbv_coverage_rate = nbv[0], nbv[1] nbv_path = DataLoadUtil.get_path(self.root_dir, scene_name, nbv_idx) - nbv_pts = DataLoadUtil.load_depth(nbv_path) + nbv_depth = DataLoadUtil.load_depth(nbv_path) cam_info = DataLoadUtil.load_cam_info(nbv_path) - nbv_cam_pose = cam_info["cam_to_world"] - nbv_cam_pose_6d = PoseUtil.matrix_to_rotation_6d_numpy(np.asarray(nbv_cam_pose[:3,:3])) - translation = nbv_cam_pose[:3,3] - nbv_cam_pose_9d = np.concatenate([nbv_cam_pose_6d, translation], axis=0) + nbv_mask = DataLoadUtil.load_seg(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_target_point_cloud = DataLoadUtil.get_target_point_cloud(nbv_depth, cam_info["cam_intrinsic"], best_to_1_pose, nbv_mask)["points_world"] + downsampled_best_target_point_cloud = PtsUtil.random_downsample_point_cloud(best_target_point_cloud, self.pts_num) + 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) + data_item = { - "scanned_views_pts": np.asarray(scanned_views_pts,dtype=np.float32), - "scanned_coverages_rate": np.asarray(scanned_coverages_rate,dtype=np.float32), - "scanned_cam_pose": np.asarray(scanned_cam_pose,dtype=np.float32), - "nbv_pts": np.asarray(nbv_pts,dtype=np.float32), - "nbv_coverage_rate": nbv_coverage_rate, - "nbv_cam_pose": nbv_cam_pose_9d, + "scanned_pts": np.asarray(scanned_views_pts,dtype=np.float32), + "scanned_coverage_rate": np.asarray(scanned_coverages_rate,dtype=np.float32), + "scanned_n_to_1_pose_9d": np.asarray(scanned_n_to_1_pose,dtype=np.float32), + "best_pts": np.asarray(downsampled_best_target_point_cloud,dtype=np.float32), + "best_coverage_rate": nbv_coverage_rate, + "best_to_1_pose_9d": best_to_1_9d, "max_coverage_rate": max_coverage_rate, "scene_name": scene_name } - + return data_item def __len__(self): @@ -91,7 +99,7 @@ if __name__ == "__main__": import torch config = { "root_dir": "/media/hofee/data/data/nbv_rec/sample", - "ratio": 0.1, + "ratio": 0.05, "batch_size": 1, "num_workers": 0, } @@ -99,7 +107,18 @@ if __name__ == "__main__": print(len(ds)) dl = ds.get_loader(shuffle=True) for idx, data in enumerate(dl): + cnt=0 + print(data["scene_name"]) + print(data["scanned_coverage_rate"]) + print(data["best_coverage_rate"]) + for pts in data["scanned_pts"][0]: + #np.savetxt(f"pts_{cnt}.txt", pts) + cnt+=1 + best_pts = data["best_pts"][0] + #np.savetxt("best_pts.txt", best_pts) for key, value in data.items(): if isinstance(value, torch.Tensor): print(key, ":" ,value.shape) + + print() \ No newline at end of file diff --git a/core/evaluation.py b/core/evaluation.py index c304bee..d5b20e2 100644 --- a/core/evaluation.py +++ b/core/evaluation.py @@ -6,7 +6,7 @@ import PytorchBoot.namespace as namespace def get_view_data(cam_pose, scene_name): pass -@stereotype.evaluation_method("pose_diff", comment="not tested") +@stereotype.evaluation_method("pose_diff") class PoseDiff: def __init__(self, _): pass @@ -16,7 +16,7 @@ class PoseDiff: rot_angle_list = [] trans_dist_list = [] for output, data in zip(output_list, data_list): - gt_pose_9d = data['nbv_cam_pose'] + gt_pose_9d = data['best_to_1_pose_9d'] pred_pose_9d = output['pred_pose_9d'] gt_rot_6d = gt_pose_9d[:, :6] gt_trans = gt_pose_9d[:, 6:] @@ -49,9 +49,9 @@ class ConverageRateIncrease: cr_diff_list = [] for output, data in zip(output_list, data_list): scanned_cr = data['scanned_coverages_rate'] - gt_cr = data["nbv_coverage_rate"] + gt_cr = data["best_coverage_rate"] scene_name_list = data['scene_name'] - scanned_view_pts_list = data['scanned_views_pts'] + scanned_view_pts_list = data['scanned_pts'] pred_pose_9ds = output['pred_pose_9d'] pred_rot_mats = PoseUtil.rotation_6d_to_matrix_tensor_batch(pred_pose_9ds[:, :6]) pred_pose_mats = torch.cat([pred_rot_mats, pred_pose_9ds[:, 6:]], dim=-1) diff --git a/utils/data_load.py b/utils/data_load.py index 08db804..c43cfe6 100644 --- a/utils/data_load.py +++ b/utils/data_load.py @@ -129,11 +129,11 @@ class DataLoadUtil: target_points_camera_aug = np.concatenate([target_points_camera, np.ones((target_points_camera.shape[0], 1))], axis=-1) target_points_world = np.dot(cam_extrinsic, target_points_camera_aug.T).T[:, :3] - return { "points_world": target_points_world, "points_camera": target_points_camera } + @staticmethod def get_point_cloud_world_from_path(path): diff --git a/utils/pts.py b/utils/pts.py index 3d3de90..c32c92a 100644 --- a/utils/pts.py +++ b/utils/pts.py @@ -2,7 +2,7 @@ import numpy as np import open3d as o3d class PtsUtil: - + @staticmethod def voxel_downsample_point_cloud(point_cloud, voxel_size=0.005): o3d_pc = o3d.geometry.PointCloud() @@ -14,4 +14,9 @@ class PtsUtil: def transform_point_cloud(points, pose_mat): points_h = np.concatenate([points, np.ones((points.shape[0], 1))], axis=1) points_h = np.dot(pose_mat, points_h.T).T - return points_h[:, :3] \ No newline at end of file + return points_h[:, :3] + + @staticmethod + def random_downsample_point_cloud(point_cloud, num_points): + idx = np.random.choice(len(point_cloud), num_points, replace=False) + return point_cloud[idx] \ No newline at end of file