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

add target_pts_num into dataset

Browse Source
This commit is contained in:
hofee 2024-09-29 18:11:55 +08:00
parent cb983fdc74
commit 99e57c3f4c
2 changed files with 144 additions and 114 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

227
core/nbv_dataset.py
View File

@ -7,6 +7,7 @@ from PytorchBoot.utils.log_util import Log
import torch import torch
import os import os
import sys import sys
sys.path.append(r"/home/data/hofee/project/nbv_rec/nbv_reconstruction") sys.path.append(r"/home/data/hofee/project/nbv_rec/nbv_reconstruction")
from utils.data_load import DataLoadUtil from utils.data_load import DataLoadUtil
@ -29,20 +30,17 @@ class NBVReconstructionDataset(BaseDataset):
self.cache = config.get("cache") self.cache = config.get("cache")
self.load_from_preprocess = config.get("load_from_preprocess", False) self.load_from_preprocess = config.get("load_from_preprocess", False)
if self.type == namespace.Mode.TEST: if self.type == namespace.Mode.TEST:
self.model_dir = config["model_dir"] self.model_dir = config["model_dir"]
self.filter_degree = config["filter_degree"] self.filter_degree = config["filter_degree"]
if self.type == namespace.Mode.TRAIN: if self.type == namespace.Mode.TRAIN:
scale_ratio = 100 scale_ratio = 100
self.datalist = self.datalist*scale_ratio self.datalist = self.datalist * scale_ratio
if self.cache: if self.cache:
expr_root = ConfigManager.get("runner", "experiment", "root_dir") expr_root = ConfigManager.get("runner", "experiment", "root_dir")
expr_name = ConfigManager.get("runner", "experiment", "name") expr_name = ConfigManager.get("runner", "experiment", "name")
self.cache_dir = os.path.join(expr_root, expr_name, "cache") self.cache_dir = os.path.join(expr_root, expr_name, "cache")
#self.preprocess_cache() # self.preprocess_cache()
def load_scene_name_list(self): def load_scene_name_list(self):
scene_name_list = [] scene_name_list = []
@ -60,7 +58,9 @@ class NBVReconstructionDataset(BaseDataset):
max_coverage_rate_list = [] max_coverage_rate_list = []
for seq_idx in range(seq_num): for seq_idx in range(seq_num):
label_path = DataLoadUtil.get_label_path(self.root_dir, scene_name, seq_idx) label_path = DataLoadUtil.get_label_path(
self.root_dir, scene_name, seq_idx
)
label_data = DataLoadUtil.load_label(label_path) label_data = DataLoadUtil.load_label(label_path)
max_coverage_rate = label_data["max_coverage_rate"] max_coverage_rate = label_data["max_coverage_rate"]
if max_coverage_rate > scene_max_coverage_rate: if max_coverage_rate > scene_max_coverage_rate:
@ -69,20 +69,24 @@ class NBVReconstructionDataset(BaseDataset):
mean_coverage_rate = np.mean(max_coverage_rate_list) mean_coverage_rate = np.mean(max_coverage_rate_list)
for seq_idx in range(seq_num): for seq_idx in range(seq_num):
label_path = DataLoadUtil.get_label_path(self.root_dir, scene_name, seq_idx) label_path = DataLoadUtil.get_label_path(
self.root_dir, scene_name, seq_idx
)
label_data = DataLoadUtil.load_label(label_path) label_data = DataLoadUtil.load_label(label_path)
if max_coverage_rate_list[seq_idx] > mean_coverage_rate - 0.1: if max_coverage_rate_list[seq_idx] > mean_coverage_rate - 0.1:
for data_pair in label_data["data_pairs"]: for data_pair in label_data["data_pairs"]:
scanned_views = data_pair[0] scanned_views = data_pair[0]
next_best_view = data_pair[1] next_best_view = data_pair[1]
datalist.append({ datalist.append(
"scanned_views": scanned_views, {
"next_best_view": next_best_view, "scanned_views": scanned_views,
"seq_max_coverage_rate": max_coverage_rate, "next_best_view": next_best_view,
"scene_name": scene_name, "seq_max_coverage_rate": max_coverage_rate,
"label_idx": seq_idx, "scene_name": scene_name,
"scene_max_coverage_rate": scene_max_coverage_rate "label_idx": seq_idx,
}) "scene_max_coverage_rate": scene_max_coverage_rate,
}
)
return datalist return datalist
def preprocess_cache(self): def preprocess_cache(self):
@ -107,9 +111,6 @@ class NBVReconstructionDataset(BaseDataset):
np.savetxt(cache_path, data) np.savetxt(cache_path, data)
except Exception as e: except Exception as e:
Log.error(f"Save cache failed: {e}") Log.error(f"Save cache failed: {e}")
# ----- Debug Trace ----- #
import ipdb; ipdb.set_trace()
# ------------------------ #
def __getitem__(self, index): def __getitem__(self, index):
data_item_info = self.datalist[index] data_item_info = self.datalist[index]
@ -117,18 +118,28 @@ class NBVReconstructionDataset(BaseDataset):
nbv = data_item_info["next_best_view"] nbv = data_item_info["next_best_view"]
max_coverage_rate = data_item_info["seq_max_coverage_rate"] max_coverage_rate = data_item_info["seq_max_coverage_rate"]
scene_name = data_item_info["scene_name"] scene_name = data_item_info["scene_name"]
scanned_views_pts, scanned_coverages_rate, scanned_n_to_world_pose = [], [], [] (
scanned_views_pts,
scanned_coverages_rate,
scanned_n_to_world_pose,
scanned_target_pts_num,
) = ([], [], [], [])
target_pts_num_dict = DataLoadUtil.load_target_pts_num_dict(
self.root_dir, scene_name
)
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]
view_path = DataLoadUtil.get_path(self.root_dir, scene_name, frame_idx) view_path = DataLoadUtil.get_path(self.root_dir, scene_name, frame_idx)
cam_info = DataLoadUtil.load_cam_info(view_path, binocular=True) cam_info = DataLoadUtil.load_cam_info(view_path, binocular=True)
target_pts_num = target_pts_num_dict[frame_idx]
n_to_world_pose = cam_info["cam_to_world"] n_to_world_pose = cam_info["cam_to_world"]
nR_to_world_pose = cam_info["cam_to_world_R"] nR_to_world_pose = cam_info["cam_to_world_R"]
if self.load_from_preprocess: if self.load_from_preprocess:
downsampled_target_point_cloud = DataLoadUtil.load_from_preprocessed_pts(view_path) downsampled_target_point_cloud = (
DataLoadUtil.load_from_preprocessed_pts(view_path)
)
else: else:
cached_data = None cached_data = None
if self.cache: if self.cache:
@ -136,72 +147,90 @@ class NBVReconstructionDataset(BaseDataset):
if cached_data is None: if cached_data is None:
print("load depth") print("load depth")
depth_L, depth_R = DataLoadUtil.load_depth(view_path, cam_info['near_plane'], cam_info['far_plane'], binocular=True) depth_L, depth_R = DataLoadUtil.load_depth(
point_cloud_L = DataLoadUtil.get_point_cloud(depth_L, cam_info['cam_intrinsic'], n_to_world_pose)['points_world'] view_path,
point_cloud_R = DataLoadUtil.get_point_cloud(depth_R, cam_info['cam_intrinsic'], nR_to_world_pose)['points_world'] 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_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_L = PtsUtil.random_downsample_point_cloud(
point_cloud_R = PtsUtil.random_downsample_point_cloud(point_cloud_R, 65536) point_cloud_L, 65536
overlap_points = DataLoadUtil.get_overlapping_points(point_cloud_L, point_cloud_R) )
downsampled_target_point_cloud = PtsUtil.random_downsample_point_cloud(overlap_points, self.pts_num) point_cloud_R = PtsUtil.random_downsample_point_cloud(
point_cloud_R, 65536
)
overlap_points = DataLoadUtil.get_overlapping_points(
point_cloud_L, point_cloud_R
)
downsampled_target_point_cloud = (
PtsUtil.random_downsample_point_cloud(
overlap_points, self.pts_num
)
)
if self.cache: if self.cache:
self.save_to_cache(scene_name, frame_idx, downsampled_target_point_cloud) self.save_to_cache(
scene_name, frame_idx, downsampled_target_point_cloud
)
else: else:
downsampled_target_point_cloud = cached_data downsampled_target_point_cloud = cached_data
scanned_views_pts.append(downsampled_target_point_cloud) scanned_views_pts.append(downsampled_target_point_cloud)
scanned_coverages_rate.append(coverage_rate) scanned_coverages_rate.append(coverage_rate)
n_to_world_6d = PoseUtil.matrix_to_rotation_6d_numpy(np.asarray(n_to_world_pose[:3,:3])) n_to_world_6d = PoseUtil.matrix_to_rotation_6d_numpy(
n_to_world_trans = n_to_world_pose[:3,3] 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) 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) scanned_n_to_world_pose.append(n_to_world_9d)
scanned_target_pts_num.append(target_pts_num)
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)
cam_info = DataLoadUtil.load_cam_info(nbv_path) cam_info = DataLoadUtil.load_cam_info(nbv_path)
best_frame_to_world = cam_info["cam_to_world"] best_frame_to_world = cam_info["cam_to_world"]
best_to_world_6d = PoseUtil.matrix_to_rotation_6d_numpy(np.asarray(best_frame_to_world[:3,:3])) best_to_world_6d = PoseUtil.matrix_to_rotation_6d_numpy(
best_to_world_trans = best_frame_to_world[:3,3] np.asarray(best_frame_to_world[:3, :3])
best_to_world_9d = np.concatenate([best_to_world_6d, best_to_world_trans], axis=0) )
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
)
combined_scanned_views_pts = np.concatenate(scanned_views_pts, axis=0) combined_scanned_views_pts = np.concatenate(scanned_views_pts, axis=0)
voxel_downsampled_combined_scanned_pts_np = PtsUtil.voxel_downsample_point_cloud(combined_scanned_views_pts, 0.002) voxel_downsampled_combined_scanned_pts_np = (
random_downsampled_combined_scanned_pts_np = PtsUtil.random_downsample_point_cloud(voxel_downsampled_combined_scanned_pts_np, self.pts_num) PtsUtil.voxel_downsample_point_cloud(combined_scanned_views_pts, 0.002)
)
random_downsampled_combined_scanned_pts_np = (
PtsUtil.random_downsample_point_cloud(
voxel_downsampled_combined_scanned_pts_np, self.pts_num
)
)
data_item = { data_item = {
"scanned_pts": np.asarray(scanned_views_pts,dtype=np.float32), "scanned_pts": np.asarray(scanned_views_pts, dtype=np.float32),
"combined_scanned_pts": np.asarray(random_downsampled_combined_scanned_pts_np,dtype=np.float32), "combined_scanned_pts": np.asarray(
random_downsampled_combined_scanned_pts_np, dtype=np.float32
),
"scanned_coverage_rate": scanned_coverages_rate, "scanned_coverage_rate": scanned_coverages_rate,
"scanned_n_to_world_pose_9d": np.asarray(scanned_n_to_world_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_coverage_rate": nbv_coverage_rate,
"best_to_world_pose_9d": np.asarray(best_to_world_9d,dtype=np.float32), "best_to_world_pose_9d": np.asarray(best_to_world_9d, dtype=np.float32),
"seq_max_coverage_rate": max_coverage_rate, "seq_max_coverage_rate": max_coverage_rate,
"scene_name": scene_name "scene_name": scene_name,
"scanned_target_points_num": np.asarray(
scanned_target_pts_num, dtype=np.int32
),
} }
# if self.type == namespace.Mode.TEST:
# diag = DataLoadUtil.get_bbox_diag(self.model_dir, scene_name)
# voxel_threshold = diag*0.02
# model_points_normals = DataLoadUtil.load_points_normals(self.root_dir, scene_name)
# pts_list = []
# for view in scanned_views:
# frame_idx = view[0]
# view_path = DataLoadUtil.get_path(self.root_dir, scene_name, frame_idx)
# point_cloud = DataLoadUtil.get_target_point_cloud_world_from_path(view_path, binocular=True)
# cam_params = DataLoadUtil.load_cam_info(view_path, binocular=True)
# sampled_point_cloud = ReconstructionUtil.filter_points(point_cloud, model_points_normals, cam_pose=cam_params["cam_to_world"], voxel_size=voxel_threshold, theta=self.filter_degree)
# pts_list.append(sampled_point_cloud)
# nL_to_world_pose = cam_params["cam_to_world"]
# nO_to_world_pose = cam_params["cam_to_world_O"]
# nO_to_nL_pose = np.dot(np.linalg.inv(nL_to_world_pose), nO_to_world_pose)
# data_item["scanned_target_pts_list"] = pts_list
# data_item["model_points_normals"] = model_points_normals
# data_item["voxel_threshold"] = voxel_threshold
# data_item["filter_degree"] = self.filter_degree
# data_item["scene_path"] = os.path.join(self.root_dir, scene_name)
# data_item["first_frame_to_world"] = np.asarray(first_frame_to_world, dtype=np.float32)
# data_item["nO_to_nL_pose"] = np.asarray(nO_to_nL_pose, dtype=np.float32)
return data_item return data_item
def __len__(self): def __len__(self):
@ -210,22 +239,44 @@ class NBVReconstructionDataset(BaseDataset):
def get_collate_fn(self): def get_collate_fn(self):
def collate_fn(batch): def collate_fn(batch):
collate_data = {} collate_data = {}
collate_data["scanned_pts"] = [torch.tensor(item['scanned_pts']) for item in batch] collate_data["scanned_pts"] = [
collate_data["scanned_n_to_world_pose_9d"] = [torch.tensor(item['scanned_n_to_world_pose_9d']) for item in batch] torch.tensor(item["scanned_pts"]) 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]) ]
collate_data["combined_scanned_pts"] = torch.stack([torch.tensor(item['combined_scanned_pts']) 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["scanned_target_points_num"] = [
torch.tensor(item["scanned_target_points_num"]) 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]
)
collate_data["combined_scanned_pts"] = torch.stack(
[torch.tensor(item["combined_scanned_pts"]) for item in batch]
)
if "first_frame_to_world" in batch[0]: 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]) 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(): for key in batch[0].keys():
if key not in ["scanned_pts", "scanned_n_to_world_pose_9d", "best_to_world_pose_9d", "first_frame_to_world", "combined_scanned_pts"]: if key not in [
"scanned_pts",
"scanned_n_to_world_pose_9d",
"best_to_world_pose_9d",
"first_frame_to_world",
"combined_scanned_pts",
"scanned_target_points_num",
]:
collate_data[key] = [item[key] for item in batch] collate_data[key] = [item[key] for item in batch]
return collate_data return collate_data
return collate_fn return collate_fn
# -------------- Debug ---------------- # # -------------- Debug ---------------- #
if __name__ == "__main__": if __name__ == "__main__":
import torch import torch
seed = 0 seed = 0
torch.manual_seed(seed) torch.manual_seed(seed)
np.random.seed(seed) np.random.seed(seed)
@ -244,41 +295,13 @@ if __name__ == "__main__":
} }
ds = NBVReconstructionDataset(config) ds = NBVReconstructionDataset(config)
print(len(ds)) print(len(ds))
#ds.__getitem__(10) # ds.__getitem__(10)
dl = ds.get_loader(shuffle=True) dl = ds.get_loader(shuffle=True)
for idx, data in enumerate(dl): for idx, data in enumerate(dl):
data = ds.process_batch(data, "cuda:0") data = ds.process_batch(data, "cuda:0")
print(data) print(data)
# ------ Debug Start ------ # ------ Debug Start ------
import ipdb;ipdb.set_trace() import ipdb
ipdb.set_trace()
# ------ Debug End ------ # ------ Debug End ------
#
# 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
# #np.savetxt("best_pts.txt", best_pts)
# for key, value in data.items():
# if isinstance(value, torch.Tensor):
# print(key, ":" ,value.shape)
# else:
# print(key, ":" ,len(value))
# if key == "scanned_n_to_world_pose_9d":
# for val in value:
# print(val.shape)
# if key == "scanned_pts":
# print("scanned_pts")
# for val in value:
# print(val.shape)
# cnt = 0
# for v in val:
# import ipdb;ipdb.set_trace()
# np.savetxt(f"pts_{cnt}.txt", v)
# cnt+=1
# print()

7
utils/data_load.py
View File

@ -98,6 +98,13 @@ class DataLoadUtil:
scene_info = json.load(f) scene_info = json.load(f)
return scene_info return scene_info
@staticmethod
def load_target_pts_num_dict(root, scene_name):
target_pts_num_path = os.path.join(root, scene_name, "target_pts_num.json")
with open(target_pts_num_path, "r") as f:
target_pts_num_dict = json.load(f)
return target_pts_num_dict
@staticmethod @staticmethod
def load_target_object_pose(root, scene_name): def load_target_object_pose(root, scene_name):
scene_info = DataLoadUtil.load_scene_info(root, scene_name) scene_info = DataLoadUtil.load_scene_info(root, scene_name)