From 2c8ef203218cf6cec284b27104994499d32b5a09 Mon Sep 17 00:00:00 2001 From: hofee Date: Wed, 20 Nov 2024 15:24:45 +0800 Subject: [PATCH] upd ab_global_only --- configs/local/inference_config.yaml | 12 ++--- core/seq_dataset.py | 84 ++++++++++++----------------- core/seq_dataset_preprocessed.py | 5 +- runners/inferencer.py | 69 ++++++++++++------------ utils/render.py | 9 +--- 5 files changed, 80 insertions(+), 99 deletions(-) diff --git a/configs/local/inference_config.yaml b/configs/local/inference_config.yaml index 1e6d89b..64bccbf 100644 --- a/configs/local/inference_config.yaml +++ b/configs/local/inference_config.yaml @@ -14,11 +14,11 @@ runner: dataset_list: - OmniObject3d_test - blender_script_path: "C:\\Document\\Local Project\\nbv_rec\\blender\\data_renderer.py" - output_dir: "C:\\Document\\Datasets\\debug_output" + 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 - + min_new_area: 1.0 dataset: # OmniObject3d_train: # root_dir: "C:\\Document\\Datasets\\inference_test1" @@ -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/core/seq_dataset.py b/core/seq_dataset.py index 68774d5..e27563b 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,6 +47,8 @@ class SeqReconstructionDataset(BaseDataset): with open(self.split_file_path, "r") as f: for line in f: scene_name = line.strip() + if not os.path.exists(os.path.join(self.root_dir, scene_name)): + continue scene_name_list.append(scene_name) return scene_name_list @@ -58,29 +60,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 +123,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 +135,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 +152,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 @@ -180,39 +170,35 @@ class SeqReconstructionDataset(BaseDataset): # -------------- Debug ---------------- # if __name__ == "__main__": 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) diff --git a/core/seq_dataset_preprocessed.py b/core/seq_dataset_preprocessed.py index 9f84bf4..c92b5db 100644 --- a/core/seq_dataset_preprocessed.py +++ b/core/seq_dataset_preprocessed.py @@ -15,21 +15,19 @@ from utils.data_load import DataLoadUtil from utils.pose import PoseUtil from utils.pts import PtsUtil - @stereotype.dataset("seq_reconstruction_dataset_preprocessed") class SeqReconstructionDatasetPreprocessed(BaseDataset): def __init__(self, config): 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 @@ -43,7 +41,6 @@ class SeqReconstructionDatasetPreprocessed(BaseDataset): def __len__(self): return len(self.item_list) - # -------------- Debug ---------------- # if __name__ == "__main__": import torch diff --git a/runners/inferencer.py b/runners/inferencer.py index 238da0b..5d0cf16 100644 --- a/runners/inferencer.py +++ b/runners/inferencer.py @@ -23,11 +23,15 @@ from utils.data_load import DataLoadUtil @stereotype.runner("inferencer") class Inferencer(Runner): def __init__(self, config_path): + super().__init__(config_path) self.script_path = ConfigManager.get(namespace.Stereotype.RUNNER, "blender_script_path") self.output_dir = ConfigManager.get(namespace.Stereotype.RUNNER, "output_dir") self.voxel_size = ConfigManager.get(namespace.Stereotype.RUNNER, "voxel_size") + self.min_new_area = ConfigManager.get(namespace.Stereotype.RUNNER, "min_new_area") + CM = 0.01 + self.min_new_pts_num = self.min_new_area * (CM / self.voxel_size) **2 ''' Pipeline ''' self.pipeline_name = self.config[namespace.Stereotype.PIPELINE] self.pipeline:torch.nn.Module = ComponentFactory.create(namespace.Stereotype.PIPELINE, self.pipeline_name) @@ -74,22 +78,24 @@ class Inferencer(Runner): total=int(len(test_set)) 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-suitcase_001": + try: + data = test_set.__getitem__(i) + scene_name = data["scene_name"] + 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}") + continue + + status_manager.set_progress("inference", "inferencer", f"Batch[{test_set_name}]", i+1, total) + output = self.predict_sequence(data) + self.save_inference_result(test_set_name, data["scene_name"], output) + except Exception as e: + Log.error(f"Error in scene {scene_name}, {e}") 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}") - continue - - status_manager.set_progress("inference", "inferencer", f"Batch[{test_set_name}]", i+1, total) - output = self.predict_sequence(data) - self.save_inference_result(test_set_name, data["scene_name"], output) 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 = 5): + def predict_sequence(self, data, cr_increase_threshold=0, overlap_area_threshold=25, scan_points_threshold=10, max_iter=50, max_retry = 10, max_success=3): scene_name = data["scene_name"] Log.info(f"Processing scene: {scene_name}") status_manager.set_status("inference", "inferencer", "scene", scene_name) @@ -128,13 +134,11 @@ 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] + last_pts_num = PtsUtil.voxel_downsample_point_cloud(data["first_scanned_pts"][0], voxel_threshold).shape[0] import time - while len(pred_cr_seq) < max_iter and retry < max_retry: - start_time = time.time() + while len(pred_cr_seq) < max_iter and retry < max_retry and success < max_success: + Log.green(f"iter: {len(pred_cr_seq)}, retry: {retry}/{max_retry}, success: {success}/{max_success}") output = self.pipeline(input_data) - end_time = time.time() - print(f"Time taken for inference: {end_time - start_time} seconds") pred_pose_9d = output["pred_pose_9d"] pred_pose = torch.eye(4, device=pred_pose_9d.device) @@ -142,7 +146,6 @@ class Inferencer(Runner): pred_pose[:3,3] = pred_pose_9d[0,6:] try: - start_time = time.time() new_target_pts, new_target_normals, new_scan_points_indices = RenderUtil.render_pts(pred_pose, scene_path, self.script_path, scan_points, voxel_threshold=voxel_threshold, filter_degree=filter_degree, nO_to_nL_pose=O_to_L_pose) #import ipdb; ipdb.set_trace() if not ReconstructionUtil.check_scan_points_overlap(history_indices, new_scan_points_indices, scan_points_threshold): @@ -153,15 +156,14 @@ class Inferencer(Runner): downsampled_new_target_pts = PtsUtil.voxel_downsample_point_cloud(new_target_pts, voxel_threshold) 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: + Log.yellow("no overlap!") retry += 1 retry_overlap_pose.append(pred_pose.cpu().numpy().tolist()) continue history_indices.append(new_scan_points_indices) - end_time = time.time() - print(f"Time taken for rendering: {end_time - start_time} seconds") except Exception as e: - Log.warning(f"Error in scene {scene_path}, {e}") + Log.error(f"Error in scene {scene_path}, {e}") print("current pose: ", pred_pose) print("curr_pred_cr: ", last_pred_cr) retry_no_pts_pose.append(pred_pose.cpu().numpy().tolist()) @@ -169,40 +171,41 @@ class Inferencer(Runner): continue if new_target_pts.shape[0] == 0: - print("no pts in new target") + Log.red("no pts in new target") retry_no_pts_pose.append(pred_pose.cpu().numpy().tolist()) retry += 1 continue - start_time = time.time() 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"]) + Log.yellow(f"{pred_cr}, {last_pred_cr}, max: , {data['seq_max_coverage_rate']}") if pred_cr >= data["seq_max_coverage_rate"] - 1e-3: print("max coverage rate reached!: ", pred_cr) - success += 1 + - retry = 0 pred_cr_seq.append(pred_cr) scanned_view_pts.append(new_target_pts) 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.vstack(scanned_view_pts) - voxel_downsampled_combined_scanned_pts_np = PtsUtil.voxel_downsample_point_cloud(combined_scanned_pts, 0.002) + voxel_downsampled_combined_scanned_pts_np = PtsUtil.voxel_downsample_point_cloud(combined_scanned_pts, voxel_threshold) 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) - if success > 3: - break + last_pred_cr = pred_cr 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: + Log.info(f"delta pts num:,{pts_num - last_pts_num },{pts_num}, {last_pts_num}") + + if pts_num - last_pts_num < self.min_new_pts_num and pred_cr <= data["seq_max_coverage_rate"] - 1e-2: retry += 1 retry_duplication_pose.append(pred_pose.cpu().numpy().tolist()) - print("delta pts num < 10:", pts_num, last_pts_num) + Log.red(f"delta pts num < {self.min_new_pts_num}:, {pts_num}, {last_pts_num}") + elif pts_num - last_pts_num < self.min_new_pts_num and pred_cr > data["seq_max_coverage_rate"] - 1e-2: + success += 1 + Log.success(f"delta pts num < {self.min_new_pts_num}:, {pts_num}, {last_pts_num}") + last_pts_num = pts_num diff --git a/utils/render.py b/utils/render.py index 504a10d..aa2f4e2 100644 --- a/utils/render.py +++ b/utils/render.py @@ -84,12 +84,10 @@ class RenderUtil: params_data_path = os.path.join(temp_dir, "params.json") with open(params_data_path, 'w') as f: 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") + # print(result) path = os.path.join(temp_dir, "tmp") cam_info = DataLoadUtil.load_cam_info(path, binocular=True) depth_L, depth_R = DataLoadUtil.load_depth( @@ -97,7 +95,6 @@ class RenderUtil: cam_info["far_plane"], binocular=True ) - start_time = time.time() mask_L, mask_R = DataLoadUtil.load_seg(path, binocular=True) normal_L = DataLoadUtil.load_normal(path, binocular=True, left_only=True) ''' target points ''' @@ -134,7 +131,5 @@ class RenderUtil: if not has_points: target_points = np.zeros((0, 3)) target_normals = np.zeros((0, 3)) - end_time = time.time() - print(f"-- Time taken for processing: {end_time - start_time} seconds") #import ipdb; ipdb.set_trace() return target_points, target_normals, scan_points_indices \ No newline at end of file