diff --git a/configs/local/inference_config.yaml b/configs/local/inference_config.yaml index f8a33a4..8d78c32 100644 --- a/configs/local/inference_config.yaml +++ b/configs/local/inference_config.yaml @@ -14,17 +14,17 @@ runner: dataset_list: - OmniObject3d_test - blender_script_path: "/data/hofee/project/nbv_rec/blender/data_renderer.py" - output_dir: "/data/hofee/data/inference_global_full_on_testset" + 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" pipeline: nbv_reconstruction_pipeline voxel_size: 0.003 dataset: OmniObject3d_train: - root_dir: "/data/hofee/data/new_full_data" - model_dir: "/data/hofee/data/scaled_object_meshes" - source: seq_reconstruction_dataset - split_file: "/data/hofee/data/sample.txt" + root_dir: "C:\\Document\\Datasets\\inference_test" + model_dir: "C:\\Document\\Datasets\\scaled_object_meshes" + source: seq_reconstruction_dataset_preprocessed + split_file: "C:\\Document\\Datasets\\data_list\\sample.txt" type: test filter_degree: 75 ratio: 1 @@ -34,10 +34,10 @@ dataset: load_from_preprocess: True OmniObject3d_test: - root_dir: "/data/hofee/data/new_full_data" - model_dir: "/data/hofee/data/scaled_object_meshes" - source: seq_reconstruction_dataset - split_file: "/data/hofee/data/new_full_data_list/OmniObject3d_test.txt" + root_dir: "C:\\Document\\Datasets\\inference_test" + model_dir: "C:\\Document\\Datasets\\scaled_object_meshes" + source: seq_reconstruction_dataset_preprocessed + 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 1c129bc..68774d5 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"/data/hofee/project/nbv_rec/nbv_reconstruction") +sys.path.append(r"C:\Document\Local Project\nbv_rec\nbv_reconstruction") from utils.data_load import DataLoadUtil from utils.pose import PoseUtil @@ -55,7 +55,9 @@ class SeqReconstructionDataset(BaseDataset): def get_datalist(self): datalist = [] - for scene_name in self.scene_name_list: + 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 = [] @@ -182,17 +184,33 @@ if __name__ == "__main__": 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": "/data/hofee/data/new_full_data", + "root_dir": "H:\\AI\\Datasets\\packed_test_data", "source": "seq_reconstruction_dataset", - "split_file": "/data/hofee/data/sample.txt", + "split_file": "H:\\AI\\Datasets\\data_list\\OmniObject3d_test.txt", "load_from_preprocess": True, - "ratio": 0.5, - "batch_size": 2, + "ratio": 1, "filter_degree": 75, "num_workers": 0, - "pts_num": 4096, - "type": namespace.Mode.TRAIN, + "pts_num": 8192, + "type": "test", } ds = SeqReconstructionDataset(config) print(len(ds)) diff --git a/core/seq_dataset_preprocessed.py b/core/seq_dataset_preprocessed.py new file mode 100644 index 0000000..9f84bf4 --- /dev/null +++ b/core/seq_dataset_preprocessed.py @@ -0,0 +1,85 @@ +import numpy as np +from PytorchBoot.dataset import BaseDataset +import PytorchBoot.namespace as namespace +import PytorchBoot.stereotype as stereotype +from PytorchBoot.config import ConfigManager +from PytorchBoot.utils.log_util import Log +import pickle +import torch +import os +import sys + +sys.path.append(r"C:\Document\Local Project\nbv_rec\nbv_reconstruction") + +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.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 + "scene_name": data["scene_name"], # String + "gt_pts": np.asarray(data["gt_pts"], dtype=np.float32), # Ndarray(N x 3) + "scene_path": os.path.join(self.real_root_dir, data["scene_name"]), # String + "O_to_L_pose": np.asarray(data["O_to_L_pose"], dtype=np.float32), + } + return data_item + + def __len__(self): + return len(self.item_list) + + +# -------------- Debug ---------------- # +if __name__ == "__main__": + import torch + + 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", + "source": "seq_reconstruction_dataset", + "split_file": "H:\\AI\\Datasets\\data_list\\OmniObject3d_test.txt", + "load_from_preprocess": True, + "ratio": 1, + "filter_degree": 75, + "num_workers": 0, + "pts_num": 8192, + "type": "test", + } + ds = SeqReconstructionDataset(config) + print(len(ds)) + print(ds.__getitem__(10)) + diff --git a/runners/inferencer.py b/runners/inferencer.py index 50f838c..0b990b5 100644 --- a/runners/inferencer.py +++ b/runners/inferencer.py @@ -19,7 +19,7 @@ from PytorchBoot.dataset import BaseDataset from PytorchBoot.runners.runner import Runner from PytorchBoot.utils import Log from PytorchBoot.status import status_manager - +from utils.data_load import DataLoadUtil @stereotype.runner("inferencer") class Inferencer(Runner): def __init__(self, config_path): @@ -35,7 +35,12 @@ class Inferencer(Runner): ''' Experiment ''' self.load_experiment("nbv_evaluator") - self.stat_result = {} + self.stat_result_path = os.path.join(self.output_dir, "stat.json") + if os.path.exists(self.stat_result_path): + with open(self.stat_result_path, "r") as f: + self.stat_result = json.load(f) + else: + self.stat_result = {} ''' Test ''' self.test_config = ConfigManager.get(namespace.Stereotype.RUNNER, namespace.Mode.TEST) @@ -68,22 +73,21 @@ class Inferencer(Runner): test_set_name = test_set.get_name() total=int(len(test_set)) - scene_name_list = test_set.get_scene_name_list() - for i in range(total): - scene_name = scene_name_list[i] + for i in tqdm(range(total), desc=f"Processing {test_set_name}", ncols=100): + 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 - data = test_set.__getitem__(i) + status_manager.set_progress("inference", "inferencer", f"Batch[{test_set_name}]", i+1, total) - scene_name = data["scene_name"] 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, max_iter=50, max_retry=5): + def predict_sequence(self, data, cr_increase_threshold=0.001, overlap_area_threshold=25, scan_points_threshold=10, max_iter=50, max_retry = 7): scene_name = data["scene_name"] Log.info(f"Processing scene: {scene_name}") status_manager.set_status("inference", "inferencer", "scene", scene_name) @@ -102,16 +106,23 @@ class Inferencer(Runner): ''' data for inference ''' input_data = {} + scanned_pts = [] input_data["combined_scanned_pts"] = torch.tensor(data["first_scanned_pts"][0], dtype=torch.float32).to(self.device).unsqueeze(0) 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] - first_frame_target_pts, first_frame_target_normals = RenderUtil.render_pts(first_frame_to_world, scene_path, self.script_path, voxel_threshold=voxel_threshold, filter_degree=filter_degree, nO_to_nL_pose=O_to_L_pose) + 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] last_pred_cr, added_pts_num = self.compute_coverage_rate(scanned_view_pts, None, down_sampled_model_pts, threshold=voxel_threshold) - + scanned_pts.append(first_frame_target_pts) retry_duplication_pose = [] retry_no_pts_pose = [] + retry_overlap_pose = [] retry = 0 pred_cr_seq = [last_pred_cr] success = 0 @@ -129,7 +140,22 @@ class Inferencer(Runner): try: start_time = time.time() - new_target_pts, new_target_normals = RenderUtil.render_pts(pred_pose, scene_path, self.script_path, voxel_threshold=voxel_threshold, filter_degree=filter_degree, nO_to_nL_pose=O_to_L_pose) + 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): + curr_overlap_area_threshold = overlap_area_threshold + else: + 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) + if not overlap: + retry += 1 + retry_overlap_pose.append(pred_pose.cpu().numpy().tolist()) + continue + + scanned_pts.append(new_target_pts) + 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: @@ -147,14 +173,16 @@ class Inferencer(Runner): continue start_time = time.time() - pred_cr, new_added_pts_num = self.compute_coverage_rate(scanned_view_pts, new_target_pts, down_sampled_model_pts, threshold=voxel_threshold) + pred_cr, covered_pts_num = 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 < 10: + 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 @@ -180,6 +208,7 @@ class Inferencer(Runner): input_data["scanned_n_to_world_pose_9d"] = input_data["scanned_n_to_world_pose_9d"][0].cpu().numpy().tolist() result = { + "scanned_pts": scanned_pts, "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, @@ -189,6 +218,7 @@ class Inferencer(Runner): "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] = { @@ -216,7 +246,7 @@ class Inferencer(Runner): os.makedirs(dataset_dir) output_path = os.path.join(dataset_dir, f"{scene_name}.pkl") pickle.dump(output, open(output_path, "wb")) - with open(os.path.join(dataset_dir, "stat.json"), "w") as f: + with open(self.stat_result_path, "w") as f: json.dump(self.stat_result, f) diff --git a/utils/reconstruction.py b/utils/reconstruction.py index c478ad4..6645ee9 100644 --- a/utils/reconstruction.py +++ b/utils/reconstruction.py @@ -32,13 +32,15 @@ class ReconstructionUtil: @staticmethod - def check_overlap(new_point_cloud, combined_point_cloud, overlap_area_threshold=25, voxel_size=0.01): + def check_overlap(new_point_cloud, combined_point_cloud, overlap_area_threshold=25, voxel_size=0.01, require_new_added_pts_num=False): kdtree = cKDTree(combined_point_cloud) distances, _ = kdtree.query(new_point_cloud) - overlapping_points = np.sum(distances < voxel_size*2) + overlapping_points_num = 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 + overlap_area = overlapping_points_num * voxel_size_cm * voxel_size_cm + if require_new_added_pts_num: + return overlap_area > overlap_area_threshold, len(new_point_cloud)-np.sum(distances < voxel_size*1.2) return overlap_area > overlap_area_threshold diff --git a/utils/render.py b/utils/render.py index 5c1f5a2..504a10d 100644 --- a/utils/render.py +++ b/utils/render.py @@ -54,7 +54,22 @@ class RenderUtil: return points_camera_world @staticmethod - def render_pts(cam_pose, scene_path, script_path, voxel_threshold=0.005, filter_degree=75, nO_to_nL_pose=None, require_full_scene=False): + def get_scan_points_indices(scan_points, mask, display_table_mask_label, cam_intrinsic, cam_extrinsic): + scan_points_homogeneous = np.hstack((scan_points, np.ones((scan_points.shape[0], 1)))) + points_camera = np.dot(np.linalg.inv(cam_extrinsic), scan_points_homogeneous.T).T[:, :3] + points_image_homogeneous = np.dot(cam_intrinsic, points_camera.T).T + points_image_homogeneous /= points_image_homogeneous[:, 2:] + pixel_x = points_image_homogeneous[:, 0].astype(int) + pixel_y = points_image_homogeneous[:, 1].astype(int) + h, w = mask.shape[:2] + valid_indices = (pixel_x >= 0) & (pixel_x < w) & (pixel_y >= 0) & (pixel_y < h) + mask_colors = mask[pixel_y[valid_indices], pixel_x[valid_indices]] + selected_points_indices = np.where((mask_colors == display_table_mask_label).all(axis=-1))[0] + selected_points_indices = np.where(valid_indices)[0][selected_points_indices] + return selected_points_indices + + @staticmethod + def render_pts(cam_pose, scene_path, script_path, scan_points, voxel_threshold=0.005, filter_degree=75, nO_to_nL_pose=None, require_full_scene=False): nO_to_world_pose = DataLoadUtil.get_real_cam_O_from_cam_L(cam_pose, nO_to_nL_pose, scene_path=scene_path) @@ -74,12 +89,7 @@ class RenderUtil: 'blender', '-b', '-P', script_path, '--', temp_dir ], capture_output=True, text=True) end_time = time.time() - print(result) print(f"-- Time taken for blender: {end_time - start_time} seconds") - if result.returncode != 0: - print("Blender script failed:") - print(result.stderr) - return None path = os.path.join(temp_dir, "tmp") cam_info = DataLoadUtil.load_cam_info(path, binocular=True) depth_L, depth_R = DataLoadUtil.load_depth( @@ -116,11 +126,15 @@ class RenderUtil: target_points, target_normals = PtsUtil.filter_points( target_points, sampled_target_normal_L, cam_info["cam_to_world"], theta_limit = filter_degree, z_range=(RenderUtil.min_z, RenderUtil.max_z) ) + + scan_points_indices_L = RenderUtil.get_scan_points_indices(scan_points, mask_img_L, RenderUtil.display_table_mask_label, cam_info["cam_intrinsic"], cam_info["cam_to_world"]) + scan_points_indices_R = RenderUtil.get_scan_points_indices(scan_points, mask_img_R, RenderUtil.display_table_mask_label, cam_info["cam_intrinsic"], cam_info["cam_to_world_R"]) + scan_points_indices = np.intersect1d(scan_points_indices_L, scan_points_indices_R) 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 \ No newline at end of file + return target_points, target_normals, scan_points_indices \ No newline at end of file