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

update scan_points strategy

Browse Source
This commit is contained in:
hofee
2024-10-02 16:24:13 +08:00
parent 551282a0ec
commit c8b8a44252
6 changed files with 120 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files

82
runners/strategy_generator.py
View File

@@ -35,7 +35,7 @@ class StrategyGenerator(Runner):
def run(self):
dataset_name_list = ConfigManager.get("runner", "generate", "dataset_list")
voxel_threshold, overlap_threshold = ConfigManager.get("runner","generate","voxel_threshold"), ConfigManager.get("runner","generate","overlap_threshold")
voxel_threshold, soft_overlap_threshold, hard_overlap_threshold = ConfigManager.get("runner","generate","voxel_threshold"), ConfigManager.get("runner","generate","soft_overlap_threshold"), ConfigManager.get("runner","generate","hard_overlap_threshold")
for dataset_idx in range(len(dataset_name_list)):
dataset_name = dataset_name_list[dataset_idx]
status_manager.set_progress("generate_strategy", "strategy_generator", "dataset", dataset_idx, len(dataset_name_list))
@@ -47,23 +47,23 @@ class StrategyGenerator(Runner):
if to_idx == -1:
to_idx = len(scene_name_list)
cnt = 0
total = len(scene_name_list)
total = len(scene_name_list[from_idx:to_idx])
Log.info(f"Processing Dataset: {dataset_name}, From: {from_idx}, To: {to_idx}")
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 = diag*0.02
#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)
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")
cnt += 1
continue
try:
self.generate_sequence(root_dir, model_dir, scene_name,voxel_threshold, overlap_threshold)
except Exception as e:
Log.error(f"Scene <{scene_name}> Failed, Error: {e}")
self.generate_sequence(root_dir, model_dir, scene_name,voxel_threshold, soft_overlap_threshold, hard_overlap_threshold)
# except Exception as e:
# Log.error(f"Scene <{scene_name}> Failed, Error: {e}")
cnt += 1
status_manager.set_progress("generate_strategy", "strategy_generator", "scene", total, total)
status_manager.set_progress("generate_strategy", "strategy_generator", "dataset", len(dataset_name_list), len(dataset_name_list))
@@ -76,7 +76,7 @@ class StrategyGenerator(Runner):
def load_experiment(self, backup_name=None):
super().load_experiment(backup_name)
def generate_sequence(self, root, model_dir, scene_name, voxel_threshold, overlap_threshold):
def generate_sequence(self, root, model_dir, scene_name, voxel_threshold, soft_overlap_threshold, hard_overlap_threshold):
status_manager.set_status("generate_strategy", "strategy_generator", "scene", scene_name)
frame_num = DataLoadUtil.get_scene_seq_length(root, scene_name)
model_points_normals = DataLoadUtil.load_points_normals(root, scene_name)
@@ -84,47 +84,83 @@ class StrategyGenerator(Runner):
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"]
top = DataLoadUtil.get_display_table_top(root, scene_name)
scan_points = ReconstructionUtil.generate_scan_points(display_table_top=top,display_table_radius=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):
if self.load_pts and os.path.exists(os.path.join(root,scene_name, "pts", f"{frame_idx}.txt")):
sampled_point_cloud = np.loadtxt(os.path.join(root,scene_name, "pts", f"{frame_idx}.txt"))
indices = np.loadtxt(os.path.join(root,scene_name, "pts", f"{frame_idx}_indices.txt")).astype(np.int32).tolist()
status_manager.set_progress("generate_strategy", "strategy_generator", "loading frame", frame_idx, 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)
status_manager.set_progress("generate_strategy", "strategy_generator", "loading frame", frame_idx, frame_num)
point_cloud, display_table_pts = DataLoadUtil.get_target_point_cloud_world_from_path(path, binocular=True, get_display_table_pts=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)
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")
covered_pts_dir = os.path.join(pts_dir, "covered_scan_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(pts_dir, f"{frame_idx}_indices.txt"), indices)
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)
status_manager.set_progress("generate_strategy", "strategy_generator", "loading frame", frame_num, frame_num)
seq_num = min(self.seq_num, len(pts_list))
init_view_list = range(seq_num)
seq_num = min(self.seq_num, non_zero_cnt)
init_view_list = []
for i in range(seq_num):
if pts_list[i].shape[0] < 100:
continue
init_view_list.append(i)
seq_idx = 0
for init_view in init_view_list:
status_manager.set_progress("generate_strategy", "strategy_generator", "computing sequence", seq_idx, len(init_view_list))
limited_useful_view, _, _ = ReconstructionUtil.compute_next_best_view_sequence_with_overlap(down_sampled_model_pts, pts_list,init_view=init_view, threshold=voxel_threshold, overlap_threshold=overlap_threshold, status_info=self.status_info)
limited_useful_view, _, _ = ReconstructionUtil.compute_next_best_view_sequence_with_overlap(down_sampled_model_pts, pts_list, scan_points_indices_list = scan_points_indices_list,init_view=init_view,
threshold=voxel_threshold, soft_overlap_threshold=soft_overlap_threshold, hard_overlap_threshold= hard_overlap_threshold, scan_points_threshold=10, status_info=self.status_info)
data_pairs = self.generate_data_pairs(limited_useful_view)
seq_save_data = {
"data_pairs": data_pairs,