nbv_rec_control/utils/control_util.py
2024-10-14 19:37:34 +08:00

227 lines
8.7 KiB
Python

import numpy as np
from frankapy import FrankaArm
from autolab_core import RigidTransform
import serial
import time
class ControlUtil:
__fa:FrankaArm = None
__ser: serial.Serial = None
curr_rotation = 0
BASE_TO_WORLD:np.ndarray = np.asarray([
[1, 0, 0, -0.61091665],
[0, 1, 0, -0.00309726],
[0, 0, 1, -0.1136743],
[0, 0, 0, 1]
])
CAMERA_TO_GRIPPER:np.ndarray = np.asarray([
[0, -1, 0, 0.01],
[1, 0, 0, 0],
[0, 0, 1, 0.08],
[0, 0, 0, 1]
])
INIT_GRIPPER_POSE:np.ndarray = np.asarray([
[ 0.46532393, 0.62171798, 0.63002284, 0.21230963],
[ 0.43205618, -0.78075723, 0.45136491, -0.25127173],
[ 0.77251656, 0.06217437, -0.63193429, 0.499957 ],
[ 0. , 0. , 0. , 1. ],
])
@staticmethod
def connect_robot():
if ControlUtil.__fa is None:
ControlUtil.__fa = FrankaArm(robot_num=2)
if ControlUtil.__ser is None:
ControlUtil.__ser = serial.Serial(port="/dev/ttyUSB0", baudrate=115200)
@staticmethod
def franka_reset() -> None:
ControlUtil.__fa.reset_joints()
@staticmethod
def init() -> None:
ControlUtil.franka_reset()
ControlUtil.set_gripper_pose(ControlUtil.INIT_GRIPPER_POSE)
@staticmethod
def get_pose() -> np.ndarray:
gripper_to_base = ControlUtil.get_curr_gripper_to_base_pose()
cam_to_world = ControlUtil.BASE_TO_WORLD @ gripper_to_base @ ControlUtil.CAMERA_TO_GRIPPER
return cam_to_world
@staticmethod
def set_pose(cam_to_world: np.ndarray) -> None:
gripper_to_base = ControlUtil.solve_gripper_to_base(cam_to_world)
ControlUtil.set_gripper_pose(gripper_to_base)
@staticmethod
def rotate_display_table(degree):
turn_directions = {
"left": 1,
"right": 0
}
delta_degree = degree - ControlUtil.curr_rotation
ControlUtil.curr_rotation += delta_degree
print(f"Table rotated {ControlUtil.cnt_rotation} degree")
if degree >= 0:
turn_angle = delta_degree
turn_direction = turn_directions["right"]
else:
turn_angle = -delta_degree
turn_direction = turn_directions["left"]
write_len = ControlUtil.__ser.write(f"CT+TRUNSINGLE({turn_direction},{turn_angle});".encode('utf-8'))
@staticmethod
def get_curr_gripper_to_base_pose() -> np.ndarray:
return ControlUtil.__fa.get_pose().matrix
@staticmethod
def set_gripper_pose(gripper_to_base: np.ndarray) -> None:
gripper_to_base = RigidTransform(rotation=gripper_to_base[:3, :3], translation=gripper_to_base[:3, 3], from_frame="franka_tool", to_frame="world")
ControlUtil.__fa.goto_pose(gripper_to_base, duration=5, use_impedance=False, ignore_errors=False)
@staticmethod
def solve_gripper_to_base(cam_to_world: np.ndarray) -> np.ndarray:
return np.linalg.inv(ControlUtil.BASE_TO_WORLD) @ cam_to_world @ np.linalg.inv(ControlUtil.CAMERA_TO_GRIPPER)
@staticmethod
def sovle_cam_to_world(gripper_to_base: np.ndarray) -> np.ndarray:
return ControlUtil.BASE_TO_WORLD @ gripper_to_base @ ControlUtil.CAMERA_TO_GRIPPER
@staticmethod
def check_limit(new_cam_to_world):
if new_cam_to_world[0,3] > 0 or new_cam_to_world[1,3] > 0:
return False
x = abs(new_cam_to_world[0,3])
y = abs(new_cam_to_world[1,3])
tan_y_x = y/x
if tan_y_x < np.sqrt(3)/3 or tan_y_x > np.sqrt(3):
return False
return True
@staticmethod
def solve_display_table_rot_and_cam_to_world(cam_to_world: np.ndarray) -> tuple:
if ControlUtil.check_limit(cam_to_world):
return 0, cam_to_world
else:
min_display_table_rot = 180
min_new_cam_to_world = None
for display_table_rot in np.linspace(0.1,360, 1800):
display_table_rot_z_pose = ControlUtil.get_z_axis_rot_mat(display_table_rot)
new_cam_to_world = np.linalg.inv(display_table_rot_z_pose) @ cam_to_world
if ControlUtil.check_limit(new_cam_to_world):
if display_table_rot < min_display_table_rot:
min_display_table_rot, min_new_cam_to_world = display_table_rot, new_cam_to_world
if abs(display_table_rot - 360) < min_display_table_rot:
min_display_table_rot, min_new_cam_to_world = display_table_rot - 360, new_cam_to_world
if min_new_cam_to_world is None:
raise ValueError("No valid display table rotation found")
return min_display_table_rot, min_new_cam_to_world
@staticmethod
def get_z_axis_rot_mat(degree):
radian = np.radians(degree)
return np.array([
[np.cos(radian), -np.sin(radian), 0, 0],
[np.sin(radian), np.cos(radian), 0, 0],
[0, 0, 1, 0],
[0, 0, 0, 1]
])
@staticmethod
def get_gripper_to_base_axis_angle(gripper_to_base: np.ndarray) -> bool:
rot_mat = gripper_to_base[:3,:3]
gripper_z_axis = rot_mat[:,2]
base_x_axis = np.array([1,0,0])
angle = np.arccos(np.dot(gripper_z_axis, base_x_axis))
return angle
@staticmethod
def move_to(pose: np.ndarray):
rot_degree, cam_to_world = ControlUtil.solve_display_table_rot_and_cam_to_world(pose)
exec_time = abs(rot_degree)/9
start_time = time.time()
ControlUtil.rotate_display_table(rot_degree)
ControlUtil.set_pose(cam_to_world)
end_time = time.time()
print(f"Move to pose with rotation {rot_degree} degree, exec time: {end_time - start_time}|exec time: {exec_time}")
if end_time - start_time < exec_time:
time.sleep(exec_time - (end_time - start_time))
# ----------- Debug Test -------------
if __name__ == "__main__":
ControlUtil.connect_robot()
# ControlUtil.franka_reset()
def main_test():
print(ControlUtil.get_curr_gripper_to_base_pose())
ControlUtil.init()
def rotate_back(rotation):
ControlUtil.rotate_display_table(-rotation)
#main_test()
import sys; sys.path.append("/home/yan20/nbv_rec/project/franka_control")
from utils.communicate_util import CommunicateUtil
import ipdb
ControlUtil.init()
view_data_0 = CommunicateUtil.get_view_data(init=True)
depth_extrinsics_0 = view_data_0["depth_extrinsics"]
cam_to_world_0 = ControlUtil.get_pose()
print("cam_extrinsics_0")
print(depth_extrinsics_0)
print("cam_to_world_0")
print(cam_to_world_0)
ipdb.set_trace()
TEST_POSE:np.ndarray = np.asarray([
[ 0.46532393, 0.62171798, 0.63002284, 0.30230963],
[ 0.43205618, -0.78075723, 0.45136491, -0.29127173],
[ 0.77251656, 0.06217437, -0.63193429, 0.559957 ],
[ 0. , 0. , 0. , 1. ],
])
TEST_POSE_CAM_TO_WORLD = ControlUtil.BASE_TO_WORLD @ TEST_POSE @ ControlUtil.CAMERA_TO_GRIPPER
ControlUtil.move_to(TEST_POSE_CAM_TO_WORLD)
view_data_1 = CommunicateUtil.get_view_data()
depth_extrinsics_1 = view_data_1["depth_extrinsics"]
depth_extrinsics_1[:3,3] = depth_extrinsics_1[:3,3] / 1000
cam_to_world_1 = ControlUtil.get_pose()
print("cam_extrinsics_1")
print(depth_extrinsics_1)
print("cam_to_world_1")
print(TEST_POSE_CAM_TO_WORLD)
actual_cam_to_world_1 = cam_to_world_0 @ depth_extrinsics_1
print("actual_cam_to_world_1")
print(actual_cam_to_world_1)
ipdb.set_trace()
TEST_POSE_2:np.ndarray = np.asarray(
[[ 0.74398544, -0.61922696, 0.251049, 0.47000935],
[-0.47287207, -0.75338888, -0.45692666, 0.20843903],
[ 0.47207883 , 0.22123272, -0.85334192, 0.57863381],
[ 0. , 0. , 0. , 1. , ]]
)
TEST_POSE_CAM_TO_WORLD_2 = ControlUtil.BASE_TO_WORLD @ TEST_POSE_2 @ ControlUtil.CAMERA_TO_GRIPPER
#ControlUtil.move_to(TEST_POSE_CAM_TO_WORLD_2)
ControlUtil.set_pose(TEST_POSE_CAM_TO_WORLD_2)
view_data_2 = CommunicateUtil.get_view_data()
depth_extrinsics_2 = view_data_2["depth_extrinsics"]
depth_extrinsics_2[:3,3] = depth_extrinsics_2[:3,3] / 1000
cam_to_world_2 = ControlUtil.get_pose()
print("cam_extrinsics_2")
print(depth_extrinsics_2)
print("cam_to_world_2")
print(TEST_POSE_CAM_TO_WORLD_2)
actual_cam_to_world_2 = cam_to_world_0 @ depth_extrinsics_2
print("actual_cam_to_world_2")
print(actual_cam_to_world_2)
ipdb.set_trace()