From 825af0226d778b11e6fcb001194e9ca73de618de Mon Sep 17 00:00:00 2001 From: wxchen Date: Wed, 10 May 2023 16:08:34 +0800 Subject: [PATCH] use pointcloud to computer rotation modified: src/maintain/launch/maintain.launch new file: src/maintain/scripts/maintain.py modified: src/maintain/scripts/test.py modified: src/yolov5_ros/launch/yolov5.launch --- src/maintain/launch/maintain.launch | 2 +- src/maintain/scripts/maintain.py | 136 ++++++++++++++++++++++++++++ src/maintain/scripts/test.py | 53 +++++++---- src/yolov5_ros/launch/yolov5.launch | 4 +- 4 files changed, 173 insertions(+), 22 deletions(-) create mode 100755 src/maintain/scripts/maintain.py diff --git a/src/maintain/launch/maintain.launch b/src/maintain/launch/maintain.launch index fbb70b0..905ab12 100644 --- a/src/maintain/launch/maintain.launch +++ b/src/maintain/launch/maintain.launch @@ -1,4 +1,4 @@ - + diff --git a/src/maintain/scripts/maintain.py b/src/maintain/scripts/maintain.py new file mode 100755 index 0000000..14d2f0f --- /dev/null +++ b/src/maintain/scripts/maintain.py @@ -0,0 +1,136 @@ +#! /home/da/miniconda3/envs/gsmini/bin/python + +import rospy +import numpy as np +import open3d as o3d +from sensor_msgs.msg import Image , CameraInfo, PointCloud2 +from detection_msgs.msg import BoundingBox, BoundingBoxes +import sensor_msgs.point_cloud2 as pc2 +import cv_bridge +from cv_bridge import CvBridge, CvBridgeError +import cv2 +import tf2_ros +import tf +from geometry_msgs.msg import PoseStamped, TransformStamped + +bridge = CvBridge() +color_intrinsics = None +cloud = None +box = None +d_width = 100 + +def camera_info_callback(msg): + global color_intrinsics + color_intrinsics = msg + +def depth_image_callback(msg): + global depth_image + depth_image = bridge.imgmsg_to_cv2(msg, '16UC1') + +def point_cloud_callback(msg): + global cloud + cloud = pc2.read_points(msg, field_names=("x", "y", "z"), skip_nans=True) + +def bounding_boxes_callback(msg): + global box + for bounding_box in msg.bounding_boxes: + # Assuming there's only one box, you can add a condition to filter the boxes if needed + box = [bounding_box.xmin - d_width, bounding_box.ymin - d_width, bounding_box.xmax + d_width, bounding_box.ymax + d_width] + +def main(): + rospy.init_node("plane_fitting_node") + + rospy.Subscriber("/camera/color/camera_info", CameraInfo, camera_info_callback) + rospy.Subscriber("/camera/aligned_depth_to_color/image_raw", Image, depth_image_callback) + rospy.Subscriber("/camera/depth/color/points", PointCloud2, point_cloud_callback) + rospy.Subscriber("/yolov5/detections", BoundingBoxes, bounding_boxes_callback) + tf_broadcaster = tf2_ros.TransformBroadcaster() + + plane_pub = rospy.Publisher("/plane_pose", PoseStamped, queue_size=10) + + rate = rospy.Rate(10) # 10 Hz + + while not rospy.is_shutdown(): + if color_intrinsics is not None and cloud is not None and box is not None: + # Get the 3D points corresponding to the box + fx, fy = color_intrinsics.K[0], color_intrinsics.K[4] + cx, cy = color_intrinsics.K[2], color_intrinsics.K[5] + points = [] + + center_x = (box[0] + box[2]) / 2 + center_y = (box[1] + box[3]) / 2 + + depth_array = np.array(depth_image, dtype=np.float32) + pz = depth_array[int(center_y), int(center_x)] / 1000.0 + px = (center_x - color_intrinsics.K[2]) * pz / color_intrinsics.K[0] + py = (center_y - color_intrinsics.K[5]) * pz / color_intrinsics.K[4] + + rospy.loginfo("Center point: {}".format([px, py, pz])) + + screw_point = None + for x, y, z in cloud: + if z != 0: + u = int(np.round((x * fx) / z + cx)) + v = int(np.round((y * fy) / z + cy)) + if u == center_x and v == center_y: + screw_point = [x, y, z] + if u >= box[0] and u <= box[2] and v >= box[1] and v <= box[3]: + points.append([x, y, z]) + points = np.array(points) + + if px != 0 and py != 0 and pz != 0: + + # rospy.loginfo("Screw point: {}".format(screw_point)) + + # Fit a plane to the points + pcd = o3d.geometry.PointCloud() + pcd.points = o3d.utility.Vector3dVector(points) + plane_model, inliers = pcd.segment_plane(distance_threshold=0.02, ransac_n=3, num_iterations=100) + [a, b, c, d] = plane_model + + # Calculate the rotation between the plane normal and the Z axis + normal = np.array([a, b, c]) + z_axis = np.array([0, 0, 1]) + cos_theta = np.dot(normal, z_axis) / (np.linalg.norm(normal) * np.linalg.norm(z_axis)) + theta = np.arccos(cos_theta) + rotation_axis = np.cross(z_axis, normal) + rotation_axis = rotation_axis / np.linalg.norm(rotation_axis) + quaternion = np.hstack((rotation_axis * np.sin(theta / 2), [np.cos(theta / 2)])) + + + # Publish the plane pose + # plane_pose = PoseStamped() + # plane_pose.header.stamp = rospy.Time.now() + # plane_pose.header.frame_id = "camera_color_optical_frame" + # plane_pose.pose.position.x = screw_point[0] + # plane_pose.pose.position.y = screw_point[1] + # plane_pose.pose.position.z = -d / np.linalg.norm(normal) + # plane_pose.pose.orientation.x = quaternion[0] + # plane_pose.pose.orientation.y = quaternion[1] + # plane_pose.pose.orientation.z = quaternion[2] + # plane_pose.pose.orientation.w = quaternion[3] + # plane_pub.publish(plane_pose) + + # publish screw tf + screw_tf = TransformStamped() + screw_tf.header.stamp = rospy.Time.now() + screw_tf.header.frame_id = "camera_color_optical_frame" + screw_tf.child_frame_id = "screw" + screw_tf.transform.translation.x = px + screw_tf.transform.translation.y = py + screw_tf.transform.translation.z = pz + screw_tf.transform.rotation.x = quaternion[0] + screw_tf.transform.rotation.y = quaternion[1] + screw_tf.transform.rotation.z = quaternion[2] + screw_tf.transform.rotation.w = quaternion[3] + + tf_broadcaster.sendTransform(screw_tf) + + + rate.sleep() + +if __name__ == "__main__": + try: + main() + except rospy.ROSInterruptException: + pass \ No newline at end of file diff --git a/src/maintain/scripts/test.py b/src/maintain/scripts/test.py index 0a36cd2..606321e 100755 --- a/src/maintain/scripts/test.py +++ b/src/maintain/scripts/test.py @@ -19,7 +19,7 @@ from rostopic import get_topic_type from detection_msgs.msg import BoundingBox, BoundingBoxes bridge = CvBridge() -annulus_width = 10 +annulus_width = 20 # 2d to 3d def computer_2d_3d(x, y, depth_roi, color_intrinsics): @@ -37,7 +37,7 @@ def compute_plane_normal(box, depth, color_intrinsics): # 计算矩形中心点坐标 x_center = (box[0] + box[2]) / 2 y_center = (box[1] + box[3]) / 2 - z = depth[int(y_center), int(x_center)] + z = depth[int(y_center), int(x_center)] / 1000 x = (x_center - cx) * z / fx y = (y_center - cy) * z / fy # 计算四个顶点坐标 @@ -50,21 +50,28 @@ def compute_plane_normal(box, depth, color_intrinsics): x4 = (box[0] - cx) * z / fx y4 = (box[3] - cy) * z / fy # 计算矩形边缘向量 - v1 = np.array([x2 - x1, y2 - y1, depth[int(box[1]), int(box[0])] - z]) - v2 = np.array([x3 - x2, y3 - y2, depth[int(box[1]), int(box[2])] - z]) - v3 = np.array([x4 - x3, y4 - y3, depth[int(box[3]), int(box[2])] - z]) - v4 = np.array([x1 - x4, y1 - y4, depth[int(box[3]), int(box[0])] - z]) + v1 = np.array([x2 - x1, y2 - y1, depth[int(box[1]), int(box[0])] / 1000 - z]) + v2 = np.array([x3 - x2, y3 - y2, depth[int(box[1]), int(box[2])] / 1000 - z]) + v3 = np.array([x4 - x3, y4 - y3, depth[int(box[3]), int(box[2])] / 1000 - z]) + v4 = np.array([x1 - x4, y1 - y4, depth[int(box[3]), int(box[0])] / 1000 - z]) # 计算平面法向量 normal = np.cross(v1, v2) normal += np.cross(v2, v3) normal += np.cross(v3, v4) normal += np.cross(v4, v1) normal /= np.linalg.norm(normal) - # 将法向量转换为四元数表示 - theta = math.acos(normal[2]) - sin_theta_2 = math.sin(theta/2) - quaternion = [math.cos(theta/2), sin_theta_2 * normal[0], sin_theta_2 * normal[1], sin_theta_2 * normal[2]] - return quaternion + + # 计算法向量相对于参考向量的旋转角度和旋转轴 + ref_vector = np.array([0, 0, 1]) + normal_vector = normal + angle = math.acos(np.dot(ref_vector, normal_vector) / (np.linalg.norm(ref_vector) * np.linalg.norm(normal_vector))) + axis = np.cross(ref_vector, normal_vector) + axis = axis / np.linalg.norm(axis) + + # 将旋转角度和旋转轴转换为四元数 + qx, qy, qz, qw = tf.transformations.quaternion_about_axis(angle, axis) + quaternion = [qx, qy, qz, qw] + return quaternion def compute_normal_vector(p1, p2, p3, p4): # Compute two vectors in the plane @@ -79,9 +86,16 @@ def compute_normal_vector(p1, p2, p3, p4): n = -n # Normalize the normal vector to obtain a unit vector n = n / np.linalg.norm(n) - theta = math.acos(n[2]) - sin_theta_2 = math.sin(theta/2) - quaternion = [math.cos(theta/2), sin_theta_2 * n[0], sin_theta_2 * n[1], sin_theta_2 * n[2]] + # 计算法向量相对于参考向量的旋转角度和旋转轴 + ref_vector = np.array([0, 0, 1]) + normal_vector = n + angle = math.acos(np.dot(ref_vector, normal_vector) / (np.linalg.norm(ref_vector) * np.linalg.norm(normal_vector))) + axis = np.cross(ref_vector, normal_vector) + axis = axis / np.linalg.norm(axis) + + # 将旋转角度和旋转轴转换为四元数 + qx, qy, qz, qw = tf.transformations.quaternion_about_axis(angle, axis) + quaternion = [qx, qy, qz, qw] return quaternion def filter_quaternion(quat, quat_prev, alpha): @@ -114,7 +128,7 @@ def box_callback(box, depth, color_info): x, y, z = computer_2d_3d(screw_x, screw_y, depth_array, color_intrinsics) # rospy.loginfo("screw pose: x: %f, y: %f, z: %f", x, y, z) # calculate normal direction of screw area - box = [boundingBox.ymin - annulus_width, boundingBox.xmin - annulus_width, boundingBox.ymax + annulus_width, boundingBox.xmax + annulus_width] + box = [boundingBox.xmin - annulus_width, boundingBox.ymin - annulus_width, boundingBox.xmax + annulus_width, boundingBox.ymax + annulus_width] # p1x, p1y, p1z = computer_2d_3d(boundingBox.xmin-annulus_width, boundingBox.ymin-annulus_width, depth_array, color_intrinsics) # p2x, p2y, p2z = computer_2d_3d(boundingBox.xmax+annulus_width, boundingBox.ymin-annulus_width, depth_array, color_intrinsics) # p3x, p3y, p3z = computer_2d_3d(boundingBox.xmax+annulus_width, boundingBox.ymax+annulus_width, depth_array, color_intrinsics) @@ -141,6 +155,7 @@ def box_callback(box, depth, color_info): screw_euler = tf.transformations.euler_from_quaternion(screw_quat) screw_quat_zero_z = tf.transformations.quaternion_from_euler(screw_euler[0], screw_euler[1], 0) + print(screw_euler) # Apply low-pass filter to screw quaternion alpha = 0.4 @@ -157,10 +172,10 @@ def box_callback(box, depth, color_info): screw_tf.transform.translation.x = x screw_tf.transform.translation.y = y screw_tf.transform.translation.z = z - screw_tf.transform.rotation.x = screw_quat[0] - screw_tf.transform.rotation.y = screw_quat[1] - screw_tf.transform.rotation.z = screw_quat[2] - screw_tf.transform.rotation.w = screw_quat[3] + screw_tf.transform.rotation.x = screw_quat_filtered[0] + screw_tf.transform.rotation.y = screw_quat_filtered[1] + screw_tf.transform.rotation.z = screw_quat_filtered[2] + screw_tf.transform.rotation.w = screw_quat_filtered[3] tf_broadcaster.sendTransform(screw_tf) diff --git a/src/yolov5_ros/launch/yolov5.launch b/src/yolov5_ros/launch/yolov5.launch index 5cade72..7a85c2e 100644 --- a/src/yolov5_ros/launch/yolov5.launch +++ b/src/yolov5_ros/launch/yolov5.launch @@ -2,7 +2,7 @@ - + @@ -23,7 +23,7 @@ - +