diff --git a/configs/local/inference_config.yaml b/configs/local/inference_config.yaml index 13a7acc..3029fb3 100644 --- a/configs/local/inference_config.yaml +++ b/configs/local/inference_config.yaml @@ -14,8 +14,8 @@ runner: dataset_list: - OmniObject3d_test - blender_script_path: "C:\\Document\\Local Project\\nbv_rec\\blender\\data_renderer.py" - output_dir: "C:\\Document\\Datasets\\inference_scan_pts_overlap_global_full_on_testset" + blender_script_path: "/media/hofee/data/project/python/nbv_reconstruction/blender/data_renderer.py" + output_dir: "/media/hofee/data/data/new_inference_test_output" pipeline: nbv_reconstruction_pipeline voxel_size: 0.003 @@ -34,10 +34,10 @@ dataset: # load_from_preprocess: True OmniObject3d_test: - root_dir: "C:\\Document\\Datasets\\inference_test" - model_dir: "C:\\Document\\Datasets\\scaled_object_meshes" + root_dir: "/media/hofee/data/data/new_testset_output" + model_dir: "/media/hofee/data/data/scaled_object_meshes" source: seq_reconstruction_dataset_preprocessed - split_file: "C:\\Document\\Datasets\\data_list\\OmniObject3d_test.txt" + # split_file: "C:\\Document\\Datasets\\data_list\\OmniObject3d_test.txt" type: test filter_degree: 75 eval_list: diff --git a/configs/local/view_generate_config.yaml b/configs/local/view_generate_config.yaml index b58d672..3bb9c50 100644 --- a/configs/local/view_generate_config.yaml +++ b/configs/local/view_generate_config.yaml @@ -7,17 +7,19 @@ runner: name: debug root_dir: experiments generate: - port: 5002 - from: 600 + port: 5000 + from: 0 to: -1 # -1 means all - object_dir: /media/hofee/data/data/object_meshes_part1 + object_dir: /media/hofee/data/data/scaled_object_meshes table_model_path: "/media/hofee/data/data/others/table.obj" - output_dir: /media/hofee/repository/data_part_1 + output_dir: /media/hofee/data/data/new_testset + object_list_path: /media/hofee/data/data/OmniObject3d_test.txt + use_list: True binocular_vision: true plane_size: 10 max_views: 512 min_views: 128 - random_view_ratio: 0.02 + random_view_ratio: 0.01 min_cam_table_included_degree: 20 max_diag: 0.7 min_diag: 0.01 diff --git a/core/seq_dataset.py b/core/seq_dataset.py index 68774d5..6e71a9b 100644 --- a/core/seq_dataset.py +++ b/core/seq_dataset.py @@ -8,7 +8,7 @@ import torch import os import sys -sys.path.append(r"C:\Document\Local Project\nbv_rec\nbv_reconstruction") +sys.path.append(r"/media/hofee/data/project/python/nbv_reconstruction/nbv_reconstruction") from utils.data_load import DataLoadUtil from utils.pose import PoseUtil @@ -47,7 +47,8 @@ class SeqReconstructionDataset(BaseDataset): with open(self.split_file_path, "r") as f: for line in f: scene_name = line.strip() - scene_name_list.append(scene_name) + if os.path.exists(os.path.join(self.root_dir, scene_name)): + scene_name_list.append(scene_name) return scene_name_list def get_scene_name_list(self): @@ -58,29 +59,19 @@ class SeqReconstructionDataset(BaseDataset): total = len(self.scene_name_list) for idx, scene_name in enumerate(self.scene_name_list): print(f"processing {scene_name} ({idx}/{total})") - seq_num = DataLoadUtil.get_label_num(self.root_dir, scene_name) - scene_max_coverage_rate = 0 - max_coverage_rate_list = [] scene_max_cr_idx = 0 - for seq_idx in range(seq_num): - label_path = DataLoadUtil.get_label_path( - self.root_dir, scene_name, seq_idx - ) - label_data = DataLoadUtil.load_label(label_path) - max_coverage_rate = label_data["max_coverage_rate"] - if max_coverage_rate > scene_max_coverage_rate: - scene_max_coverage_rate = max_coverage_rate - scene_max_cr_idx = seq_idx - max_coverage_rate_list.append(max_coverage_rate) - best_label_path = DataLoadUtil.get_label_path(self.root_dir, scene_name, scene_max_cr_idx) - best_label_data = DataLoadUtil.load_label(best_label_path) - first_frame = best_label_data["best_sequence"][0] - best_seq_len = len(best_label_data["best_sequence"]) + frame_len = DataLoadUtil.get_scene_seq_length(self.root_dir, scene_name) + + for i in range(frame_len): + path = DataLoadUtil.get_path(self.root_dir, scene_name, i) + pts = DataLoadUtil.load_from_preprocessed_pts(path, "npy") + if pts.shape[0] == 0: + continue datalist.append({ "scene_name": scene_name, - "first_frame": first_frame, - "best_seq_len": best_seq_len, - "max_coverage_rate": scene_max_coverage_rate, + "first_frame": i, + "best_seq_len": -1, + "max_coverage_rate": 1.0, "label_idx": scene_max_cr_idx, }) return datalist @@ -131,8 +122,7 @@ class SeqReconstructionDataset(BaseDataset): scanned_n_to_world_pose, ) = ([], [], []) view = data_item_info["first_frame"] - frame_idx = view[0] - coverage_rate = view[1] + frame_idx = view view_path = DataLoadUtil.get_path(self.root_dir, scene_name, frame_idx) cam_info = DataLoadUtil.load_cam_info(view_path, binocular=True) @@ -144,7 +134,7 @@ class SeqReconstructionDataset(BaseDataset): target_point_cloud, self.pts_num ) scanned_views_pts.append(downsampled_target_point_cloud) - scanned_coverages_rate.append(coverage_rate) + n_to_world_6d = PoseUtil.matrix_to_rotation_6d_numpy( np.asarray(n_to_world_pose[:3, :3]) ) @@ -161,7 +151,6 @@ class SeqReconstructionDataset(BaseDataset): gt_pts = self.seq_combined_pts(scene_name, frame_list) data_item = { "first_scanned_pts": np.asarray(scanned_views_pts, dtype=np.float32), # Ndarray(S x Nv x 3) - "first_scanned_coverage_rate": scanned_coverages_rate, # List(S): Float, range(0, 1) "first_scanned_n_to_world_pose_9d": np.asarray(scanned_n_to_world_pose, dtype=np.float32), # Ndarray(S x 9) "seq_max_coverage_rate": max_coverage_rate, # Float, range(0, 1) "best_seq_len": best_seq_len, # Int @@ -179,40 +168,37 @@ class SeqReconstructionDataset(BaseDataset): # -------------- Debug ---------------- # if __name__ == "__main__": + #import ipdb; ipdb.set_trace() import torch + from tqdm import tqdm + import pickle + import os seed = 0 torch.manual_seed(seed) np.random.seed(seed) - ''' - OmniObject3d_test: - root_dir: "H:\\AI\\Datasets\\packed_test_data" - model_dir: "H:\\AI\\Datasets\\scaled_object_meshes" - source: seq_reconstruction_dataset - split_file: "H:\\AI\\Datasets\\data_list\\OmniObject3d_test.txt" - type: test - filter_degree: 75 - eval_list: - - pose_diff - - coverage_rate_increase - ratio: 0.1 - batch_size: 1 - num_workers: 12 - pts_num: 8192 - load_from_preprocess: True - ''' + config = { - "root_dir": "H:\\AI\\Datasets\\packed_test_data", + "root_dir": "/media/hofee/data/data/new_testset", "source": "seq_reconstruction_dataset", - "split_file": "H:\\AI\\Datasets\\data_list\\OmniObject3d_test.txt", + "split_file": "/media/hofee/data/data/OmniObject3d_test.txt", "load_from_preprocess": True, - "ratio": 1, "filter_degree": 75, "num_workers": 0, "pts_num": 8192, - "type": "test", + "type": namespace.Mode.TEST, } - ds = SeqReconstructionDataset(config) - print(len(ds)) - print(ds.__getitem__(10)) + output_dir = "/media/hofee/data/data/new_testset_output" + os.makedirs(output_dir, exist_ok=True) + + ds = SeqReconstructionDataset(config) + for i in tqdm(range(len(ds)), desc="processing dataset"): + output_path = os.path.join(output_dir, f"item_{i}.pkl") + item = ds.__getitem__(i) + for key, value in item.items(): + if isinstance(value, np.ndarray): + item[key] = value.tolist() + import ipdb; ipdb.set_trace() + with open(output_path, "wb") as f: + pickle.dump(item, f) \ No newline at end of file diff --git a/core/seq_dataset_preprocessed.py b/core/seq_dataset_preprocessed.py index 9f84bf4..a113c9f 100644 --- a/core/seq_dataset_preprocessed.py +++ b/core/seq_dataset_preprocessed.py @@ -22,14 +22,13 @@ class SeqReconstructionDatasetPreprocessed(BaseDataset): super(SeqReconstructionDatasetPreprocessed, self).__init__(config) self.config = config self.root_dir = config["root_dir"] - self.real_root_dir = r"H:\AI\Datasets\packed_test_data" + self.real_root_dir = r"/media/hofee/data/data/new_testset" self.item_list = os.listdir(self.root_dir) def __getitem__(self, index): data = pickle.load(open(os.path.join(self.root_dir, self.item_list[index]), "rb")) data_item = { "first_scanned_pts": np.asarray(data["first_scanned_pts"], dtype=np.float32), # Ndarray(S x Nv x 3) - "first_scanned_coverage_rate": data["first_scanned_coverage_rate"], # List(S): Float, range(0, 1) "first_scanned_n_to_world_pose_9d": np.asarray(data["first_scanned_n_to_world_pose_9d"], dtype=np.float32), # Ndarray(S x 9) "seq_max_coverage_rate": data["seq_max_coverage_rate"], # Float, range(0, 1) "best_seq_len": data["best_seq_len"], # Int diff --git a/preprocess/pack_upload_data.py b/preprocess/pack_upload_data.py index 2e5de1d..a4ec228 100644 --- a/preprocess/pack_upload_data.py +++ b/preprocess/pack_upload_data.py @@ -29,8 +29,8 @@ def pack_all_scenes(root, scene_list, output_dir): pack_scene_data(root, scene, output_dir) if __name__ == "__main__": - root = r"H:\AI\Datasets\nbv_rec_part2" - output_dir = r"H:\AI\Datasets\upload_part2" + root = r"/media/hofee/repository/data_part_1" + output_dir = r"/media/hofee/repository/upload_part1" scene_list = os.listdir(root) from_idx = 0 to_idx = len(scene_list) diff --git a/preprocess/preprocessor.py b/preprocess/preprocessor.py index b1ff29a..18c2299 100644 --- a/preprocess/preprocessor.py +++ b/preprocess/preprocessor.py @@ -164,10 +164,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"H:\AI\Datasets\nbv_rec_part2" + root = "/media/hofee/data/data/new_testset" scene_list = os.listdir(root) from_idx = 0 # 1000 - to_idx = 600 # 1500 + to_idx = len(scene_list) # 1500 cnt = 0 @@ -179,7 +179,11 @@ if __name__ == "__main__": print(f"Scene {scene} has been processed") cnt+=1 continue - save_scene_data(root, scene, cnt, total, file_type="npy") + try: + save_scene_data(root, scene, cnt, total, file_type="npy") + except Exception as e: + print(f"Error occurred when processing scene {scene}") + print(e) cnt+=1 end = time.time() print(f"Time cost: {end-start}") diff --git a/runners/inferencer.py b/runners/inferencer.py index a38b9e8..ef66858 100644 --- a/runners/inferencer.py +++ b/runners/inferencer.py @@ -76,6 +76,8 @@ class Inferencer(Runner): for i in tqdm(range(total), desc=f"Processing {test_set_name}", ncols=100): data = test_set.__getitem__(i) scene_name = data["scene_name"] + if scene_name != "omniobject3d-book_004": + continue inference_result_path = os.path.join(self.output_dir, test_set_name, f"{scene_name}.pkl") if os.path.exists(inference_result_path): Log.info(f"Inference result already exists for scene: {scene_name}") @@ -87,7 +89,7 @@ class Inferencer(Runner): status_manager.set_progress("inference", "inferencer", f"dataset", len(self.test_set_list), len(self.test_set_list)) - def predict_sequence(self, data, cr_increase_threshold=0, overlap_area_threshold=25, scan_points_threshold=10, max_iter=50, max_retry = 7): + def predict_sequence(self, data, cr_increase_threshold=0, overlap_area_threshold=25, scan_points_threshold=10, max_iter=50, max_retry = 5): scene_name = data["scene_name"] Log.info(f"Processing scene: {scene_name}") status_manager.set_status("inference", "inferencer", "scene", scene_name) @@ -110,10 +112,13 @@ class Inferencer(Runner): input_data["scanned_n_to_world_pose_9d"] = [torch.tensor(data["first_scanned_n_to_world_pose_9d"], dtype=torch.float32).to(self.device)] input_data["mode"] = namespace.Mode.TEST input_pts_N = input_data["combined_scanned_pts"].shape[1] + root = os.path.dirname(scene_path) + display_table_info = DataLoadUtil.get_display_table_info(root, scene_name) radius = display_table_info["radius"] scan_points = np.asarray(ReconstructionUtil.generate_scan_points(display_table_top=0,display_table_radius=radius)) + first_frame_target_pts, first_frame_target_normals, first_frame_scan_points_indices = RenderUtil.render_pts(first_frame_to_world, scene_path, self.script_path, scan_points, voxel_threshold=voxel_threshold, filter_degree=filter_degree, nO_to_nL_pose=O_to_L_pose) scanned_view_pts = [first_frame_target_pts] history_indices = [first_frame_scan_points_indices] @@ -124,6 +129,7 @@ class Inferencer(Runner): retry = 0 pred_cr_seq = [last_pred_cr] success = 0 + last_pts_num = PtsUtil.voxel_downsample_point_cloud(data["first_scanned_pts"][0], 0.002).shape[0] import time while len(pred_cr_seq) < max_iter and retry < max_retry: start_time = time.time() @@ -146,7 +152,7 @@ class Inferencer(Runner): curr_overlap_area_threshold = overlap_area_threshold * 0.5 downsampled_new_target_pts = PtsUtil.voxel_downsample_point_cloud(new_target_pts, voxel_threshold) - overlap, new_added_pts_num = ReconstructionUtil.check_overlap(downsampled_new_target_pts, down_sampled_model_pts, overlap_area_threshold = curr_overlap_area_threshold, voxel_size=voxel_threshold, require_new_added_pts_num = True) + overlap, _ = ReconstructionUtil.check_overlap(downsampled_new_target_pts, down_sampled_model_pts, overlap_area_threshold = curr_overlap_area_threshold, voxel_size=voxel_threshold, require_new_added_pts_num = True) if not overlap: retry += 1 retry_overlap_pose.append(pred_pose.cpu().numpy().tolist()) @@ -170,31 +176,22 @@ class Inferencer(Runner): continue start_time = time.time() - pred_cr, covered_pts_num = self.compute_coverage_rate(scanned_view_pts, new_target_pts, down_sampled_model_pts, threshold=voxel_threshold) + pred_cr, _ = self.compute_coverage_rate(scanned_view_pts, new_target_pts, down_sampled_model_pts, threshold=voxel_threshold) end_time = time.time() print(f"Time taken for coverage rate computation: {end_time - start_time} seconds") print(pred_cr, last_pred_cr, " max: ", data["seq_max_coverage_rate"]) - print("new added pts num: ", new_added_pts_num) if pred_cr >= data["seq_max_coverage_rate"] - 1e-3: print("max coverage rate reached!: ", pred_cr) success += 1 - elif new_added_pts_num < 5: - #success += 1 - print("min added pts num reached!: ", new_added_pts_num) - if pred_cr <= last_pred_cr + cr_increase_threshold: - retry += 1 - retry_duplication_pose.append(pred_pose.cpu().numpy().tolist()) - continue + retry = 0 pred_cr_seq.append(pred_cr) scanned_view_pts.append(new_target_pts) - down_sampled_new_pts_world = PtsUtil.random_downsample_point_cloud(new_target_pts, input_pts_N) - new_pts = down_sampled_new_pts_world input_data["scanned_n_to_world_pose_9d"] = [torch.cat([input_data["scanned_n_to_world_pose_9d"][0], pred_pose_9d], dim=0)] - combined_scanned_pts = np.concatenate([input_data["combined_scanned_pts"][0].cpu().numpy(), new_pts], axis=0) + combined_scanned_pts = np.vstack(scanned_view_pts) voxel_downsampled_combined_scanned_pts_np = PtsUtil.voxel_downsample_point_cloud(combined_scanned_pts, 0.002) random_downsampled_combined_scanned_pts_np = PtsUtil.random_downsample_point_cloud(voxel_downsampled_combined_scanned_pts_np, input_pts_N) input_data["combined_scanned_pts"] = torch.tensor(random_downsampled_combined_scanned_pts_np, dtype=torch.float32).unsqueeze(0).to(self.device) @@ -202,29 +199,12 @@ class Inferencer(Runner): if success > 3: break last_pred_cr = pred_cr - - input_data["scanned_n_to_world_pose_9d"] = input_data["scanned_n_to_world_pose_9d"][0].cpu().numpy().tolist() - result = { - "pred_pose_9d_seq": input_data["scanned_n_to_world_pose_9d"], - "combined_scanned_pts": input_data["combined_scanned_pts"], - "target_pts_seq": scanned_view_pts, - "coverage_rate_seq": pred_cr_seq, - "max_coverage_rate": data["seq_max_coverage_rate"], - "pred_max_coverage_rate": max(pred_cr_seq), - "scene_name": scene_name, - "retry_no_pts_pose": retry_no_pts_pose, - "retry_duplication_pose": retry_duplication_pose, - "retry_overlap_pose": retry_overlap_pose, - "best_seq_len": data["best_seq_len"], - } - self.stat_result[scene_name] = { - "coverage_rate_seq": pred_cr_seq, - "pred_max_coverage_rate": max(pred_cr_seq), - "pred_seq_len": len(pred_cr_seq), - } - print('success rate: ', max(pred_cr_seq)) - - return result + pts_num = voxel_downsampled_combined_scanned_pts_np.shape[0] + if pts_num - last_pts_num < 10 and pred_cr < data["seq_max_coverage_rate"] - 1e-3: + retry += 1 + retry_duplication_pose.append(pred_pose.cpu().numpy().tolist()) + print("delta pts num < 10:", pts_num, last_pts_num) + last_pts_num = pts_num def compute_coverage_rate(self, scanned_view_pts, new_pts, model_pts, threshold=0.005): if new_pts is not None: diff --git a/utils/render.py b/utils/render.py index 504a10d..9d23cc3 100644 --- a/utils/render.py +++ b/utils/render.py @@ -86,9 +86,10 @@ class RenderUtil: json.dump(params, f) start_time = time.time() result = subprocess.run([ - 'blender', '-b', '-P', script_path, '--', temp_dir + '/home/hofee/blender-4.0.2-linux-x64/blender', '-b', '-P', script_path, '--', temp_dir ], capture_output=True, text=True) end_time = time.time() + print(f"-- Time taken for blender: {end_time - start_time} seconds") path = os.path.join(temp_dir, "tmp") cam_info = DataLoadUtil.load_cam_info(path, binocular=True)