diff --git a/app_generate_strategy.py b/app_generate_strategy.py index 9625297..28905e5 100644 --- a/app_generate_strategy.py +++ b/app_generate_strategy.py @@ -5,5 +5,5 @@ from runners.strategy_generator import StrategyGenerator class DataGenerateApp: @staticmethod def start(): - StrategyGenerator("configs/server/server_strategy_generate_config.yaml").run() + StrategyGenerator("configs/local/strategy_generate_config.yaml").run() \ No newline at end of file diff --git a/configs/local/strategy_generate_config.yaml b/configs/local/strategy_generate_config.yaml index 4a2a0e7..7195639 100644 --- a/configs/local/strategy_generate_config.yaml +++ b/configs/local/strategy_generate_config.yaml @@ -12,8 +12,9 @@ runner: generate: voxel_threshold: 0.003 - soft_overlap_threshold: 0.3 - hard_overlap_threshold: 0.6 + overlap_area_threshold: 25 + compute_with_normal: True + scan_points_threshold: 10 overwrite: False seq_num: 15 dataset_list: @@ -21,8 +22,8 @@ runner: datasets: OmniObject3d: - root_dir: /media/hofee/repository/full_data_output + root_dir: C:\\Document\\Local Project\\nbv_rec\\nbv_reconstruction\\temp from: 0 - to: -1 # -1 means end + to: 1 # -1 means end diff --git a/configs/local/view_generate_config.yaml b/configs/local/view_generate_config.yaml index 06e353c..b58d672 100644 --- a/configs/local/view_generate_config.yaml +++ b/configs/local/view_generate_config.yaml @@ -7,21 +7,12 @@ runner: name: debug root_dir: experiments generate: -<<<<<<< HEAD port: 5002 from: 600 to: -1 # -1 means all object_dir: /media/hofee/data/data/object_meshes_part1 table_model_path: "/media/hofee/data/data/others/table.obj" output_dir: /media/hofee/repository/data_part_1 -======= - port: 5000 - from: 0 - to: -1 # -1 means all - object_dir: H:\\AI\\Datasets\\object_meshes_part2 - table_model_path: "H:\\AI\\Datasets\\table.obj" - output_dir: C:\\Document\\Datasets\\nbv_rec_part2 ->>>>>>> c55a398b6d5c347497b528bdd460e26ffdd184e8 binocular_vision: true plane_size: 10 max_views: 512 diff --git a/preprocess/preprocessor.py b/preprocess/preprocessor.py index aa46a02..2f01c0a 100644 --- a/preprocess/preprocessor.py +++ b/preprocess/preprocessor.py @@ -163,16 +163,10 @@ def save_scene_data(root, scene, scene_idx=0, scene_total=1,file_type="txt"): if __name__ == "__main__": #root = "/media/hofee/repository/new_data_with_normal" - root = r"/media/hofee/repository/data_part_1" - # list_path = r"/media/hofee/repository/full_list.txt" - # scene_list = [] - - # with open(list_path, "r") as f: - # for line in f: - # scene_list.append(line.strip()) + root = r"C:\Document\Local Project\nbv_rec\nbv_reconstruction\temp" scene_list = os.listdir(root) from_idx = 0 # 1000 - to_idx = 600 # 1500 + to_idx = len(scene_list) # 1500 cnt = 0 diff --git a/runners/strategy_generator.py b/runners/strategy_generator.py index cd17860..60a5cfc 100644 --- a/runners/strategy_generator.py +++ b/runners/strategy_generator.py @@ -24,12 +24,15 @@ class StrategyGenerator(Runner): } self.overwrite = ConfigManager.get("runner", "generate", "overwrite") self.seq_num = ConfigManager.get("runner","generate","seq_num") + self.overlap_area_threshold = ConfigManager.get("runner","generate","overlap_area_threshold") + self.compute_with_normal = ConfigManager.get("runner","generate","compute_with_normal") + self.scan_points_threshold = ConfigManager.get("runner","generate","scan_points_threshold") def run(self): dataset_name_list = ConfigManager.get("runner", "generate", "dataset_list") - 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") + voxel_threshold = ConfigManager.get("runner","generate","voxel_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)) @@ -51,7 +54,7 @@ class StrategyGenerator(Runner): cnt += 1 continue - self.generate_sequence(root_dir, scene_name,voxel_threshold, soft_overlap_threshold, hard_overlap_threshold) + self.generate_sequence(root_dir, scene_name,voxel_threshold) 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)) @@ -64,28 +67,34 @@ class StrategyGenerator(Runner): def load_experiment(self, backup_name=None): super().load_experiment(backup_name) - def generate_sequence(self, root, scene_name, voxel_threshold, soft_overlap_threshold, hard_overlap_threshold): + def generate_sequence(self, root, scene_name, voxel_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) model_pts = model_points_normals[:,:3] - down_sampled_model_pts = PtsUtil.voxel_downsample_point_cloud(model_pts, voxel_threshold) + down_sampled_model_pts, idx = PtsUtil.voxel_downsample_point_cloud(model_pts, voxel_threshold, require_idx=True) + down_sampled_model_nrm = model_points_normals[idx, 3:] pts_list = [] + nrm_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}.npy") + nrm_path = os.path.join(root,scene_name, "nrm", f"{frame_idx}.npy") idx_path = os.path.join(root,scene_name, "scan_points_indices", f"{frame_idx}.npy") - point_cloud = np.load(pts_path) - sampled_point_cloud = PtsUtil.voxel_downsample_point_cloud(point_cloud, voxel_threshold) + pts = np.load(pts_path) + if pts.shape[0] == 0: + nrm = np.zeros((0,3)) + else: + nrm = np.load(nrm_path) indices = np.load(idx_path) - pts_list.append(sampled_point_cloud) - + pts_list.append(pts) + nrm_list.append(nrm) scan_points_indices_list.append(indices) - if sampled_point_cloud.shape[0] > 0: + if pts.shape[0] > 0: non_zero_cnt += 1 status_manager.set_progress("generate_strategy", "strategy_generator", "loading frame", frame_num, frame_num) @@ -93,7 +102,7 @@ class StrategyGenerator(Runner): init_view_list = [] idx = 0 while len(init_view_list) < seq_num and idx < len(pts_list): - if pts_list[idx].shape[0] > 100: + if pts_list[idx].shape[0] > 50: init_view_list.append(idx) idx += 1 @@ -102,8 +111,13 @@ class StrategyGenerator(Runner): for init_view in init_view_list: status_manager.set_progress("generate_strategy", "strategy_generator", "computing sequence", seq_idx, len(init_view_list)) start = time.time() - 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) + + if not self.compute_with_normal: + limited_useful_view, _, _ = ReconstructionUtil.compute_next_best_view_sequence(down_sampled_model_pts, pts_list, scan_points_indices_list = scan_points_indices_list,init_view=init_view, + threshold=voxel_threshold, scan_points_threshold=self.scan_points_threshold, overlap_area_threshold=self.overlap_area_threshold, status_info=self.status_info) + else: + limited_useful_view, _, _ = ReconstructionUtil.compute_next_best_view_sequence_with_normal(down_sampled_model_pts, down_sampled_model_nrm, pts_list, nrm_list, scan_points_indices_list = scan_points_indices_list,init_view=init_view, + threshold=voxel_threshold, scan_points_threshold=self.scan_points_threshold, overlap_area_threshold=self.overlap_area_threshold, status_info=self.status_info) end = time.time() print(f"Time: {end-start}") data_pairs = self.generate_data_pairs(limited_useful_view) diff --git a/utils/pts.py b/utils/pts.py index e5cb436..58336c8 100644 --- a/utils/pts.py +++ b/utils/pts.py @@ -5,7 +5,7 @@ import torch class PtsUtil: @staticmethod - def voxel_downsample_point_cloud(point_cloud, voxel_size=0.005): + def voxel_downsample_point_cloud(point_cloud, voxel_size=0.005, require_idx=False): voxel_indices = np.floor(point_cloud / voxel_size).astype(np.int32) unique_voxels = np.unique(voxel_indices, axis=0, return_inverse=True) return unique_voxels[0]*voxel_size diff --git a/utils/reconstruction.py b/utils/reconstruction.py index fed95e1..6266a84 100644 --- a/utils/reconstruction.py +++ b/utils/reconstruction.py @@ -8,14 +8,15 @@ class ReconstructionUtil: def compute_coverage_rate(target_point_cloud, combined_point_cloud, threshold=0.01): kdtree = cKDTree(combined_point_cloud) distances, _ = kdtree.query(target_point_cloud) - covered_points_num = np.sum(distances < threshold) + covered_points_num = np.sum(distances < threshold*2) coverage_rate = covered_points_num / target_point_cloud.shape[0] return coverage_rate, covered_points_num + @staticmethod def compute_coverage_rate_with_normal(target_point_cloud, combined_point_cloud, target_normal, combined_normal, threshold=0.01, normal_threshold=0.1): kdtree = cKDTree(combined_point_cloud) distances, indices = kdtree.query(target_point_cloud) - is_covered_by_distance = distances < threshold + is_covered_by_distance = distances < threshold*2 normal_dots = np.einsum('ij,ij->i', target_normal, combined_normal[indices]) is_covered_by_normal = normal_dots > normal_threshold covered_points_num = np.sum(is_covered_by_distance & is_covered_by_normal) @@ -25,15 +26,14 @@ class ReconstructionUtil: @staticmethod - def compute_overlap_rate(new_point_cloud, combined_point_cloud, threshold=0.01): + def check_overlap(new_point_cloud, combined_point_cloud, overlap_area_threshold=25, voxel_size=0.01): kdtree = cKDTree(combined_point_cloud) distances, _ = kdtree.query(new_point_cloud) - overlapping_points = np.sum(distances < threshold) - if new_point_cloud.shape[0] == 0: - overlap_rate = 0 - else: - overlap_rate = overlapping_points / new_point_cloud.shape[0] - return overlap_rate + overlapping_points = np.sum(distances < voxel_size*2) + cm = 0.01 + voxel_size_cm = voxel_size / cm + overlap_area = overlapping_points * voxel_size_cm * voxel_size_cm + return overlap_area > overlap_area_threshold @staticmethod @@ -49,7 +49,7 @@ class ReconstructionUtil: return new_added_points @staticmethod - def compute_next_best_view_sequence_with_overlap(target_point_cloud, point_cloud_list, scan_points_indices_list, threshold=0.01, soft_overlap_threshold=0.5, hard_overlap_threshold=0.7, init_view = 0, scan_points_threshold=5, status_info=None): + def compute_next_best_view_sequence(target_point_cloud, point_cloud_list, scan_points_indices_list, threshold=0.01, overlap_area_threshold=25, init_view = 0, scan_points_threshold=5, status_info=None): selected_views = [init_view] combined_point_cloud = point_cloud_list[init_view] history_indices = [scan_points_indices_list[init_view]] @@ -83,22 +83,16 @@ class ReconstructionUtil: if selected_views: new_scan_points_indices = scan_points_indices_list[view_index] if not ReconstructionUtil.check_scan_points_overlap(history_indices, new_scan_points_indices, scan_points_threshold): - overlap_threshold = hard_overlap_threshold + curr_overlap_area_threshold = overlap_area_threshold else: - overlap_threshold = soft_overlap_threshold - start = time.time() - overlap_rate = ReconstructionUtil.compute_overlap_rate(point_cloud_list[view_index],combined_point_cloud, threshold) - end = time.time() - # print(f"overlap_rate Time: {end-start}") - if overlap_rate < overlap_threshold: + curr_overlap_area_threshold = overlap_area_threshold * 0.5 + + if not ReconstructionUtil.check_overlap(point_cloud_list[view_index], combined_point_cloud, overlap_area_threshold = curr_overlap_area_threshold, voxel_size=threshold): continue - start = time.time() new_combined_point_cloud = np.vstack([combined_point_cloud, point_cloud_list[view_index]]) new_downsampled_combined_point_cloud = PtsUtil.voxel_downsample_point_cloud(new_combined_point_cloud,threshold) new_coverage, new_covered_num = ReconstructionUtil.compute_coverage_rate(downsampled_max_rec_pts, new_downsampled_combined_point_cloud, threshold) - end = time.time() - #print(f"compute_coverage_rate Time: {end-start}") coverage_increase = new_coverage - current_coverage if coverage_increase > best_coverage_increase: best_coverage_increase = coverage_increase @@ -107,6 +101,101 @@ class ReconstructionUtil: best_combined_point_cloud = new_downsampled_combined_point_cloud + if best_view is not None: + if best_coverage_increase <=1e-3 or best_covered_num - current_covered_num <= 5: + break + + selected_views.append(best_view) + best_rec_pts_num = best_combined_point_cloud.shape[0] + print(f"Current rec pts num: {curr_rec_pts_num}, Best rec pts num: {best_rec_pts_num}, Best cover pts: {best_covered_num}, Max rec pts num: {max_rec_pts_num}") + print(f"Current coverage: {current_coverage+best_coverage_increase}, Best coverage increase: {best_coverage_increase}, Max Real coverage: {max_real_rec_pts_coverage}") + current_covered_num = best_covered_num + curr_rec_pts_num = best_rec_pts_num + combined_point_cloud = best_combined_point_cloud + remaining_views.remove(best_view) + history_indices.append(scan_points_indices_list[best_view]) + current_coverage += best_coverage_increase + cnt_processed_view += 1 + if status_info is not None: + sm = status_info["status_manager"] + app_name = status_info["app_name"] + runner_name = status_info["runner_name"] + sm.set_status(app_name, runner_name, "current coverage", current_coverage) + sm.set_progress(app_name, runner_name, "processed view", cnt_processed_view, len(point_cloud_list)) + + view_sequence.append((best_view, current_coverage)) + + else: + break + if status_info is not None: + sm = status_info["status_manager"] + app_name = status_info["app_name"] + runner_name = status_info["runner_name"] + sm.set_progress(app_name, runner_name, "processed view", len(point_cloud_list), len(point_cloud_list)) + return view_sequence, remaining_views, combined_point_cloud + + @staticmethod + def compute_next_best_view_sequence_with_normal(target_point_cloud, target_normal, point_cloud_list, normal_list, scan_points_indices_list, threshold=0.01, overlap_area_threshold=25, init_view = 0, scan_points_threshold=5, status_info=None): + selected_views = [init_view] + combined_point_cloud = point_cloud_list[init_view] + combined_normal = normal_list[init_view] + history_indices = [scan_points_indices_list[init_view]] + + max_rec_pts = np.vstack(point_cloud_list) + max_rec_nrm = np.vstack(normal_list) + downsampled_max_rec_pts, idx = PtsUtil.voxel_downsample_point_cloud(max_rec_pts, threshold, require_idx=True) + downsampled_max_rec_nrm = max_rec_nrm[idx] + + max_rec_pts_num = downsampled_max_rec_pts.shape[0] + try: + max_real_rec_pts_coverage, _ = ReconstructionUtil.compute_coverage_rate_with_normal(target_point_cloud, downsampled_max_rec_pts, target_normal, downsampled_max_rec_nrm, threshold) + except: + import ipdb; ipdb.set_trace() + + new_coverage, new_covered_num = ReconstructionUtil.compute_coverage_rate_with_normal(downsampled_max_rec_pts, combined_point_cloud, downsampled_max_rec_nrm, combined_normal, threshold) + current_coverage = new_coverage + current_covered_num = new_covered_num + + remaining_views = list(range(len(point_cloud_list))) + view_sequence = [(init_view, current_coverage)] + cnt_processed_view = 0 + remaining_views.remove(init_view) + curr_rec_pts_num = combined_point_cloud.shape[0] + + while remaining_views: + best_view = None + best_coverage_increase = -1 + best_combined_point_cloud = None + best_combined_normal = None + best_covered_num = 0 + + for view_index in remaining_views: + if point_cloud_list[view_index].shape[0] == 0: + continue + if selected_views: + new_scan_points_indices = scan_points_indices_list[view_index] + if not ReconstructionUtil.check_scan_points_overlap(history_indices, new_scan_points_indices, scan_points_threshold): + curr_overlap_area_threshold = overlap_area_threshold + else: + curr_overlap_area_threshold = overlap_area_threshold * 0.5 + + if not ReconstructionUtil.check_overlap(point_cloud_list[view_index], combined_point_cloud, overlap_area_threshold = curr_overlap_area_threshold, voxel_size=threshold): + continue + + new_combined_point_cloud = np.vstack([combined_point_cloud, point_cloud_list[view_index]]) + new_combined_normal = np.vstack([combined_normal, normal_list[view_index]]) + new_downsampled_combined_point_cloud = PtsUtil.voxel_downsample_point_cloud(new_combined_point_cloud,threshold) + new_downsampled_combined_normal = new_combined_normal[idx] + new_coverage, new_covered_num = ReconstructionUtil.compute_coverage_rate_with_normal(downsampled_max_rec_pts, new_downsampled_combined_point_cloud, downsampled_max_rec_nrm, new_downsampled_combined_normal, threshold) + coverage_increase = new_coverage - current_coverage + if coverage_increase > best_coverage_increase: + best_coverage_increase = coverage_increase + best_view = view_index + best_covered_num = new_covered_num + best_combined_point_cloud = new_downsampled_combined_point_cloud + best_combined_normal = new_downsampled_combined_normal + + if best_view is not None: if best_coverage_increase <=1e-3 or best_covered_num - current_covered_num <= 5: break @@ -118,6 +207,7 @@ class ReconstructionUtil: current_covered_num = best_covered_num curr_rec_pts_num = best_rec_pts_num combined_point_cloud = best_combined_point_cloud + combined_normal = best_combined_normal remaining_views.remove(best_view) history_indices.append(scan_points_indices_list[best_view]) current_coverage += best_coverage_increase diff --git a/utils/vis.py b/utils/vis.py index a6ef750..61fe554 100644 --- a/utils/vis.py +++ b/utils/vis.py @@ -47,6 +47,42 @@ class visualizeUtil: all_combined_pts = np.vstack(all_combined_pts) downsampled_all_pts = PtsUtil.voxel_downsample_point_cloud(all_combined_pts, 0.001) np.savetxt(os.path.join(output_dir, "all_combined_pts.txt"), downsampled_all_pts) + + @staticmethod + def save_seq_cam_pos_and_cam_axis(root, scene, frame_idx_list, output_dir): + all_cam_pos = [] + all_cam_axis = [] + for i in frame_idx_list: + path = DataLoadUtil.get_path(root, scene, i) + cam_info = DataLoadUtil.load_cam_info(path, binocular=True) + cam_pose = cam_info["cam_to_world"] + cam_pos = cam_pose[:3, 3] + cam_axis = cam_pose[:3, 2] + + num_samples = 10 + sample_points = [cam_pos + 0.02*t * cam_axis for t in range(num_samples)] + sample_points = np.array(sample_points) + + all_cam_pos.append(cam_pos) + all_cam_axis.append(sample_points) + + all_cam_pos = np.array(all_cam_pos) + all_cam_axis = np.array(all_cam_axis).reshape(-1, 3) + np.savetxt(os.path.join(output_dir, "seq_cam_pos.txt"), all_cam_pos) + np.savetxt(os.path.join(output_dir, "seq_cam_axis.txt"), all_cam_axis) + + @staticmethod + def save_seq_combined_pts(root, scene, frame_idx_list, output_dir): + all_combined_pts = [] + for i in frame_idx_list: + path = DataLoadUtil.get_path(root, scene, i) + pts = DataLoadUtil.load_from_preprocessed_pts(path,"npy") + if pts.shape[0] == 0: + continue + all_combined_pts.append(pts) + all_combined_pts = np.vstack(all_combined_pts) + downsampled_all_pts = PtsUtil.voxel_downsample_point_cloud(all_combined_pts, 0.001) + np.savetxt(os.path.join(output_dir, "seq_combined_pts.txt"), downsampled_all_pts) @staticmethod def save_target_mesh_at_world_space( @@ -126,12 +162,14 @@ class visualizeUtil: # ------ Debug ------ if __name__ == "__main__": - root = r"/home/yan20/nbv_rec/project/franka_control/temp" + root = r"C:\Document\Local Project\nbv_rec\nbv_reconstruction\temp" model_dir = r"H:\\AI\\Datasets\\scaled_object_box_meshes" scene = "box" output_dir = r"C:\Document\Local Project\nbv_rec\nbv_reconstruction\test" #visualizeUtil.save_all_cam_pos_and_cam_axis(root, scene, output_dir) visualizeUtil.save_all_combined_pts(root, scene, output_dir) + visualizeUtil.save_seq_combined_pts(root, scene, [0, 121, 286, 175, 111,366,45,230,232,225,255,17,199,78,60], output_dir) + visualizeUtil.save_seq_cam_pos_and_cam_axis(root, scene, [0, 121, 286, 175, 111,366,45,230,232,225,255,17,199,78,60], output_dir) visualizeUtil.save_target_mesh_at_world_space(root, model_dir, scene) - #visualizeUtil.save_points_and_normals(root, scene,"10", output_dir, binocular=True) \ No newline at end of file + #visualizeUtil.save_points_and_normals(root, scene,"10", output_dir, binocular=True)