This commit is contained in:
2024-10-10 21:48:55 +08:00
parent 8fd2d6b1e1
commit cd85fed3a0
10 changed files with 73 additions and 33004 deletions

View File

@@ -14,7 +14,7 @@ from utils.control_util import ControlUtil
from utils.communicate_util import CommunicateUtil
from utils.pts_util import PtsUtil
from utils.reconstruction_util import ReconstructionUtil
from utils.preprocess_util import save_scene_data
from utils.preprocess_util import save_scene_data, save_scene_data_multithread
from utils.data_load import DataLoadUtil
from utils.view_util import ViewUtil
@@ -59,6 +59,8 @@ class CADStrategyRunner(Runner):
return scan_pts, obj_pts
def run_one_model(self, model_name):
result = dict()
shot_pts_list = []
ControlUtil.connect_robot()
''' init robot '''
@@ -72,30 +74,26 @@ class CADStrategyRunner(Runner):
Log.info("[Part 1/5] take first view data")
view_data = CommunicateUtil.get_view_data(init=True)
first_cam_pts = ViewUtil.get_pts(view_data)
#shot_pts_list.append(first_cam_pts)
first_cam_to_real_world = ControlUtil.get_pose()
''' register '''
Log.info("[Part 1/5] do registeration")
cam_to_cad = PtsUtil.register(first_cam_pts, cad_model)
first_cam_to_world = ControlUtil.get_pose()
cad_to_world = first_cam_to_world @ np.linalg.inv(cam_to_cad)
cad_to_real_world = first_cam_to_real_world @ np.linalg.inv(cam_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)
world_to_blender_world = np.eye(4)
world_to_blender_world[:3, 3] = np.asarray([0, 0, 0.9215])
cad_to_blender_world = world_to_blender_world @ cad_to_world
temp_dir = "/home/yan20/nbv_rec/project/franka_control/temp_output"
cad_model:trimesh.Trimesh = cad_model.apply_transform(cad_to_blender_world)
cad_model.export(os.path.join(temp_dir, f"{temp_name}.obj"))
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"))
cad_model.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")
result = subprocess.run([
subprocess.run([
self.blender_bin_path, '-b', '-P', self.generator_script_path, '--', temp_dir
], capture_output=True, text=True)
Log.success("[Part 2/5] finish running renderer")
@@ -104,7 +102,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(temp_dir, temp_name)
save_scene_data_multithread(temp_dir, temp_name)
Log.success("[Part 3/5] finish preprocessing data")
pts_dir = os.path.join(temp_dir,temp_name,"pts")
@@ -114,11 +112,11 @@ class CADStrategyRunner(Runner):
for frame_idx in range(frame_num):
pts_path = os.path.join(temp_dir,temp_name, "pts", f"{frame_idx}.txt")
idx_path = os.path.join(temp_dir,temp_name, "scan_points_indices", f"{frame_idx}.txt")
idx_path = os.path.join(temp_dir,temp_name, "scan_points_indices", f"{frame_idx}.npy")
point_cloud = np.loadtxt(pts_path)
if point_cloud.shape[0] != 0:
sampled_point_cloud = PtsUtil.voxel_downsample_point_cloud(point_cloud, self.voxel_size)
indices = np.loadtxt(idx_path, dtype=np.int32)
indices = np.load(idx_path)
try:
len(indices)
except:
@@ -149,7 +147,7 @@ class CADStrategyRunner(Runner):
for idx, coverage_rate in limited_useful_view:
path = DataLoadUtil.get_path(temp_dir, temp_name, idx)
cam_info = DataLoadUtil.load_cam_info(path, binocular=True)
cam_to_world_seq.append(cam_info["cam_to_world"])
cam_to_world_seq.append(cam_info["cam_to_world_O"])
coveraget_rate_seq.append(coverage_rate)
idx_seq.append(idx)
render_pts.append(sample_view_pts_list[idx])
@@ -157,32 +155,55 @@ 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)
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)
step = 1
result["real_coverage_rate_seq"] = []
for cam_to_world in cam_to_world_seq:
ControlUtil.move_to(cam_to_world)
''' get world pts '''
time.sleep(1)
view_data = CommunicateUtil.get_view_data()
if view_data is None:
Log.error("Failed to get view data")
continue
cam_pts = ViewUtil.get_pts(view_data)
shot_pts_list.append(cam_pts)
try:
ControlUtil.move_to(cam_to_world)
''' get world pts '''
time.sleep(1)
view_data = CommunicateUtil.get_view_data()
if view_data is None:
Log.error("Failed to get view data")
continue
cam_pts = ViewUtil.get_pts(view_data)
shot_pts_list.append(cam_pts)
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)
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)
step += 1
except Exception as e:
Log.error(f"Failed to move to {cam_to_world}")
Log.error(e)
#import ipdb;ipdb.set_trace()
print(idx_seq)
for idx in range(len(shot_pts_list)):
if not os.path.exists(os.path.join(temp_dir, temp_name, "shot_pts")):
os.makedirs(os.path.join(temp_dir, temp_name, "shot_pts"))
if not os.path.exists(os.path.join(temp_dir, temp_name, "render_pts")):
os.makedirs(os.path.join(temp_dir, temp_name, "render_pts"))
shot_pts = PtsUtil.transform_point_cloud(shot_pts_list[idx], first_cam_to_world)
shot_pts = PtsUtil.transform_point_cloud(shot_pts_list[idx], first_cam_to_real_world)
np.savetxt(os.path.join(temp_dir, temp_name, "shot_pts", f"{idx}.txt"), shot_pts)
np.savetxt(os.path.join(temp_dir, temp_name, "render_pts", f"{idx}.txt"), render_pts[idx])
Log.success("[Part 5/5] finish running robot")
Log.debug(result)
def run(self):
total = len(os.listdir(self.model_dir))