From c2f252dce6ce6f100f44a9e9a892c72f7acb63f0 Mon Sep 17 00:00:00 2001 From: wxchen Date: Mon, 8 May 2023 19:36:52 +0800 Subject: [PATCH] Set the z-axis rotation to zero, Apply low-pass filter to screw quaternion modified: src/maintain/scripts/test.py --- src/maintain/scripts/test.py | 113 ++++++++++++++++++++++++++++++----- 1 file changed, 97 insertions(+), 16 deletions(-) diff --git a/src/maintain/scripts/test.py b/src/maintain/scripts/test.py index d308176..bf48f94 100755 --- a/src/maintain/scripts/test.py +++ b/src/maintain/scripts/test.py @@ -5,8 +5,9 @@ import cv2 as cv from matplotlib import pyplot as plt import rospy import tf2_ros +import tf from sensor_msgs.msg import Image , CameraInfo -from geometry_msgs.msg import PoseStamped, TransformStamped +from geometry_msgs.msg import PoseStamped, TransformStamped, Quaternion import message_filters from cv_bridge import CvBridge, CvBridgeError import rospkg @@ -17,6 +18,73 @@ from rostopic import get_topic_type from detection_msgs.msg import BoundingBox, BoundingBoxes bridge = CvBridge() +annulus_width = 10 + +def calculate_image_edge_plane_normal(depth_roi): + # Get the shape of the depth_roi + height, width = depth_roi.shape + + # Get the edges of the ROI + left_edge = [(0, y) for y in range(height)] + right_edge = [(width-1, y) for y in range(height)] + top_edge = [(x, 0) for x in range(width)] + bottom_edge = [(x, height-1) for x in range(width)] + edges = left_edge + right_edge + top_edge + bottom_edge + + # Create a 2D grid of X and Y coordinates + X, Y = np.meshgrid(np.arange(width), np.arange(height)) + + # Reshape the X, Y, and depth_roi arrays into one-dimensional arrays + X = X.reshape(-1) + Y = Y.reshape(-1) + Z = depth_roi.reshape(-1) + + # Stack the X, Y, and depth_roi arrays vertically to create a 3D array of points in the form of [X, Y, Z] + points = np.vstack([X, Y, Z]).T + + # Compute the mean depth value of the edges + edge_depths = [] + for edge_point in edges: + edge_depths.append(depth_roi[edge_point[1], edge_point[0]]) + mean_depth = np.mean(edge_depths) + + # Create a mask to extract the points on the edges + mask = np.zeros_like(depth_roi, dtype=np.uint8) + for edge_point in edges: + mask[edge_point[1], edge_point[0]] = 1 + masked_depth_roi = depth_roi * mask + + # Extract the 3D coordinates of the points on the edges + edge_points = [] + for edge_point in edges: + edge_points.append([edge_point[0], edge_point[1], masked_depth_roi[edge_point[1], edge_point[0]]]) + + # Convert the list of edge points to a numpy array + edge_points = np.array(edge_points) + + # Shift the edge points so that the mean depth value is at the origin + edge_points = edge_points - np.array([width/2, height/2, mean_depth]) + + # Compute the singular value decomposition (SVD) of the edge points + U, S, V = np.linalg.svd(edge_points) + + # Extract the normal vector of the plane that best fits the edge points from the right-singular vector corresponding to the smallest singular value + normal = V[2] + + return normal + +def filter_quaternion(quat, quat_prev, alpha): + if quat_prev is None: + quat_prev = quat + # Apply low-pass filter to quaternion + quat_filtered = np.zeros(4) + for i in range(4): + quat_filtered[i] = alpha * quat[i] + (1-alpha) * quat_prev[i] + # Normalize the quaternion + quat_filtered = quat_filtered / np.linalg.norm(quat_filtered) + return quat_filtered + + def box_callback(box, depth, color_info): try: color_intrinsics = color_info.K @@ -33,18 +101,11 @@ def box_callback(box, depth, color_info): z = np.mean(depth_roi) * 0.001 x = (screw_x - color_intrinsics[2]) * z / color_intrinsics[0] y = (screw_y - color_intrinsics[5]) * z / color_intrinsics[4] - rospy.loginfo("screw pose: x: %f, y: %f, z: %f", x, y, z) + # rospy.loginfo("screw pose: x: %f, y: %f, z: %f", x, y, z) # calculate normal direction of screw area - - X,Y = np.meshgrid(np.arange(depth_roi.shape[1]), np.arange(depth_roi.shape[0])) - X = X.reshape(-1) - Y = Y.reshape(-1) - Z = depth_roi.reshape(-1) - points = np.vstack([X, Y, Z]).T - center = np.mean(points, axis=0) - points = points - center - U, S, V = np.linalg.svd(points) - normal = V[2] + + annulus_roi = depth_array[boundingBox.ymin-annulus_width:boundingBox.ymax+annulus_width, boundingBox.xmin-annulus_width:boundingBox.xmax+annulus_width] + normal = calculate_image_edge_plane_normal(annulus_roi) # print(normal) # publish screw pose @@ -61,6 +122,24 @@ def box_callback(box, depth, color_info): # pose_pub.publish(screw_pose) + # normal vector to quaternion + screw_quat = tf.transformations.quaternion_from_euler(0, 0, 0) + screw_quat[0] = normal[0] + screw_quat[1] = normal[1] + screw_quat[2] = normal[2] + screw_quat[3] = 0 + # quaternion to euler + screw_euler = tf.transformations.euler_from_quaternion(screw_quat) + screw_quat = tf.transformations.quaternion_from_euler(screw_euler[0], screw_euler[1], 0) + + + # Apply low-pass filter to screw quaternion + alpha = 0.4 + global screw_quat_prev + screw_quat_filtered = filter_quaternion(screw_quat, screw_quat_prev, alpha) + screw_quat_prev = screw_quat_filtered + + # publish screw tf screw_tf = TransformStamped() screw_tf.header.stamp = rospy.Time.now() @@ -69,10 +148,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 = normal[0] - screw_tf.transform.rotation.y = normal[1] - screw_tf.transform.rotation.z = normal[2] - screw_tf.transform.rotation.w = 0 + 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) @@ -84,6 +163,8 @@ def box_callback(box, depth, color_info): if __name__ == "__main__": + global screw_quat_prev + screw_quat_prev = None rospy.init_node("maintain") rospy.loginfo("maintain task start ......")