diff --git a/configs/local/inference_config.yaml b/configs/local/inference_config.yaml index f8a33a4..64bccbf 100644 --- a/configs/local/inference_config.yaml +++ b/configs/local/inference_config.yaml @@ -14,30 +14,30 @@ 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: "/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: "/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" - type: test - filter_degree: 75 - ratio: 1 - batch_size: 1 - num_workers: 12 - pts_num: 8192 - load_from_preprocess: True + # OmniObject3d_train: + # root_dir: "C:\\Document\\Datasets\\inference_test1" + # 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 + # batch_size: 1 + # num_workers: 12 + # pts_num: 8192 + # 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: "/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" 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..4396de7 100644 --- a/configs/local/view_generate_config.yaml +++ b/configs/local/view_generate_config.yaml @@ -8,11 +8,11 @@ runner: root_dir: experiments generate: 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 + from: 1 + to: 50 # -1 means all + object_dir: C:\\Document\\Datasets\\scaled_object_meshes + table_model_path: C:\\Document\\Datasets\\table.obj + output_dir: C:\\Document\\Datasets\\debug_generate_view binocular_vision: true plane_size: 10 max_views: 512 diff --git a/core/seq_dataset.py b/core/seq_dataset.py index 7fd50f3..e91a52f 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"/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 @@ -55,33 +57,23 @@ class SeqReconstructionDataset(BaseDataset): def get_datalist(self): - datalist = [] + datalist = [] 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 @@ -132,8 +124,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) @@ -145,7 +136,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]) ) @@ -162,7 +153,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 @@ -190,9 +180,9 @@ if __name__ == "__main__": np.random.seed(seed) config = { - "root_dir": "/data/hofee/data/new_full_data", + "root_dir": "/media/hofee/data/data/new_testset", "source": "seq_reconstruction_dataset", - "split_file": "/data/hofee/data/new_full_data_list/OmniObject3d_test.txt", + "split_file": "/media/hofee/data/data/OmniObject3d_test.txt", "load_from_preprocess": True, "filter_degree": 75, "num_workers": 0, @@ -200,22 +190,16 @@ if __name__ == "__main__": "type": namespace.Mode.TEST, } - output_dir = "/data/hofee/trash_can/output_inference_test" - new_output_dir = "/data/hofee/inference_test" + output_dir = "/media/hofee/data/data/new_testset_output" os.makedirs(output_dir, exist_ok=True) - os.makedirs(new_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") - if os.path.exists(output_path): - item = pickle.load(open(output_path, "rb")) - else: - item = ds.__getitem__(i) + item = ds.__getitem__(i) for key, value in item.items(): if isinstance(value, np.ndarray): item[key] = value.tolist() - new_output_path = os.path.join(new_output_dir, f"item_{i}.pkl") - with open(new_output_path, "wb") as f: + #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 new file mode 100644 index 0000000..c92b5db --- /dev/null +++ b/core/seq_dataset_preprocessed.py @@ -0,0 +1,82 @@ +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"/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_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/inference_server.py b/runners/inference_server.py index 4ad7ed1..ac62fc9 100644 --- a/runners/inference_server.py +++ b/runners/inference_server.py @@ -41,17 +41,8 @@ class InferencerServer(Runner): fps_downsampled_combined_scanned_pts, fps_idx = PtsUtil.fps_downsample_point_cloud( voxel_downsampled_combined_scanned_pts, self.pts_num, require_idx=True ) - # combined_scanned_views_pts_mask = np.zeros(len(scanned_pts), dtype=np.uint8) - # start_idx = 0 - # for i in range(len(scanned_pts)): - # end_idx = start_idx + len(scanned_pts[i]) - # combined_scanned_views_pts_mask[start_idx:end_idx] = i - # start_idx = end_idx - - # fps_downsampled_combined_scanned_pts_mask = combined_scanned_views_pts_mask[fps_idx] input_data["scanned_pts"] = scanned_pts - # input_data["scanned_pts_mask"] = np.asarray(fps_downsampled_combined_scanned_pts_mask, dtype=np.uint8) input_data["scanned_n_to_world_pose_9d"] = np.asarray(scanned_n_to_world_pose_9d, dtype=np.float32) input_data["combined_scanned_pts"] = np.asarray(fps_downsampled_combined_scanned_pts, dtype=np.float32) return input_data diff --git a/runners/inferencer.py b/runners/inferencer.py index 50f838c..5d0cf16 100644 --- a/runners/inferencer.py +++ b/runners/inferencer.py @@ -19,15 +19,19 @@ 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): + 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) @@ -35,7 +39,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 +77,25 @@ 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] - 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}") + for i in tqdm(range(total), desc=f"Processing {test_set_name}", ncols=100): + 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 - 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, 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) @@ -106,21 +118,27 @@ 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] - 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) - scanned_view_pts = [first_frame_target_pts] - last_pred_cr, added_pts_num = self.compute_coverage_rate(scanned_view_pts, None, down_sampled_model_pts, threshold=voxel_threshold) + + 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) retry_duplication_pose = [] retry_no_pts_pose = [] + retry_overlap_pose = [] retry = 0 pred_cr_seq = [last_pred_cr] success = 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) @@ -128,12 +146,24 @@ class Inferencer(Runner): pred_pose[:3,3] = pred_pose_9d[0,6:] 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) - end_time = time.time() - print(f"Time taken for rendering: {end_time - start_time} seconds") + 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, _ = 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) 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()) @@ -141,42 +171,43 @@ 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, new_added_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"]) + pred_cr, _ = self.compute_coverage_rate(scanned_view_pts, new_target_pts, down_sampled_model_pts, threshold=voxel_threshold) + 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 - elif new_added_pts_num < 10: - 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) - voxel_downsampled_combined_scanned_pts_np = PtsUtil.voxel_downsample_point_cloud(combined_scanned_pts, 0.002) + combined_scanned_pts = np.vstack(scanned_view_pts) + 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] + 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()) + 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 + input_data["scanned_n_to_world_pose_9d"] = input_data["scanned_n_to_world_pose_9d"][0].cpu().numpy().tolist() result = { @@ -189,6 +220,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 +248,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/runners/view_generator.py b/runners/view_generator.py index 634ccbf..2e28c44 100644 --- a/runners/view_generator.py +++ b/runners/view_generator.py @@ -9,7 +9,7 @@ class ViewGenerator(Runner): self.config_path = config_path def run(self): - result = subprocess.run(['/home/hofee/blender-4.0.2-linux-x64/blender', '-b', '-P', '../blender/run_blender.py', '--', self.config_path]) + result = subprocess.run(['blender', '-b', '-P', '../blender/run_blender.py', '--', self.config_path]) print() def create_experiment(self, backup_name=None): 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..aa2f4e2 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) @@ -69,17 +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(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 + # 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( @@ -87,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 ''' @@ -116,11 +123,13 @@ 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