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

finish PoseDiff and NBVDataset

Browse Source
This commit is contained in:
hofee 2024-08-30 19:21:18 +08:00
parent be5a2d57fa
commit f58360c0c0
5 changed files with 57 additions and 33 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

65
core/dataset.py
View File

@ -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.data_load import DataLoadUtil
from utils.pose import PoseUtil 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): class NBVReconstructionDataset(BaseDataset):
def __init__(self, config): def __init__(self, config):
super(NBVReconstructionDataset, self).__init__(config) super(NBVReconstructionDataset, self).__init__(config)
self.config = config self.config = config
self.root_dir = config["root_dir"] self.root_dir = config["root_dir"]
self.datalist = self.get_datalist() self.datalist = self.get_datalist()
self.pts_num = 1024
def get_datalist(self): def get_datalist(self):
datalist = [] datalist = []
@ -43,9 +45,9 @@ class NBVReconstructionDataset(BaseDataset):
nbv = data_item_info["next_best_view"] nbv = data_item_info["next_best_view"]
max_coverage_rate = data_item_info["max_coverage_rate"] max_coverage_rate = data_item_info["max_coverage_rate"]
scene_name = data_item_info["scene_name"] 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_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: for view in scanned_views:
frame_idx = view[0] frame_idx = view[0]
coverage_rate = view[1] coverage_rate = view[1]
@ -53,31 +55,37 @@ class NBVReconstructionDataset(BaseDataset):
depth = DataLoadUtil.load_depth(view_path) depth = DataLoadUtil.load_depth(view_path)
cam_info = DataLoadUtil.load_cam_info(view_path) cam_info = DataLoadUtil.load_cam_info(view_path)
mask = DataLoadUtil.load_seg(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) frame_curr_to_world = cam_info["cam_to_world"]
scanned_views_pts.append(target_point_cloud) 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) scanned_coverages_rate.append(coverage_rate)
cam_pose = DataLoadUtil.load_cam_info(view_path)["cam_to_world"] 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]
cam_pose_6d = PoseUtil.matrix_to_rotation_6d_numpy(np.asarray(cam_pose[:3,:3])) n_to_1_9d = np.concatenate([n_to_1_6d, n_to_1_trans], axis=0)
translation = cam_pose[:3,3] scanned_n_to_1_pose.append(n_to_1_9d)
cam_pose_9d = np.concatenate([cam_pose_6d, translation], axis=0)
scanned_cam_pose.append(cam_pose_9d)
nbv_idx, nbv_coverage_rate = nbv[0], nbv[1] nbv_idx, nbv_coverage_rate = nbv[0], nbv[1]
nbv_path = DataLoadUtil.get_path(self.root_dir, scene_name, nbv_idx) 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) cam_info = DataLoadUtil.load_cam_info(nbv_path)
nbv_cam_pose = cam_info["cam_to_world"] nbv_mask = DataLoadUtil.load_seg(nbv_path)
nbv_cam_pose_6d = PoseUtil.matrix_to_rotation_6d_numpy(np.asarray(nbv_cam_pose[:3,:3])) best_frame_to_world = cam_info["cam_to_world"]
translation = nbv_cam_pose[:3,3] best_to_1_pose = np.dot(np.linalg.inv(first_frame_to_world), best_frame_to_world)
nbv_cam_pose_9d = np.concatenate([nbv_cam_pose_6d, translation], axis=0) 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 = { data_item = {
"scanned_views_pts": np.asarray(scanned_views_pts,dtype=np.float32), "scanned_pts": np.asarray(scanned_views_pts,dtype=np.float32),
"scanned_coverages_rate": np.asarray(scanned_coverages_rate,dtype=np.float32), "scanned_coverage_rate": np.asarray(scanned_coverages_rate,dtype=np.float32),
"scanned_cam_pose": np.asarray(scanned_cam_pose,dtype=np.float32), "scanned_n_to_1_pose_9d": np.asarray(scanned_n_to_1_pose,dtype=np.float32),
"nbv_pts": np.asarray(nbv_pts,dtype=np.float32), "best_pts": np.asarray(downsampled_best_target_point_cloud,dtype=np.float32),
"nbv_coverage_rate": nbv_coverage_rate, "best_coverage_rate": nbv_coverage_rate,
"nbv_cam_pose": nbv_cam_pose_9d, "best_to_1_pose_9d": best_to_1_9d,
"max_coverage_rate": max_coverage_rate, "max_coverage_rate": max_coverage_rate,
"scene_name": scene_name "scene_name": scene_name
} }
@ -91,7 +99,7 @@ if __name__ == "__main__":
import torch import torch
config = { config = {
"root_dir": "/media/hofee/data/data/nbv_rec/sample", "root_dir": "/media/hofee/data/data/nbv_rec/sample",
"ratio": 0.1, "ratio": 0.05,
"batch_size": 1, "batch_size": 1,
"num_workers": 0, "num_workers": 0,
} }
@ -99,7 +107,18 @@ if __name__ == "__main__":
print(len(ds)) print(len(ds))
dl = ds.get_loader(shuffle=True) dl = ds.get_loader(shuffle=True)
for idx, data in enumerate(dl): 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(): for key, value in data.items():
if isinstance(value, torch.Tensor): if isinstance(value, torch.Tensor):
print(key, ":" ,value.shape) print(key, ":" ,value.shape)
print() print()

8
core/evaluation.py
View File

@ -6,7 +6,7 @@ import PytorchBoot.namespace as namespace
def get_view_data(cam_pose, scene_name): def get_view_data(cam_pose, scene_name):
pass pass
@stereotype.evaluation_method("pose_diff", comment="not tested") @stereotype.evaluation_method("pose_diff")
class PoseDiff: class PoseDiff:
def __init__(self, _): def __init__(self, _):
pass pass
@ -16,7 +16,7 @@ class PoseDiff:
rot_angle_list = [] rot_angle_list = []
trans_dist_list = [] trans_dist_list = []
for output, data in zip(output_list, data_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'] pred_pose_9d = output['pred_pose_9d']
gt_rot_6d = gt_pose_9d[:, :6] gt_rot_6d = gt_pose_9d[:, :6]
gt_trans = gt_pose_9d[:, 6:] gt_trans = gt_pose_9d[:, 6:]
@ -49,9 +49,9 @@ class ConverageRateIncrease:
cr_diff_list = [] cr_diff_list = []
for output, data in zip(output_list, data_list): for output, data in zip(output_list, data_list):
scanned_cr = data['scanned_coverages_rate'] scanned_cr = data['scanned_coverages_rate']
gt_cr = data["nbv_coverage_rate"] gt_cr = data["best_coverage_rate"]
scene_name_list = data['scene_name'] 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_pose_9ds = output['pred_pose_9d']
pred_rot_mats = PoseUtil.rotation_6d_to_matrix_tensor_batch(pred_pose_9ds[:, :6]) 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) pred_pose_mats = torch.cat([pred_rot_mats, pred_pose_9ds[:, 6:]], dim=-1)

2
utils/data_load.py
View File

@ -129,12 +129,12 @@ class DataLoadUtil:
target_points_camera_aug = np.concatenate([target_points_camera, np.ones((target_points_camera.shape[0], 1))], axis=-1) 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] target_points_world = np.dot(cam_extrinsic, target_points_camera_aug.T).T[:, :3]
return { return {
"points_world": target_points_world, "points_world": target_points_world,
"points_camera": target_points_camera "points_camera": target_points_camera
} }
@staticmethod @staticmethod
def get_point_cloud_world_from_path(path): def get_point_cloud_world_from_path(path):
cam_info = DataLoadUtil.load_cam_info(path) cam_info = DataLoadUtil.load_cam_info(path)

5
utils/pts.py
View File

@ -15,3 +15,8 @@ class PtsUtil:
points_h = np.concatenate([points, np.ones((points.shape[0], 1))], axis=1) points_h = np.concatenate([points, np.ones((points.shape[0], 1))], axis=1)
points_h = np.dot(pose_mat, points_h.T).T points_h = np.dot(pose_mat, points_h.T).T
return points_h[:, :3] 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]