diff --git a/configs/render_config.yaml b/configs/render_config.yaml new file mode 100644 index 0000000..bac4f56 --- /dev/null +++ b/configs/render_config.yaml @@ -0,0 +1,23 @@ + +runner: + general: + seed: 1 + device: cpu + cuda_visible_devices: "0,1,2,3,4,5,6,7" + + experiment: + name: debug + root_dir: "experiments" + + web: + host: “0.0.0.0” + port: 11111 + render: + model_dir: "/home/yan20/nbv_rec/data/test_CAD/test_model" + + + reconstruct: + soft_overlap_threshold: 0.3 + hard_overlap_threshold: 0.6 + scan_points_threshold: 10 + \ No newline at end of file diff --git a/runners/cad_strategy.py b/runners/cad_strategy.py index c87f17e..a59bb9b 100644 --- a/runners/cad_strategy.py +++ b/runners/cad_strategy.py @@ -61,7 +61,7 @@ class CADStrategyRunner(Runner): view_data = CommunicateUtil.get_view_data() first_cam_pts = None ''' register ''' - cad_to_cam = PtsUtil.register_icp(first_cam_pts, cad_model) + cad_to_cam = PtsUtil.register(first_cam_pts, cad_model) cam_to_world = ControlUtil.get_pose() cad_to_world = cam_to_world @ cad_to_cam cad_model:trimesh.Trimesh = cad_model.apply_transform(cad_to_world) @@ -94,6 +94,8 @@ class CADStrategyRunner(Runner): scan_points_idx_list.append(scan_points_idx) ''' generate strategy ''' + scan_points = ReconstructionUtil.generate_scan_points(display_table_top=0, display_table_radius=0.25) + limited_useful_view, _, _ = ReconstructionUtil.compute_next_best_view_sequence_with_overlap( world_model_points, sample_view_pts_list, scan_points_indices_list = scan_points_idx_list, diff --git a/runners/renderer.py b/runners/renderer.py new file mode 100644 index 0000000..00cc3ed --- /dev/null +++ b/runners/renderer.py @@ -0,0 +1,127 @@ +import os +import trimesh +import numpy as np +from PytorchBoot.runners.runner import Runner +from PytorchBoot.config import ConfigManager +import PytorchBoot.stereotype as stereotype +from PytorchBoot.utils.log_util import Log +from PytorchBoot.status import status_manager + +#from utils.control_util import ControlUtil +from utils.communicate_util import CommunicateUtil +from utils.pts_util import PtsUtil +from utils.view_sample_util import ViewSampleUtil +from utils.reconstruction_util import ReconstructionUtil + + +@stereotype.runner("online_renderer") +class OnlineRenderer(Runner): + + def __init__(self, config_path: str): + super().__init__(config_path) + self.load_experiment("cad_strategy") + self.host = ConfigManager.get("runner", "web", "host") + self.port = ConfigManager.get("runner", "web", "port") + + def create_experiment(self, backup_name=None): + super().create_experiment(backup_name) + + def load_experiment(self, backup_name=None): + super().load_experiment(backup_name) + + def run_one_model(self, model_name): + + ''' init robot ''' + ControlUtil.init() + ''' load CAD model ''' + model_path = os.path.join(self.model_dir, model_name) + cad_model = trimesh.load(model_path) + ''' take first view ''' + view_data = CommunicateUtil.get_view_data() + first_cam_pts = None + ''' register ''' + cad_to_cam = PtsUtil.register(first_cam_pts, cad_model) + cam_to_world = ControlUtil.get_pose() + cad_to_world = cam_to_world @ cad_to_cam + cad_model:trimesh.Trimesh = cad_model.apply_transform(cad_to_world) + ''' sample view ''' + min_corner = cad_model.bounds[0] + max_corner = cad_model.bounds[1] + diag = np.linalg.norm(max_corner - min_corner) + view_num = int(self.min_view + (diag - self.min_diag)/(self.max_diag - self.min_diag) * (self.max_view - self.min_view)) + sampled_view_data = ViewSampleUtil.sample_view_data_world_space( + cad_model, cad_to_world, + voxel_size = self.voxel_size, + max_views = view_num, + min_cam_table_included_degree= self.min_cam_table_included_degree, + random_view_ratio = self.random_view_ratio + ) + cam_to_world_poses = sampled_view_data["cam_to_world_poses"] + world_model_points = sampled_view_data["voxel_down_sampled_points"] + + ''' take sample view ''' + scan_points_idx_list = [] + sample_view_pts_list = [] + for cam_to_world in cam_to_world_poses: + ControlUtil.move_to(cam_to_world) + ''' get world pts ''' + view_data = CommunicateUtil.get_view_data() + cam_pts = None + scan_points_idx = None + world_pts = PtsUtil.transform_point_cloud(cam_pts, cam_to_world) + sample_view_pts_list.append(world_pts) + scan_points_idx_list.append(scan_points_idx) + + ''' generate strategy ''' + scan_points = ReconstructionUtil.generate_scan_points(display_table_top=0, display_table_radius=0.25) + + limited_useful_view, _, _ = ReconstructionUtil.compute_next_best_view_sequence_with_overlap( + world_model_points, sample_view_pts_list, + scan_points_indices_list = scan_points_idx_list, + init_view=0, + threshold=self.voxel_size, + soft_overlap_threshold = self.soft_overlap_threshold, + hard_overlap_threshold = self.hard_overlap_threshold, + scan_points_threshold = self.scan_points_threshold, + status_info=self.status_info + ) + + ''' extract cam_to_world sequence ''' + cam_to_world_seq = [] + coveraget_rate_seq = [] + + for idx, coverage_rate in limited_useful_view: + cam_to_world_seq.append(cam_to_world_poses[idx]) + coveraget_rate_seq.append(coverage_rate) + + ''' take best seq view ''' + for cam_to_world in cam_to_world_seq: + ControlUtil.move_to(cam_to_world) + ''' get world pts ''' + view_data = CommunicateUtil.get_view_data() + cam_pts = None + scan_points_idx = None + world_pts = PtsUtil.transform_point_cloud(cam_pts, cam_to_world) + sample_view_pts_list.append(world_pts) + scan_points_idx_list.append(scan_points_idx) + + + def run(self): + total = len(os.listdir(self.model_dir)) + model_start_idx = self.generate_config["model_start_idx"] + count_object = model_start_idx + for model_name in os.listdir(self.model_dir[model_start_idx:]): + Log.info(f"[{count_object}/{total}]Processing {model_name}") + self.run_one_model(model_name) + Log.success(f"[{count_object}/{total}]Finished processing {model_name}") + + +# ---------------------------- test ---------------------------- # +if __name__ == "__main__": + + model_path = r"C:\Users\hofee\Downloads\mesh.obj" + model = trimesh.load(model_path) + + + + \ No newline at end of file diff --git a/utils/control_util.py b/utils/control_util.py index 3af6b20..ee46693 100644 --- a/utils/control_util.py +++ b/utils/control_util.py @@ -4,7 +4,7 @@ from autolab_core import RigidTransform class ControlUtil: - #__fa = FrankaArm(robot_num=2) + __fa = FrankaArm(robot_num=2) BASE_TO_WORLD:np.ndarray = np.asarray([ [1, 0, 0, -0.5],