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

update

Browse Source
This commit is contained in:
hofee
2024-10-12 16:39:00 +08:00
parent cd85fed3a0
commit 3fe74eb6eb
5 changed files with 226 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files

37
runners/cad_strategy.py
View File

@@ -53,12 +53,13 @@ class CADStrategyRunner(Runner):
def load_experiment(self, backup_name=None):
super().load_experiment(backup_name)
def split_scan_pts_and_obj_pts(self, world_pts, scan_pts_z, z_threshold = 0.003):
scan_pts = world_pts[scan_pts_z < z_threshold]
obj_pts = world_pts[scan_pts_z >= z_threshold]
def split_scan_pts_and_obj_pts(self, world_pts, z_threshold = 0):
scan_pts = world_pts[world_pts[:,2] < z_threshold]
obj_pts = world_pts[world_pts[:,2] >= z_threshold]
return scan_pts, obj_pts
def run_one_model(self, model_name):
temp_dir = "/home/yan20/nbv_rec/project/franka_control/temp_output"
result = dict()
shot_pts_list = []
@@ -66,6 +67,7 @@ class CADStrategyRunner(Runner):
''' init robot '''
Log.info("[Part 1/5] start init and register")
ControlUtil.init()
''' load CAD model '''
model_path = os.path.join(self.model_dir, model_name,"mesh.ply")
temp_name = "cad_model_world"
@@ -75,21 +77,25 @@ class CADStrategyRunner(Runner):
view_data = CommunicateUtil.get_view_data(init=True)
first_cam_pts = ViewUtil.get_pts(view_data)
first_cam_to_real_world = ControlUtil.get_pose()
first_real_world_pts = PtsUtil.transform_point_cloud(first_cam_pts, first_cam_to_real_world)
_, first_splitted_real_world_pts = self.split_scan_pts_and_obj_pts(first_real_world_pts)
np.savetxt(f"first_real_pts_{model_name}.txt", first_splitted_real_world_pts)
''' register '''
Log.info("[Part 1/5] do registeration")
cam_to_cad = PtsUtil.register(first_cam_pts, cad_model)
cad_to_real_world = first_cam_to_real_world @ np.linalg.inv(cam_to_cad)
real_world_to_cad = PtsUtil.register(first_splitted_real_world_pts, cad_model)
cad_to_real_world = np.linalg.inv(real_world_to_cad)
Log.success("[Part 1/5] finish init and register")
real_world_to_blender_world = np.eye(4)
real_world_to_blender_world[:3, 3] = np.asarray([0, 0, 0.9215])
cad_to_blender_world = real_world_to_blender_world @ cad_to_real_world
cad_model_blender_world:trimesh.Trimesh = cad_model.apply_transform(cad_to_blender_world)
cad_model_real_world:trimesh.Trimesh = cad_model.apply_transform(cad_to_real_world)
cad_model_real_world.export(os.path.join(temp_dir, f"real_world_{temp_name}.obj"))
cad_model_blender_world:trimesh.Trimesh = cad_model.apply_transform(real_world_to_blender_world)
with tempfile.TemporaryDirectory() as temp_dir:
temp_dir = "/home/yan20/nbv_rec/project/franka_control/temp_output"
cad_model_blender_world.export(os.path.join(temp_dir, f"{temp_name}.obj"))
scene_dir = os.path.join(temp_dir, temp_name)
''' sample view '''
Log.info("[Part 2/5] start running renderer")
@@ -102,7 +108,7 @@ class CADStrategyRunner(Runner):
world_model_points = np.loadtxt(os.path.join(scene_dir, "points_and_normals.txt"))[:,:3]
''' preprocess '''
Log.info("[Part 3/5] start preprocessing data")
save_scene_data_multithread(temp_dir, temp_name)
save_scene_data(temp_dir, temp_name)
Log.success("[Part 3/5] finish preprocessing data")
pts_dir = os.path.join(temp_dir,temp_name,"pts")
@@ -154,16 +160,13 @@ class CADStrategyRunner(Runner):
Log.info("[Part 5/5] start running robot")
''' take best seq view '''
cad_model_real_world = cad_model_blender_world.apply_transform(np.linalg.inv(real_world_to_blender_world))
cad_model_real_world_pts = cad_model_real_world.vertices
cad_model_real_world.export(os.path.join(temp_dir, f"{temp_name}_real_world.obj"))
voxel_downsampled_cad_model_real_world_pts = PtsUtil.voxel_downsample_point_cloud(cad_model_real_world_pts, self.voxel_size)
#import ipdb; ipdb.set_trace()
target_scanned_pts = np.concatenate(sample_view_pts_list)
voxel_downsampled_target_scanned_pts = PtsUtil.voxel_downsample_point_cloud(target_scanned_pts, self.voxel_size)
result = dict()
gt_scanned_pts = np.concatenate(render_pts, axis=0)
voxel_down_sampled_gt_scanned_pts = PtsUtil.voxel_downsample_point_cloud(gt_scanned_pts, self.voxel_size)
result["gt_final_coverage_rate_cad"] = ReconstructionUtil.compute_coverage_rate(voxel_downsampled_cad_model_real_world_pts, voxel_down_sampled_gt_scanned_pts, self.voxel_size)
result["gt_final_coverage_rate_cad"] = ReconstructionUtil.compute_coverage_rate(voxel_downsampled_target_scanned_pts, voxel_down_sampled_gt_scanned_pts, self.voxel_size)
step = 1
result["real_coverage_rate_seq"] = []
for cam_to_world in cam_to_world_seq:
@@ -180,7 +183,7 @@ class CADStrategyRunner(Runner):
scanned_pts = np.concatenate(shot_pts_list, axis=0)
voxel_down_sampled_scanned_pts = PtsUtil.voxel_downsample_point_cloud(scanned_pts, self.voxel_size)
voxel_down_sampled_scanned_pts_world = PtsUtil.transform_point_cloud(voxel_down_sampled_scanned_pts, first_cam_to_real_world)
curr_CR = ReconstructionUtil.compute_coverage_rate(voxel_downsampled_cad_model_real_world_pts, voxel_down_sampled_scanned_pts_world, self.voxel_size)
curr_CR = ReconstructionUtil.compute_coverage_rate(voxel_downsampled_target_scanned_pts, voxel_down_sampled_scanned_pts_world, self.voxel_size)
Log.success(f"(step {step}/{len(cam_to_world_seq)}) current coverage: {curr_CR} | gt coverage: {result['gt_final_coverage_rate_cad']}")
result["real_final_coverage_rate"] = curr_CR
result["real_coverage_rate_seq"].append(curr_CR)