243 lines
11 KiB
Python
243 lines
11 KiB
Python
|
import os
|
||
|
|
import time
|
||
|
|
import trimesh
|
||
|
|
import tempfile
|
||
|
|
import subprocess
|
||
|
|
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.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_close_loop_strategy_runner")
|
||
|
|
class CADCloseLoopStrategyRunner(Runner):
|
||
|
|
|
||
|
|
def __init__(self, config_path: str):
|
||
|
|
super().__init__(config_path)
|
||
|
|
self.load_experiment("cad_strategy")
|
||
|
|
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_shot_view_num = self.reconstruct_config["max_shot_view_num"]
|
||
|
|
self.min_shot_new_pts_num = self.reconstruct_config["min_shot_new_pts_num"]
|
||
|
|
self.min_coverage_increase = self.reconstruct_config["min_coverage_increase"]
|
||
|
|
self.scan_points_threshold = self.reconstruct_config["scan_points_threshold"]
|
||
|
|
|
||
|
|
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 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):
|
||
|
|
time_code = time.strftime("%Y%m%d%H%M%S", time.localtime())
|
||
|
|
temp_dir = "/home/yan20/nbv_rec/project/franka_control/temp_output_" + time_code
|
||
|
|
if not os.path.exists(temp_dir):
|
||
|
|
os.makedirs(temp_dir)
|
||
|
|
ControlUtil.connect_robot()
|
||
|
|
""" 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.obj")
|
||
|
|
temp_name = "cad_model_world"
|
||
|
|
cad_model = trimesh.load(model_path)
|
||
|
|
""" take first view """
|
||
|
|
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)
|
||
|
|
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
|
||
|
|
)
|
||
|
|
|
||
|
|
debug_dir = os.path.join(temp_dir, "debug")
|
||
|
|
if not os.path.exists(debug_dir):
|
||
|
|
os.makedirs(debug_dir)
|
||
|
|
|
||
|
|
np.savetxt(os.path.join(debug_dir,"FIRST_shot.txt"), first_splitted_real_world_pts)
|
||
|
|
# np.savetxt(f"first_real_pts_{model_name}.txt", first_splitted_real_world_pts)
|
||
|
|
""" register """
|
||
|
|
Log.info("[Part 1/4] do registeration")
|
||
|
|
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/4] 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_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_" + time_code
|
||
|
|
cad_model_blender_world.export(os.path.join(temp_dir, f"{temp_name}.obj"))
|
||
|
|
""" sample view """
|
||
|
|
Log.info("[Part 2/4] start running renderer")
|
||
|
|
subprocess.run(
|
||
|
|
[
|
||
|
|
self.blender_bin_path,
|
||
|
|
"-b",
|
||
|
|
"-P",
|
||
|
|
self.generator_script_path,
|
||
|
|
"--",
|
||
|
|
temp_dir,
|
||
|
|
],
|
||
|
|
capture_output=True,
|
||
|
|
text=True,
|
||
|
|
)
|
||
|
|
Log.success("[Part 2/4] finish running renderer")
|
||
|
|
|
||
|
|
""" preprocess """
|
||
|
|
Log.info("[Part 3/4] start preprocessing data")
|
||
|
|
save_scene_data(temp_dir, temp_name)
|
||
|
|
Log.success("[Part 3/4] finish preprocessing data")
|
||
|
|
|
||
|
|
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, temp_name, "pts", 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.load(idx_path)
|
||
|
|
try:
|
||
|
|
len(indices)
|
||
|
|
except:
|
||
|
|
indices = np.array([indices])
|
||
|
|
sample_view_pts_list.append(sampled_point_cloud)
|
||
|
|
scan_points_idx_list.append(indices)
|
||
|
|
|
||
|
|
""" close-loop strategy """
|
||
|
|
scanned_pts = PtsUtil.voxel_downsample_point_cloud(
|
||
|
|
first_splitted_real_world_pts, self.voxel_size
|
||
|
|
)
|
||
|
|
shot_pts_list = [first_splitted_real_world_pts]
|
||
|
|
history_indices = []
|
||
|
|
last_coverage = 0
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
Log.info("[Part 4/4] start close-loop control")
|
||
|
|
cnt = 0
|
||
|
|
selected_view = []
|
||
|
|
while True:
|
||
|
|
#import ipdb; ipdb.set_trace()
|
||
|
|
next_best_view, next_best_coverage, next_best_covered_num = (
|
||
|
|
ReconstructionUtil.compute_next_best_view_with_overlap(
|
||
|
|
scanned_pts,
|
||
|
|
sample_view_pts_list,
|
||
|
|
selected_view,
|
||
|
|
history_indices,
|
||
|
|
scan_points_idx_list,
|
||
|
|
threshold=self.voxel_size,
|
||
|
|
overlap_area_threshold=50,
|
||
|
|
scan_points_threshold=self.scan_points_threshold,
|
||
|
|
)
|
||
|
|
)
|
||
|
|
selected_view.append(next_best_view)
|
||
|
|
nbv_path = DataLoadUtil.get_path(temp_dir, temp_name, next_best_view)
|
||
|
|
nbv_cam_info = DataLoadUtil.load_cam_info(nbv_path, binocular=True)
|
||
|
|
nbv_cam_to_world = nbv_cam_info["cam_to_world_O"]
|
||
|
|
ControlUtil.move_to(nbv_cam_to_world)
|
||
|
|
''' get world pts '''
|
||
|
|
time.sleep(0.5)
|
||
|
|
view_data = CommunicateUtil.get_view_data()
|
||
|
|
if view_data is None:
|
||
|
|
Log.error("No view data received")
|
||
|
|
continue
|
||
|
|
cam_shot_pts = ViewUtil.get_pts(view_data)
|
||
|
|
world_shot_pts = PtsUtil.transform_point_cloud(
|
||
|
|
cam_shot_pts, first_cam_to_real_world
|
||
|
|
)
|
||
|
|
_, world_splitted_shot_pts = self.split_scan_pts_and_obj_pts(
|
||
|
|
world_shot_pts
|
||
|
|
)
|
||
|
|
shot_pts_list.append(world_splitted_shot_pts)
|
||
|
|
|
||
|
|
|
||
|
|
np.savetxt(os.path.join(debug_dir, f"shot_pts_{cnt}.txt"), world_splitted_shot_pts)
|
||
|
|
np.savetxt(os.path.join(debug_dir, f"render_pts_{cnt}.txt"), sample_view_pts_list[next_best_view])
|
||
|
|
#real_world_to_cad = PtsUtil.register(first_splitted_real_world_pts, cad_model)
|
||
|
|
#import ipdb; ipdb.set_trace()
|
||
|
|
last_scanned_pts_num = scanned_pts.shape[0]
|
||
|
|
new_scanned_pts = PtsUtil.voxel_downsample_point_cloud(
|
||
|
|
np.vstack([scanned_pts, world_splitted_shot_pts]), self.voxel_size
|
||
|
|
)
|
||
|
|
new_scanned_pts_num = new_scanned_pts.shape[0]
|
||
|
|
history_indices.append(scan_points_idx_list[next_best_view])
|
||
|
|
scanned_pts = new_scanned_pts
|
||
|
|
Log.info(
|
||
|
|
f"Next Best cover pts: {next_best_covered_num}, Best coverage: {next_best_coverage}"
|
||
|
|
)
|
||
|
|
|
||
|
|
coverage_rate_increase = next_best_coverage - last_coverage
|
||
|
|
if coverage_rate_increase < self.min_coverage_increase:
|
||
|
|
Log.info(f"Coverage rate = {coverage_rate_increase} < {self.min_coverage_increase}, stop scanning")
|
||
|
|
# break
|
||
|
|
last_coverage = next_best_coverage
|
||
|
|
|
||
|
|
new_added_pts_num = new_scanned_pts_num - last_scanned_pts_num
|
||
|
|
if new_added_pts_num < self.min_shot_new_pts_num:
|
||
|
|
Log.info(f"New added pts num = {new_added_pts_num} < {self.min_shot_new_pts_num}")
|
||
|
|
#ipdb.set_trace()
|
||
|
|
if len(shot_pts_list) >= self.max_shot_view_num:
|
||
|
|
Log.info(f"Scanned view num = {len(shot_pts_list)} >= {self.max_shot_view_num}, stop scanning")
|
||
|
|
#break
|
||
|
|
cnt += 1
|
||
|
|
|
||
|
|
Log.success("[Part 4/4] finish close-loop control")
|
||
|
|
|
||
|
|
|
||
|
|
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}")
|
||
|
|
self.run_one_model("box_1")
|
||
|
|
|
||
|
|
# ---------------------------- test ---------------------------- #
|
||
|
|
if __name__ == "__main__":
|
||
|
|
|
||
|
|
model_path = r"C:\Users\hofee\Downloads\mesh.obj"
|
||
|
|
model = trimesh.load(model_path)
|