cnc-graver-assist/rbs_mill_assist/scripts/grasping_service.py

132 lines
4.7 KiB
Python
Executable file

#!/usr/bin/env python3
import rclpy
import math
from tf2_ros import Buffer, TransformListener, TransformException, TransformStamped
from geometry_msgs.msg import Pose, Quaternion, Point
from rbs_mill_interfaces.srv import GetGraspingPose
from rclpy.node import Node
from rclpy.service import Service
from scipy.spatial.transform import Rotation as R
class GetGraspPickPoses(Node):
def __init__(self) -> None:
super().__init__("get_grasp_pick_poses")
self.srv: Service = self.create_service(
GetGraspingPose, "get_grasping_pose", self.get_grasp_pick_poses
)
self.tf_buffer = Buffer()
self.tf_listener = TransformListener(self.tf_buffer, self)
# Precompute the rotation correction quaternion
self.rotation_correction = self.euler_to_quat(0, math.pi, math.pi/2)
def get_grasp_pick_poses(
self,
request: GetGraspingPose.Request,
response: GetGraspingPose.Response,
) -> GetGraspingPose.Response:
if request.action_type not in ("pick", "place"):
self.get_logger().error(f"Unsupported action type: {request.action_type}")
response.ok = False
return response
# Common logic for pick/place
tf = self.get_tf_pose(request.pose_name, request.relative_to)
#tf_base = self.get_tf_pose(request.pose_name, "base_link")
if tf is None:
response.ok = False
return response
# Apply transformations
main_pose, pre_pose, post_pose = self.process_tf_and_generate_poses(tf)
if request.action_type == "pick":
response.grasp_poses.pose = main_pose
response.grasp_poses.pre_pose = pre_pose
response.grasp_poses.post_pose = post_pose
else:
response.place_poses.pose = main_pose
response.place_poses.pre_pose = pre_pose
response.place_poses.post_pose = post_pose
response.ok = True
return response
def process_tf_and_generate_poses(self, tf: TransformStamped) -> tuple[Pose, Pose, Pose]:
"""Processes TF and generates main poses."""
# Apply rotation correction
tf.transform.rotation = self.multiply_quaternions(
tf.transform.rotation,
self.rotation_correction
)
pose = self.get_pose_from_tf(tf)
return pose, self.add_offset(pose, z=0.1), self.add_offset(pose, z=0.1)
def multiply_quaternions(self, q1: Quaternion, q2: Quaternion) -> Quaternion:
"""Multiplies two quaternions using scipy."""
r1 = R.from_quat([q1.x, q1.y, q1.z, q1.w])
r2 = R.from_quat([q2.x, q2.y, q2.z, q2.w])
result = (r1 * r2).as_quat() # type: ignore
return Quaternion(x=result[0], y=result[1], z=result[2], w=result[3])
def add_offset(self, pose: Pose, x: float = 0.0, y: float = 0.0, z: float = 0.0) -> Pose:
"""Creates a new pose with offset."""
return Pose(
position=Point(
x=pose.position.x + x,
y=pose.position.y + y,
z=pose.position.z + z
),
orientation=pose.orientation
)
def euler_to_quat(self, roll: float, pitch: float, yaw: float) -> Quaternion:
"""Converts Euler angles to quaternion."""
q = R.from_euler("xyz", [roll, pitch, yaw], degrees=False).as_quat()
return Quaternion(x=q[0], y=q[1], z=q[2], w=q[3])
def get_pose_from_tf(self, tf: TransformStamped) -> Pose:
"""Creates Pose from TransformStamped."""
return Pose(
position=Point(
x=tf.transform.translation.x,
y=tf.transform.translation.y,
z=tf.transform.translation.z
),
orientation=Quaternion(
x=tf.transform.rotation.x,
y=tf.transform.rotation.y,
z=tf.transform.rotation.z,
w=tf.transform.rotation.w
)
)
def get_tf_pose(self, source_frame: str, target_frame: str) -> TransformStamped | None:
"""Gets coordinate transform between frames."""
try:
return self.tf_buffer.lookup_transform(
target_frame,
source_frame,
rclpy.time.Time()
)
except TransformException as ex:
self.get_logger().error(
f"Transform error: {target_frame} -> {source_frame}: {ex}"
)
return None
def main():
rclpy.init()
node = GetGraspPickPoses()
try:
rclpy.spin(node)
except KeyboardInterrupt:
pass
finally:
node.destroy_node()
rclpy.shutdown()
if __name__ == "__main__":
main()