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

finish pipeline

Browse Source
This commit is contained in:
hofee
2024-10-09 21:46:13 +08:00
parent f8514564c1
commit ba36803fba
8 changed files with 325 additions and 138 deletions
Download Patch File Download Diff File Expand all files Collapse all files

125
runners/cad_strategy.py
View File

@@ -1,4 +1,5 @@
import os
import time
import trimesh
import tempfile
import subprocess
@@ -15,6 +16,8 @@ from utils.pts_util import PtsUtil
from utils.reconstruction_util import ReconstructionUtil
from utils.preprocess_util import save_scene_data
from utils.data_load import DataLoadUtil
from utils.view_util import ViewUtil
@stereotype.runner("CAD_strategy_runner")
class CADStrategyRunner(Runner):
@@ -29,6 +32,8 @@ class CADStrategyRunner(Runner):
}
self.generate_config = ConfigManager.get("runner", "generate")
self.reconstruct_config = ConfigManager.get("runner", "reconstruct")
self.blender_bin_path = self.generate_config["blender_bin_path"]
self.generator_script_path = self.generate_config["generator_script_path"]
self.model_dir = self.generate_config["model_dir"]
self.voxel_size = self.generate_config["voxel_size"]
self.max_view = self.generate_config["max_view"]
@@ -47,13 +52,6 @@ class CADStrategyRunner(Runner):
def load_experiment(self, backup_name=None):
super().load_experiment(backup_name)
def get_pts_from_view_data(self, view_data):
depth = view_data["depth_image"]
depth_intrinsics = view_data["depth_intrinsics"]
depth_extrinsics = view_data["depth_extrinsics"]
cam_pts = PtsUtil.get_pts_from_depth(depth, depth_intrinsics, depth_extrinsics)
return cam_pts
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]
@@ -61,56 +59,65 @@ class CADStrategyRunner(Runner):
return scan_pts, obj_pts
def run_one_model(self, model_name):
shot_pts_list = []
ControlUtil.connect_robot()
''' init robot '''
#ControlUtil.init()
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.obj")
model_path = os.path.join(self.model_dir, model_name,"mesh.ply")
temp_name = "cad_model_world"
cad_model = trimesh.load(model_path)
''' take first view '''
#view_data = CommunicateUtil.get_view_data(init=True)
#first_cam_pts = self.get_pts_from_view_data(view_data)
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)
''' register '''
#cad_to_cam = PtsUtil.register(first_cam_pts, cad_model)
#cam_to_world = ControlUtil.get_pose()
cad_to_world = np.eye(4) #cam_to_world @ cad_to_cam
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)
Log.success("[Part 1/5] finish init and register")
world_to_blender_world = np.eye(4)
world_to_blender_world[:3, 3] = np.asarray([0, 0, 0.9215])
cad_to_blender_world = np.dot(world_to_blender_world, cad_to_world)
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:
name = "cad_model_world"
temp_dir = "/home/yan20/nbv_rec/project/franka_control/temp_output"
cad_model.export(os.path.join(temp_dir, f"{name}.obj"))
temp_dir = "/home/user/nbv_rec/nbv_rec_control/test_output"
scene_dir = os.path.join(temp_dir, name)
script_path = "/home/user/nbv_rec/blender_app/data_generator.py"
cad_model.export(os.path.join(temp_dir, f"{temp_name}.obj"))
scene_dir = os.path.join(temp_dir, temp_name)
''' sample view '''
# import ipdb; ipdb.set_trace()
# print("start running renderer")
# result = subprocess.run([
# 'blender', '-b', '-P', script_path, '--', temp_dir
# ], capture_output=True, text=True)
# print("finish running renderer")
#
Log.info("[Part 2/5] start running renderer")
result = 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")
world_model_points = np.loadtxt(os.path.join(scene_dir, "points_and_normals.txt"))[:,:3]
''' preprocess '''
# save_scene_data(temp_dir, name)
Log.info("[Part 3/5] start preprocessing data")
save_scene_data(temp_dir, temp_name)
Log.success("[Part 3/5] finish preprocessing data")
pts_dir = os.path.join(temp_dir,name,"pts")
pts_dir = os.path.join(temp_dir,temp_name,"pts")
sample_view_pts_list = []
scan_points_idx_list = []
frame_num = len(os.listdir(pts_dir))
for frame_idx in range(frame_num):
pts_path = os.path.join(temp_dir,name, "pts", f"{frame_idx}.txt")
idx_path = os.path.join(temp_dir,name, "scan_points_indices", f"{frame_idx}.txt")
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")
point_cloud = np.loadtxt(pts_path)
sampled_point_cloud = PtsUtil.voxel_downsample_point_cloud(point_cloud, self.voxel_size)
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)
try:
len(indices)
@@ -121,6 +128,7 @@ class CADStrategyRunner(Runner):
''' generate strategy '''
Log.info("[Part 4/5] start generating strategy")
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,
@@ -131,28 +139,49 @@ class CADStrategyRunner(Runner):
scan_points_threshold = self.scan_points_threshold,
status_info=self.status_info
)
Log.success("[Part 4/5] finish generating strategy")
''' extract cam_to_world sequence '''
cam_to_world_seq = []
coveraget_rate_seq = []
from ipdb import set_trace; set_trace()
render_pts = []
idx_seq = []
for idx, coverage_rate in limited_useful_view:
path = DataLoadUtil.get_path(temp_dir, name, idx)
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"])
coveraget_rate_seq.append(coverage_rate)
idx_seq.append(idx)
render_pts.append(sample_view_pts_list[idx])
# ''' 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 = self.get_pts_from_view_data(view_data)
# 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)
Log.info("[Part 5/5] start running robot")
''' take best seq view '''
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)
#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)
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")
def run(self):