diff --git a/app_generate_strategy.py b/app_generate_strategy.py index f56d745..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/strategy_generate_config.yaml").run() + StrategyGenerator("configs/local/strategy_generate_config.yaml").run() \ No newline at end of file diff --git a/app_generate_view.py b/app_generate_view.py index e016f09..a48d0ba 100644 --- a/app_generate_view.py +++ b/app_generate_view.py @@ -13,4 +13,4 @@ class ViewGenerateApp: Trainer("path_to_your_train_config").run() Evaluator("path_to_your_eval_config").run() ''' - ViewGenerator("./configs/view_generate_config.yaml").run() + ViewGenerator("configs/local/view_generate_config.yaml").run() diff --git a/configs/local/inference_config.yaml b/configs/local/inference_config.yaml index f8fed9a..4cc4e01 100644 --- a/configs/local/inference_config.yaml +++ b/configs/local/inference_config.yaml @@ -8,19 +8,19 @@ runner: experiment: name: local_eval root_dir: "experiments" - epoch: 600 # -1 stands for last epoch + epoch: 555 # -1 stands for last epoch test: dataset_list: - OmniObject3d_train blender_script_path: "/media/hofee/data/project/python/nbv_reconstruction/blender/data_renderer.py" - output_dir: "/media/hofee/data/project/python/nbv_reconstruction/sample_for_training/inference_result" + output_dir: "/media/hofee/data/project/python/nbv_reconstruction/sample_for_training/inference_result2" pipeline: nbv_reconstruction_pipeline dataset: OmniObject3d_train: - root_dir: "/media/hofee/data/project/python/nbv_reconstruction/sample_for_training/scenes" + root_dir: "/media/hofee/data/project/python/nbv_reconstruction/sample_for_training/sample_preprocessed_scenes" model_dir: "/media/hofee/data/data/scaled_object_meshes" source: seq_nbv_reconstruction_dataset split_file: "/media/hofee/data/project/python/nbv_reconstruction/sample_for_training/OmniObject3d_train.txt" @@ -30,6 +30,7 @@ dataset: batch_size: 1 num_workers: 12 pts_num: 4096 + load_from_preprocess: True pipeline: nbv_reconstruction_pipeline: diff --git a/configs/local/strategy_generate_config.yaml b/configs/local/strategy_generate_config.yaml index c0dd913..5bcd80a 100644 --- a/configs/local/strategy_generate_config.yaml +++ b/configs/local/strategy_generate_config.yaml @@ -15,17 +15,19 @@ runner: overlap_threshold: 0.5 filter_degree: 75 to_specified_dir: True # if True, output_dir is used, otherwise, root_dir is used - save_points: False - save_best_combined_points: True + save_points: True + load_points: True + save_best_combined_points: False save_mesh: True overwrite: False + seq_num: 50 dataset_list: - OmniObject3d datasets: OmniObject3d: #"/media/hofee/data/data/temp_output" - root_dir: "/media/hofee/repository/nbv_reconstruction_data_512" + root_dir: "/media/hofee/data/data/sample_data/view_data" model_dir: "/media/hofee/data/data/scaled_object_meshes" #output_dir: "/media/hofee/data/data/label_output" diff --git a/configs/local/view_generate_config.yaml b/configs/local/view_generate_config.yaml index cb5b8a8..8b2e35e 100644 --- a/configs/local/view_generate_config.yaml +++ b/configs/local/view_generate_config.yaml @@ -9,7 +9,7 @@ runner: generate: object_dir: /media/hofee/data/data/scaled_object_meshes table_model_path: /media/hofee/data/data/others/table.obj - output_dir: /media/hofee/repository/nbv_reconstruction_data_512 + output_dir: /media/hofee/repository/new_nbv_reconstruction_data_512 binocular_vision: true plane_size: 10 max_views: 512 diff --git a/core/nbv_dataset.py b/core/nbv_dataset.py index 634a745..891a62c 100644 --- a/core/nbv_dataset.py +++ b/core/nbv_dataset.py @@ -102,9 +102,7 @@ class NBVReconstructionDataset(BaseDataset): max_coverage_rate = data_item_info["max_coverage_rate"] scene_name = data_item_info["scene_name"] scanned_views_pts, scanned_coverages_rate, scanned_n_to_world_pose = [], [], [] - first_frame_idx = scanned_views[0][0] - first_cam_info = DataLoadUtil.load_cam_info(DataLoadUtil.get_path(self.root_dir, scene_name, first_frame_idx), binocular=True) - first_frame_to_world = first_cam_info["cam_to_world"] + for view in scanned_views: frame_idx = view[0] coverage_rate = view[1] diff --git a/core/seq_dataset.py b/core/seq_dataset.py index 4cc4043..12aa763 100644 --- a/core/seq_dataset.py +++ b/core/seq_dataset.py @@ -28,6 +28,7 @@ class SeqNBVReconstructionDataset(BaseDataset): self.model_dir = config["model_dir"] self.filter_degree = config["filter_degree"] + self.load_from_preprocess = config.get("load_from_preprocess", False) def load_scene_name_list(self): @@ -38,10 +39,30 @@ class SeqNBVReconstructionDataset(BaseDataset): scene_name_list.append(scene_name) return scene_name_list + def get_datalist_new(self): + datalist = [] + for scene_name in self.scene_name_list: + label_num = DataLoadUtil.get_label_num(self.root_dir, scene_name) + for i in range(label_num): + label_path = DataLoadUtil.get_label_path(self.root_dir, scene_name, i) + label_data = DataLoadUtil.load_label(label_path) + best_seq = label_data["best_sequence"] + max_coverage_rate = label_data["max_coverage_rate"] + first_frame = best_seq[0] + best_seq_len = len(best_seq) + datalist.append({ + "scene_name": scene_name, + "first_frame": first_frame, + "max_coverage_rate": max_coverage_rate, + "best_seq_len": best_seq_len, + "label_idx": i, + }) + return datalist + def get_datalist(self): datalist = [] for scene_name in self.scene_name_list: - label_path = DataLoadUtil.get_label_path(self.root_dir, scene_name) + label_path = DataLoadUtil.get_label_path_old(self.root_dir, scene_name) label_data = DataLoadUtil.load_label(label_path) best_seq = label_data["best_sequence"] max_coverage_rate = label_data["max_coverage_rate"] @@ -52,8 +73,9 @@ class SeqNBVReconstructionDataset(BaseDataset): "first_frame": first_frame, "max_coverage_rate": max_coverage_rate, "best_seq_len": best_seq_len, + "best_seq": best_seq, }) - return datalist[5:] + return datalist def __getitem__(self, index): data_item_info = self.datalist[index] @@ -62,27 +84,27 @@ class SeqNBVReconstructionDataset(BaseDataset): max_coverage_rate = data_item_info["max_coverage_rate"] scene_name = data_item_info["scene_name"] first_cam_info = DataLoadUtil.load_cam_info(DataLoadUtil.get_path(self.root_dir, scene_name, first_frame_idx), binocular=True) - first_frame_to_world = first_cam_info["cam_to_world"] first_view_path = DataLoadUtil.get_path(self.root_dir, scene_name, first_frame_idx) first_left_cam_pose = first_cam_info["cam_to_world"] first_right_cam_pose = first_cam_info["cam_to_world_R"] first_center_cam_pose = first_cam_info["cam_to_world_O"] - - first_depth_L, first_depth_R = DataLoadUtil.load_depth(first_view_path, first_cam_info['near_plane'], first_cam_info['far_plane'], binocular=True) - first_L_to_L_pose = np.eye(4) - first_R_to_L_pose = np.dot(np.linalg.inv(first_left_cam_pose), first_right_cam_pose) - first_point_cloud_L = DataLoadUtil.get_point_cloud(first_depth_L, first_cam_info['cam_intrinsic'], first_L_to_L_pose)['points_world'] - first_point_cloud_R = DataLoadUtil.get_point_cloud(first_depth_R, first_cam_info['cam_intrinsic'], first_R_to_L_pose)['points_world'] - - first_point_cloud_L = PtsUtil.random_downsample_point_cloud(first_point_cloud_L, 65536) - first_point_cloud_R = PtsUtil.random_downsample_point_cloud(first_point_cloud_R, 65536) - first_overlap_points = DataLoadUtil.get_overlapping_points(first_point_cloud_L, first_point_cloud_R) - first_downsampled_target_point_cloud = PtsUtil.random_downsample_point_cloud(first_overlap_points, self.pts_num) - first_to_first_pose = np.eye(4) - first_to_first_rot_6d = PoseUtil.matrix_to_rotation_6d_numpy(np.asarray(first_to_first_pose[:3,:3])) - first_to_first_trans = first_to_first_pose[:3,3] - first_to_first_9d = np.concatenate([first_to_first_rot_6d, first_to_first_trans], axis=0) + if self.load_from_preprocess: + first_downsampled_target_point_cloud = DataLoadUtil.load_from_preprocessed_pts(first_view_path) + else: + first_depth_L, first_depth_R = DataLoadUtil.load_depth(first_view_path, first_cam_info['near_plane'], first_cam_info['far_plane'], binocular=True) + + first_point_cloud_L = DataLoadUtil.get_point_cloud(first_depth_L, first_cam_info['cam_intrinsic'], first_left_cam_pose)['points_world'] + first_point_cloud_R = DataLoadUtil.get_point_cloud(first_depth_R, first_cam_info['cam_intrinsic'], first_right_cam_pose)['points_world'] + + first_point_cloud_L = PtsUtil.random_downsample_point_cloud(first_point_cloud_L, 65536) + first_point_cloud_R = PtsUtil.random_downsample_point_cloud(first_point_cloud_R, 65536) + first_overlap_points = DataLoadUtil.get_overlapping_points(first_point_cloud_L, first_point_cloud_R) + first_downsampled_target_point_cloud = PtsUtil.random_downsample_point_cloud(first_overlap_points, self.pts_num) + + first_to_world_rot_6d = PoseUtil.matrix_to_rotation_6d_numpy(np.asarray(first_left_cam_pose[:3,:3])) + first_to_world_trans = first_left_cam_pose[:3,3] + first_to_world_9d = np.concatenate([first_to_world_rot_6d, first_to_world_trans], axis=0) diag = DataLoadUtil.get_bbox_diag(self.model_dir, scene_name) voxel_threshold = diag*0.02 first_O_to_first_L_pose = np.dot(np.linalg.inv(first_left_cam_pose), first_center_cam_pose) @@ -90,17 +112,17 @@ class SeqNBVReconstructionDataset(BaseDataset): model_points_normals = DataLoadUtil.load_points_normals(self.root_dir, scene_name) data_item = { "first_pts": np.asarray([first_downsampled_target_point_cloud],dtype=np.float32), - "first_to_first_9d": np.asarray([first_to_first_9d],dtype=np.float32), + "first_to_world_9d": np.asarray([first_to_world_9d],dtype=np.float32), "scene_name": scene_name, "max_coverage_rate": max_coverage_rate, "voxel_threshold": voxel_threshold, "filter_degree": self.filter_degree, - "first_frame_to_world": np.asarray(first_frame_to_world, dtype=np.float32), "O_to_L_pose": first_O_to_first_L_pose, "first_frame_coverage": first_frame_coverage, "scene_path": scene_path, "model_points_normals": model_points_normals, "best_seq_len": data_item_info["best_seq_len"], + "first_frame_id": first_frame_idx, } return data_item @@ -111,10 +133,9 @@ class SeqNBVReconstructionDataset(BaseDataset): def collate_fn(batch): collate_data = {} collate_data["first_pts"] = [torch.tensor(item['first_pts']) for item in batch] - collate_data["first_to_first_9d"] = [torch.tensor(item['first_to_first_9d']) for item in batch] - collate_data["first_frame_to_world"] = torch.stack([torch.tensor(item["first_frame_to_world"]) for item in batch]) + collate_data["first_to_world_9d"] = [torch.tensor(item['first_to_world_9d']) for item in batch] for key in batch[0].keys(): - if key not in ["first_pts", "first_to_first_9d", "first_frame_to_world"]: + if key not in ["first_pts", "first_to_world_9d"]: collate_data[key] = [item[key] for item in batch] return collate_data return collate_fn diff --git a/runners/inferencer.py b/runners/inferencer.py index aad0bc3..bfd20f1 100644 --- a/runners/inferencer.py +++ b/runners/inferencer.py @@ -92,12 +92,16 @@ class Inferencer(Runner): model_points_normals = data["model_points_normals"][0] model_pts = model_points_normals[:,:3] down_sampled_model_pts = PtsUtil.voxel_downsample_point_cloud(model_pts, voxel_threshold) - first_frame_to_world = data["first_frame_to_world"][0] + first_frame_to_world_9d = data["first_to_world_9d"][0] + first_frame_to_world = torch.eye(4, device=first_frame_to_world_9d.device) + first_frame_to_world[:3,:3] = PoseUtil.rotation_6d_to_matrix_tensor_batch(first_frame_to_world_9d[:,:6])[0] + first_frame_to_world[:3,3] = first_frame_to_world_9d[0,6:] + first_frame_to_world = first_frame_to_world.to(self.device) ''' data for inference ''' input_data = {} input_data["scanned_pts"] = [data["first_pts"][0].to(self.device)] - input_data["scanned_n_to_world_pose_9d"] = [data["first_frame_to_world"][0].to(self.device)] + input_data["scanned_n_to_world_pose_9d"] = [data["first_to_world_9d"][0].to(self.device)] input_data["mode"] = namespace.Mode.TEST input_pts_N = input_data["scanned_pts"][0].shape[1] @@ -113,20 +117,19 @@ class Inferencer(Runner): while len(pred_cr_seq) < max_iter and retry < max_retry: output = self.pipeline(input_data) - next_pose_9d = output["pred_pose_9d"] - pred_pose = torch.eye(4, device=next_pose_9d.device) + pred_pose_9d = output["pred_pose_9d"] + pred_pose = torch.eye(4, device=pred_pose_9d.device) - pred_pose[:3,:3] = PoseUtil.rotation_6d_to_matrix_tensor_batch(next_pose_9d[:,:6])[0] - pred_pose[:3,3] = next_pose_9d[0,6:] - pred_n_to_world_pose_mat = torch.matmul(first_frame_to_world, pred_pose) + pred_pose[:3,:3] = PoseUtil.rotation_6d_to_matrix_tensor_batch(pred_pose_9d[:,:6])[0] + pred_pose[:3,3] = pred_pose_9d[0,6:] try: - new_target_pts_world, new_pts_world = RenderUtil.render_pts(pred_n_to_world_pose_mat, scene_path, self.script_path, model_points_normals, voxel_threshold=voxel_threshold, filter_degree=filter_degree, nO_to_nL_pose=O_to_L_pose, require_full_scene=True) + new_target_pts_world, new_pts_world = RenderUtil.render_pts(pred_pose, scene_path, self.script_path, model_points_normals, voxel_threshold=voxel_threshold, filter_degree=filter_degree, nO_to_nL_pose=O_to_L_pose, require_full_scene=True) except Exception as e: Log.warning(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_n_to_world_pose_mat.cpu().numpy().tolist()) + retry_no_pts_pose.append(pred_pose.cpu().numpy().tolist()) retry += 1 continue @@ -138,7 +141,7 @@ class Inferencer(Runner): break if pred_cr <= last_pred_cr + cr_increase_threshold: retry += 1 - retry_duplication_pose.append(pred_n_to_world_pose_mat.cpu().numpy().tolist()) + retry_duplication_pose.append(pred_pose.cpu().numpy().tolist()) continue retry = 0 @@ -151,7 +154,7 @@ class Inferencer(Runner): new_pts_tensor = torch.tensor(new_pts, dtype=torch.float32).unsqueeze(0).to(self.device) input_data["scanned_pts"] = [torch.cat([input_data["scanned_pts"][0] , new_pts_tensor], dim=0)] - input_data["scanned_n_to_world_pose_9d"] = [torch.cat([input_data["scanned_n_to_world_pose_9d"][0], next_pose_9d], dim=0)] + input_data["scanned_n_to_world_pose_9d"] = [torch.cat([input_data["scanned_n_to_world_pose_9d"][0], pred_pose_9d], dim=0)] last_pred_cr = pred_cr diff --git a/runners/strategy_generator.py b/runners/strategy_generator.py index 5280178..d6c1598 100644 --- a/runners/strategy_generator.py +++ b/runners/strategy_generator.py @@ -25,15 +25,17 @@ class StrategyGenerator(Runner): self.to_specified_dir = ConfigManager.get("runner", "generate", "to_specified_dir") self.save_best_combined_pts = ConfigManager.get("runner", "generate", "save_best_combined_points") self.save_mesh = ConfigManager.get("runner", "generate", "save_mesh") + self.load_pts = ConfigManager.get("runner", "generate", "load_points") self.filter_degree = ConfigManager.get("runner", "generate", "filter_degree") self.overwrite = ConfigManager.get("runner", "generate", "overwrite") + self.save_pts = ConfigManager.get("runner","generate","save_points") + self.seq_num = ConfigManager.get("runner","generate","seq_num") def run(self): dataset_name_list = ConfigManager.get("runner", "generate", "dataset_list") voxel_threshold, overlap_threshold = ConfigManager.get("runner","generate","voxel_threshold"), ConfigManager.get("runner","generate","overlap_threshold") - self.save_pts = ConfigManager.get("runner","generate","save_points") 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)) @@ -48,7 +50,7 @@ class StrategyGenerator(Runner): diag = DataLoadUtil.get_bbox_diag(model_dir, scene_name) voxel_threshold = diag*0.02 status_manager.set_status("generate_strategy", "strategy_generator", "voxel_threshold", voxel_threshold) - output_label_path = DataLoadUtil.get_label_path(root_dir, scene_name) + output_label_path = DataLoadUtil.get_label_path(root_dir, scene_name,0) if os.path.exists(output_label_path) and not self.overwrite: Log.info(f"Scene <{scene_name}> Already Exists, Skip") cnt += 1 @@ -79,43 +81,52 @@ class StrategyGenerator(Runner): pts_list = [] for frame_idx in range(frame_num): - path = DataLoadUtil.get_path(root, scene_name, frame_idx) - cam_params = DataLoadUtil.load_cam_info(path, binocular=True) - status_manager.set_progress("generate_strategy", "strategy_generator", "loading frame", frame_idx, frame_num) - point_cloud = DataLoadUtil.get_target_point_cloud_world_from_path(path, binocular=True) - #display_table = None #DataLoadUtil.get_target_point_cloud_world_from_path(path, binocular=True, target_mask_label=()) #TODO - sampled_point_cloud = ReconstructionUtil.filter_points(point_cloud, model_points_normals, cam_pose=cam_params["cam_to_world"], voxel_size=voxel_threshold, theta=self.filter_degree) + if self.load_pts and os.path.exists(os.path.join(root,scene_name, "pts", f"{frame_idx}.txt")): + sampled_point_cloud = np.loadtxt(os.path.join(root,scene_name, "pts", f"{frame_idx}.txt")) + status_manager.set_progress("generate_strategy", "strategy_generator", "loading frame", frame_idx, frame_num) + pts_list.append(sampled_point_cloud) + continue + else: + path = DataLoadUtil.get_path(root, scene_name, frame_idx) + cam_params = DataLoadUtil.load_cam_info(path, binocular=True) + status_manager.set_progress("generate_strategy", "strategy_generator", "loading frame", frame_idx, frame_num) + point_cloud = DataLoadUtil.get_target_point_cloud_world_from_path(path, binocular=True) + sampled_point_cloud = ReconstructionUtil.filter_points(point_cloud, model_points_normals, cam_pose=cam_params["cam_to_world"], voxel_size=voxel_threshold, theta=self.filter_degree) - if self.save_pts: - pts_dir = os.path.join(root,scene_name, "pts") - if not os.path.exists(pts_dir): - os.makedirs(pts_dir) - np.savetxt(os.path.join(pts_dir, f"{frame_idx}.txt"), sampled_point_cloud) - pts_list.append(sampled_point_cloud) + if self.save_pts: + pts_dir = os.path.join(root,scene_name, "pts") + if not os.path.exists(pts_dir): + os.makedirs(pts_dir) + np.savetxt(os.path.join(pts_dir, f"{frame_idx}.txt"), sampled_point_cloud) + pts_list.append(sampled_point_cloud) status_manager.set_progress("generate_strategy", "strategy_generator", "loading frame", frame_num, frame_num) + + seq_num = min(self.seq_num, len(pts_list)) + init_view_list = range(seq_num) + + seq_idx = 0 + for init_view in init_view_list: + status_manager.set_progress("generate_strategy", "strategy_generator", "computing sequence", seq_idx, len(init_view_list)) + limited_useful_view, _, _ = ReconstructionUtil.compute_next_best_view_sequence_with_overlap(down_sampled_model_pts, pts_list,init_view=init_view, threshold=voxel_threshold, overlap_threshold=overlap_threshold, status_info=self.status_info) + data_pairs = self.generate_data_pairs(limited_useful_view) + seq_save_data = { + "data_pairs": data_pairs, + "best_sequence": limited_useful_view, + "max_coverage_rate": limited_useful_view[-1][1] + } - limited_useful_view, _, best_combined_pts = ReconstructionUtil.compute_next_best_view_sequence_with_overlap(down_sampled_model_pts, pts_list, threshold=voxel_threshold, overlap_threshold=overlap_threshold, status_info=self.status_info) - data_pairs = self.generate_data_pairs(limited_useful_view) - seq_save_data = { - "data_pairs": data_pairs, - "best_sequence": limited_useful_view, - "max_coverage_rate": limited_useful_view[-1][1] - } - - status_manager.set_status("generate_strategy", "strategy_generator", "max_coverage_rate", limited_useful_view[-1][1]) - Log.success(f"Scene <{scene_name}> Finished, Max Coverage Rate: {limited_useful_view[-1][1]}, Best Sequence length: {len(limited_useful_view)}") - - output_label_path = DataLoadUtil.get_label_path(root, scene_name) - output_best_reconstructed_pts_path = os.path.join(root,scene_name, f"best_reconstructed_pts.txt") - - with open(output_label_path, 'w') as f: - json.dump(seq_save_data, f) - - if self.save_best_combined_pts: - np.savetxt(output_best_reconstructed_pts_path, best_combined_pts) + status_manager.set_status("generate_strategy", "strategy_generator", "max_coverage_rate", limited_useful_view[-1][1]) + Log.success(f"Scene <{scene_name}> Finished, Max Coverage Rate: {limited_useful_view[-1][1]}, Best Sequence length: {len(limited_useful_view)}") + output_label_path = DataLoadUtil.get_label_path(root, scene_name, seq_idx) + + + with open(output_label_path, 'w') as f: + json.dump(seq_save_data, f) + seq_idx += 1 if self.save_mesh: DataLoadUtil.save_target_mesh_at_world_space(root, model_dir, scene_name) + status_manager.set_progress("generate_strategy", "strategy_generator", "computing sequence", len(init_view_list), len(init_view_list)) def generate_data_pairs(self, useful_view): diff --git a/utils/data_load.py b/utils/data_load.py index 06523c3..78ddc74 100644 --- a/utils/data_load.py +++ b/utils/data_load.py @@ -3,6 +3,7 @@ import numpy as np import json import cv2 import trimesh +import torch from utils.pts import PtsUtil class DataLoadUtil: @@ -13,8 +14,21 @@ class DataLoadUtil: return path @staticmethod - def get_label_path(root, scene_name): - path = os.path.join(root,scene_name, f"label.json") + def get_label_num(root, scene_name): + label_dir = os.path.join(root,scene_name,"label") + return len(os.listdir(label_dir)) + + @staticmethod + def get_label_path(root, scene_name, seq_idx): + label_dir = os.path.join(root,scene_name,"label") + if not os.path.exists(label_dir): + os.makedirs(label_dir) + path = os.path.join(label_dir,f"{seq_idx}.json") + return path + + @staticmethod + def get_label_path_old(root, scene_name): + path = os.path.join(root,scene_name,"label.json") return path @staticmethod @@ -45,11 +59,14 @@ class DataLoadUtil: mesh.export(model_path) @staticmethod - def save_target_mesh_at_world_space(root, model_dir, scene_name): + def save_target_mesh_at_world_space(root, model_dir, scene_name, display_table_as_world_space_origin=True): scene_info = DataLoadUtil.load_scene_info(root, scene_name) target_name = scene_info["target_name"] transformation = scene_info[target_name] - location = transformation["location"] + if display_table_as_world_space_origin: + location = transformation["location"] - DataLoadUtil.DISPLAY_TABLE_POSITION + else: + location = transformation["location"] rotation_euler = transformation["rotation_euler"] pose_mat = trimesh.transformations.euler_matrix(*rotation_euler) pose_mat[:3, 3] = location @@ -181,7 +198,9 @@ class DataLoadUtil: @staticmethod def get_real_cam_O_from_cam_L(cam_L, cam_O_to_cam_L, display_table_as_world_space_origin=True): - nO_to_display_table_pose = cam_L.cpu().numpy() @ cam_O_to_cam_L + if isinstance(cam_L, torch.Tensor): + cam_L = cam_L.cpu().numpy() + nO_to_display_table_pose = cam_L @ cam_O_to_cam_L if display_table_as_world_space_origin: display_table_to_world = np.eye(4) display_table_to_world[:3, 3] = DataLoadUtil.DISPLAY_TABLE_POSITION diff --git a/utils/reconstruction.py b/utils/reconstruction.py index 1bc3436..7b9d5cf 100644 --- a/utils/reconstruction.py +++ b/utils/reconstruction.py @@ -45,12 +45,17 @@ class ReconstructionUtil: @staticmethod - def compute_next_best_view_sequence_with_overlap(target_point_cloud, point_cloud_list, display_table_point_cloud_list = None,threshold=0.01, overlap_threshold=0.3, status_info=None): - selected_views = [] - current_coverage = 0.0 + def compute_next_best_view_sequence_with_overlap(target_point_cloud, point_cloud_list,threshold=0.01, overlap_threshold=0.3, init_view = 0, status_info=None): + selected_views = [point_cloud_list[init_view]] + combined_point_cloud = np.vstack(selected_views) + down_sampled_combined_point_cloud = PtsUtil.voxel_downsample_point_cloud(combined_point_cloud,threshold) + new_coverage = ReconstructionUtil.compute_coverage_rate(target_point_cloud, down_sampled_combined_point_cloud, threshold) + current_coverage = new_coverage remaining_views = list(range(len(point_cloud_list))) - view_sequence = [] + view_sequence = [(init_view, current_coverage)] cnt_processed_view = 0 + remaining_views.remove(init_view) + while remaining_views: best_view = None best_coverage_increase = -1 @@ -70,14 +75,13 @@ class ReconstructionUtil: down_sampled_combined_point_cloud = PtsUtil.voxel_downsample_point_cloud(combined_point_cloud,threshold) new_coverage = ReconstructionUtil.compute_coverage_rate(target_point_cloud, down_sampled_combined_point_cloud, threshold) coverage_increase = new_coverage - current_coverage - #print(f"view_index: {view_index}, coverage_increase: {coverage_increase}") if coverage_increase > best_coverage_increase: best_coverage_increase = coverage_increase best_view = view_index if best_view is not None: - if best_coverage_increase <=1e-3: + if best_coverage_increase <=3e-3: break selected_views.append(point_cloud_list[best_view]) remaining_views.remove(best_view) diff --git a/utils/render.py b/utils/render.py index 46ea97d..2282542 100644 --- a/utils/render.py +++ b/utils/render.py @@ -12,8 +12,8 @@ class RenderUtil: def render_pts(cam_pose, scene_path,script_path, model_points_normals, 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) - - + + with tempfile.TemporaryDirectory() as temp_dir: params = { "cam_pose": nO_to_world_pose.tolist(), @@ -30,7 +30,6 @@ class RenderUtil: print(result.stderr) return None path = os.path.join(temp_dir, "tmp") - point_cloud = DataLoadUtil.get_target_point_cloud_world_from_path(path, binocular=True) cam_params = DataLoadUtil.load_cam_info(path, binocular=True) filtered_point_cloud = ReconstructionUtil.filter_points(point_cloud, model_points_normals, cam_pose=cam_params["cam_to_world"], voxel_size=voxel_threshold, theta=filter_degree) @@ -44,4 +43,5 @@ class RenderUtil: point_cloud_R = PtsUtil.random_downsample_point_cloud(point_cloud_R, 65536) full_scene_point_cloud = DataLoadUtil.get_overlapping_points(point_cloud_L, point_cloud_R) + return filtered_point_cloud, full_scene_point_cloud \ No newline at end of file