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

@@ -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]))
+        annulus_roi = depth_array[boundingBox.ymin-annulus_width:boundingBox.ymax+annulus_width, boundingBox.xmin-annulus_width:boundingBox.xmax+annulus_width]
-        X = X.reshape(-1)
+        normal = calculate_image_edge_plane_normal(annulus_roi)
-        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]
         # 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.x = screw_quat_filtered[0]
-        screw_tf.transform.rotation.y = normal[1]
+        screw_tf.transform.rotation.y = screw_quat_filtered[1]
-        screw_tf.transform.rotation.z = normal[2]
+        screw_tf.transform.rotation.z = screw_quat_filtered[2]
-        screw_tf.transform.rotation.w = 0
+        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 ......")