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

add preprocess

Browse Source
This commit is contained in:
hofee
2024-10-03 01:59:13 +08:00
parent f460e6e6b2
commit d7561738c6
8 changed files with 243 additions and 142 deletions
Download Patch File Download Diff File Expand all files Collapse all files

70
runners/strategy_generator.py
View File

@@ -52,9 +52,8 @@ class StrategyGenerator(Runner):
for scene_name in scene_name_list[from_idx:to_idx]:
Log.info(f"({dataset_name})Processing [{cnt}/{total}]: {scene_name}")
status_manager.set_progress("generate_strategy", "strategy_generator", "scene", cnt, total)
#diag = DataLoadUtil.get_bbox_diag(model_dir, scene_name)
voxel_threshold = 0.002
status_manager.set_status("generate_strategy", "strategy_generator", "voxel_threshold", voxel_threshold)
diag = DataLoadUtil.get_bbox_diag(model_dir, scene_name)
status_manager.set_status("generate_strategy", "strategy_generator", "diagonal", diag)
output_label_path = DataLoadUtil.get_label_path(root_dir, scene_name,0)
if os.path.exists(output_label_path) and not self.overwrite:
Log.info(f"Scene <{scene_name}> Already Exists, Skip")
@@ -82,71 +81,16 @@ class StrategyGenerator(Runner):
model_points_normals = DataLoadUtil.load_points_normals(root, scene_name)
model_pts = model_points_normals[:,:3]
down_sampled_model_pts = PtsUtil.voxel_downsample_point_cloud(model_pts, voxel_threshold)
display_table_info = DataLoadUtil.get_display_table_info(root, scene_name)
radius = display_table_info["radius"]
scan_points_path = os.path.join(root,scene_name, "scan_points.txt")
if os.path.exists(scan_points_path):
scan_points = np.loadtxt(scan_points_path)
else:
scan_points = ReconstructionUtil.generate_scan_points(display_table_top=0,display_table_radius=radius)
np.savetxt(scan_points_path, scan_points)
pts_list = []
scan_points_indices_list = []
non_zero_cnt = 0
for frame_idx in range(frame_num):
status_manager.set_progress("generate_strategy", "strategy_generator", "loading frame", frame_idx, frame_num)
pts_path = os.path.join(root,scene_name, "pts", f"{frame_idx}.txt")
if self.load_pts and pts_path:
with open(pts_path, 'r') as f:
pts_str = f.read()
if pts_str == "":
sampled_point_cloud = np.asarray([])
else:
sampled_point_cloud = np.loadtxt(pts_path)
indices_path = os.path.join(root,scene_name, "covered_scan_pts", f"{frame_idx}_indices.txt")
with open(indices_path, 'r') as f:
indices_str = f.read()
if indices_str == "":
indices = []
else:
indices = np.loadtxt(indices_path).astype(np.int32).tolist()
if isinstance(indices, int):
indices = [indices]
pts_list.append(sampled_point_cloud)
if sampled_point_cloud.shape[0] != 0:
non_zero_cnt += 1
scan_points_indices_list.append(indices)
else:
path = DataLoadUtil.get_path(root, scene_name, frame_idx)
cam_params = DataLoadUtil.load_cam_info(path, binocular=True)
point_cloud, display_table_pts = DataLoadUtil.get_target_point_cloud_world_from_path(path, binocular=True, get_display_table_pts=True)
if point_cloud.shape[0] != 0:
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)
non_zero_cnt += 1
else:
sampled_point_cloud = point_cloud
covered_pts, indices = ReconstructionUtil.compute_covered_scan_points(scan_points, display_table_pts)
if self.save_pts:
pts_dir = os.path.join(root,scene_name, "pts")
#display_dir = os.path.join(root,scene_name, "display_pts")
covered_pts_dir = os.path.join(root,scene_name, "covered_scan_pts")
if not os.path.exists(pts_dir):
os.makedirs(pts_dir)
if not os.path.exists(covered_pts_dir):
os.makedirs(covered_pts_dir)
# if not os.path.exists(display_dir):
# os.makedirs(display_dir)
np.savetxt(os.path.join(pts_dir, f"{frame_idx}.txt"), sampled_point_cloud)
#np.savetxt(os.path.join(display_dir, f"{frame_idx}.txt"), display_table_pts)
np.savetxt(os.path.join(covered_pts_dir, f"{frame_idx}.txt"), covered_pts)
np.savetxt(os.path.join(covered_pts_dir, f"{frame_idx}_indices.txt"), indices)
pts_list.append(sampled_point_cloud)
scan_points_indices_list.append(indices)
pts_path = os.path.join(root,scene_name, "target_pts", f"{frame_idx}.txt")
sampled_point_cloud = np.loadtxt(pts_path)
indices = None # ReconstructionUtil.compute_covered_scan_points(scan_points, display_table_pts)
pts_list.append(sampled_point_cloud)
scan_points_indices_list.append(indices)
status_manager.set_progress("generate_strategy", "strategy_generator", "loading frame", frame_num, frame_num)
seq_num = min(self.seq_num, non_zero_cnt)