hofee/nbv_rec_control
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add renderer

Browse Source
This commit is contained in:
hofee 2024-10-08 16:16:54 +08:00
parent d9d2716ba7
commit 3f1ae95ff8
4 changed files with 154 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

23
configs/render_config.yaml Normal file
View File

@ -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

4
runners/cad_strategy.py
View File

@ -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,

127
runners/renderer.py Normal file
View File

@ -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)

2
utils/control_util.py
View File

@ -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],