update
This commit is contained in:
hofee
@@ -1,5 +1,7 @@
|
||||
import os
|
||||
import trimesh
|
||||
import tempfile
|
||||
import subprocess
|
||||
import numpy as np
|
||||
from PytorchBoot.runners.runner import Runner
|
||||
from PytorchBoot.config import ConfigManager
|
||||
@@ -7,16 +9,12 @@ import PytorchBoot.stereotype as stereotype
|
||||
from PytorchBoot.utils.log_util import Log
|
||||
from PytorchBoot.status import status_manager
|
||||
|
||||
|
||||
import sys
|
||||
|
||||
sys.path.append(r"C:\Document\Local Project\nbv_rec_control")
|
||||
#from utils.control_util import ControlUtil
|
||||
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
|
||||
|
||||
from utils.preprocess_util import save_scene_data
|
||||
from utils.data_load import DataLoadUtil
|
||||
|
||||
@stereotype.runner("CAD_strategy_runner")
|
||||
class CADStrategyRunner(Runner):
|
||||
@@ -50,79 +48,111 @@ 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]
|
||||
obj_pts = world_pts[scan_pts_z >= z_threshold]
|
||||
return scan_pts, obj_pts
|
||||
|
||||
def run_one_model(self, model_name):
|
||||
|
||||
''' init robot '''
|
||||
ControlUtil.init()
|
||||
#ControlUtil.init()
|
||||
''' load CAD model '''
|
||||
model_path = os.path.join(self.model_dir, model_name)
|
||||
model_path = os.path.join(self.model_dir, model_name,"mesh.obj")
|
||||
cad_model = trimesh.load(model_path)
|
||||
''' take first view '''
|
||||
view_data = CommunicateUtil.get_view_data()
|
||||
first_cam_pts = None
|
||||
#view_data = CommunicateUtil.get_view_data(init=True)
|
||||
#first_cam_pts = self.get_pts_from_view_data(view_data)
|
||||
''' register '''
|
||||
cad_to_cam = PtsUtil.register_icp(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"]
|
||||
#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
|
||||
|
||||
|
||||
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_model:trimesh.Trimesh = cad_model.apply_transform(cad_to_blender_world)
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
name = "cad_model_world"
|
||||
|
||||
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"
|
||||
''' 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")
|
||||
#
|
||||
|
||||
world_model_points = np.loadtxt(os.path.join(scene_dir, "points_and_normals.txt"))[:,:3]
|
||||
''' preprocess '''
|
||||
# save_scene_data(temp_dir, name)
|
||||
|
||||
pts_dir = os.path.join(temp_dir,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")
|
||||
point_cloud = np.loadtxt(pts_path)
|
||||
sampled_point_cloud = PtsUtil.voxel_downsample_point_cloud(point_cloud, self.voxel_size)
|
||||
indices = np.loadtxt(idx_path, dtype=np.int32)
|
||||
try:
|
||||
len(indices)
|
||||
except:
|
||||
indices = np.array([indices])
|
||||
sample_view_pts_list.append(sampled_point_cloud)
|
||||
scan_points_idx_list.append(indices)
|
||||
|
||||
''' 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 '''
|
||||
|
||||
''' generate 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,
|
||||
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)
|
||||
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 = []
|
||||
|
||||
''' 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)
|
||||
from ipdb import set_trace; set_trace()
|
||||
for idx, coverage_rate in limited_useful_view:
|
||||
path = DataLoadUtil.get_path(temp_dir, 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)
|
||||
|
||||
# ''' 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)
|
||||
|
||||
|
||||
def run(self):
|
||||
@@ -159,32 +189,5 @@ if __name__ == "__main__":
|
||||
# cad_pts_L = PtsUtil.transform_point_cloud(noisy_pts_L, cad_to_cam_L)
|
||||
# np.savetxt(r"test.txt", cad_pts_L)
|
||||
# np.savetxt(r"src.txt", noisy_pts_L)
|
||||
|
||||
''' test view sample '''
|
||||
|
||||
sampled_view_data = ViewSampleUtil.sample_view_data_world_space(
|
||||
model, np.eye(4),
|
||||
voxel_size = 0.005,
|
||||
max_views = 20,
|
||||
min_cam_table_included_degree= 20,
|
||||
random_view_ratio = 0
|
||||
)
|
||||
cam_poses = sampled_view_data["cam_to_world_poses"]
|
||||
cam_poses = np.array(cam_poses)
|
||||
print(cam_poses.shape)
|
||||
def sample_camera_direction(cam_pose, num_samples, sample_distance):
|
||||
cam_position = cam_pose[:3, 3]
|
||||
cam_direction = cam_pose[:3, 2]
|
||||
sampled_points = np.array([cam_position - (i + 1) * sample_distance * cam_direction for i in range(num_samples)])
|
||||
return sampled_points
|
||||
|
||||
all_sampled_points = []
|
||||
for i in range(cam_poses.shape[0]):
|
||||
samped_points = sample_camera_direction(cam_poses[i], 10, 0.02)
|
||||
all_sampled_points.append(samped_points)
|
||||
|
||||
all_sampled_points = np.concatenate(all_sampled_points, axis=0)
|
||||
np.savetxt(r"cam_poses.txt", cam_poses[:, :3, 3])
|
||||
np.savetxt(r"cam_directions.txt", all_sampled_points)
|
||||
|
||||
|
||||
|
||||
@@ -1,87 +0,0 @@
|
||||
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
|
||||
import PytorchBoot.namespace as namespace
|
||||
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("inferencer")
|
||||
class Inferencer(Runner):
|
||||
|
||||
def __init__(self, config_path: str):
|
||||
super().__init__(config_path)
|
||||
self.load_experiment("inferencer")
|
||||
self.reconstruct_config = ConfigManager.get("runner", "reconstruct")
|
||||
self.voxel_size = self.reconstruct_config["voxel_size"]
|
||||
self.max_iter = self.reconstruct_config["max_iter"]
|
||||
|
||||
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_inference(self, model_name):
|
||||
|
||||
''' init robot '''
|
||||
ControlUtil.init()
|
||||
''' take first view '''
|
||||
view_data = CommunicateUtil.get_view_data()
|
||||
first_cam_pts = None
|
||||
first_cam_pose = None
|
||||
combined_pts = first_cam_pts
|
||||
input_data = {
|
||||
"scanned_target_points_num": [first_cam_pts.shape[0]],
|
||||
"scanned_n_to_world_pose_9d": [first_cam_pose],
|
||||
"combined_scanned_pts": combined_pts
|
||||
}
|
||||
''' enter loop '''
|
||||
iter = 0
|
||||
while True:
|
||||
''' inference '''
|
||||
inference_result = CommunicateUtil.get_inference_data(input_data)
|
||||
cam_to_world = inference_result["cam_to_world"]
|
||||
''' set pose '''
|
||||
ControlUtil.set_pose(cam_to_world)
|
||||
''' take view '''
|
||||
view_data = CommunicateUtil.get_view_data()
|
||||
curr_cam_pts = None
|
||||
curr_cam_pose = None
|
||||
''' update combined pts '''
|
||||
combined_pts = np.concatenate([combined_pts, curr_cam_pts], axis=0)
|
||||
combined_pts = PtsUtil.voxel_downsample_point_cloud(combined_pts, voxel_size=self.voxel_size)
|
||||
''' update input data '''
|
||||
input_data["combined_scanned_pts"] = combined_pts
|
||||
input_data["scanned_target_points_num"].append(curr_cam_pts.shape[0])
|
||||
input_data["scanned_n_to_world_pose_9d"].append(curr_cam_pose)
|
||||
|
||||
''' check stop condition '''
|
||||
if iter >= self.max_iter:
|
||||
break
|
||||
|
||||
|
||||
def run(self):
|
||||
self.run_inference()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
model_path = "/home/yan20/nbv_rec/data/test_CAD/test_model/bear_scaled.ply"
|
||||
model = trimesh.load(model_path)
|
||||
test_pts_L = np.loadtxt("/home/yan20/nbv_rec/data/test_CAD/cam_pts_0_L.txt")
|
||||
test_pts_R = np.loadtxt("/home/yan20/nbv_rec/data/test_CAD/cam_pts_0_R.txt")
|
||||
cam_to_world_L = PtsUtil.register_icp(test_pts_L, model)
|
||||
cam_to_world_R = PtsUtil.register_icp(test_pts_R, model)
|
||||
print(cam_to_world_L)
|
||||
print("================================")
|
||||
print(cam_to_world_R)
|
||||
|
||||
Reference in New Issue
Block a user