add compute_plane_normal

modified:   src/maintain/scripts/test.py

src/maintain/scripts/test.py

@@ -14,8 +14,7 @@ import rospkg
 # import open3d as o3d
 # from open3d_ros_helper import open3d_ros_helper as orh
 
-import os
+import math
-import sys
 from rostopic import get_topic_type
 from detection_msgs.msg import BoundingBox, BoundingBoxes
 
@@ -24,11 +23,49 @@ annulus_width = 10
 
 # 2d to 3d
 def computer_2d_3d(x, y, depth_roi, color_intrinsics):
-    pz = np.mean(depth_roi) * 0.001
+    pz = depth_roi[int(y), int(x)] / 1000.0
     px = (x - color_intrinsics[2]) * pz / color_intrinsics[0]
     py = (y - color_intrinsics[5]) * pz / color_intrinsics[4]
     return px, py, pz
 
+def compute_plane_normal(box, depth, color_intrinsics):
+    # 计算相机内参
+    fx = color_intrinsics[0]
+    fy = color_intrinsics[4]
+    cx = color_intrinsics[2]
+    cy = color_intrinsics[5]
+    # 计算矩形中心点坐标
+    x_center = (box[0] + box[2]) / 2
+    y_center = (box[1] + box[3]) / 2
+    z = depth[int(y_center), int(x_center)]
+    x = (x_center - cx) * z / fx
+    y = (y_center - cy) * z / fy
+    # 计算四个顶点坐标
+    x1 = (box[0] - cx) * z / fx
+    y1 = (box[1] - cy) * z / fy
+    x2 = (box[2] - cx) * z / fx
+    y2 = (box[1] - cy) * z / fy
+    x3 = (box[2] - cx) * z / fx
+    y3 = (box[3] - cy) * z / fy
+    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])
+    # 计算平面法向量
+    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  
+
 def compute_normal_vector(p1, p2, p3, p4):
     # Compute two vectors in the plane
     v1 = np.array(p2) - np.array(p1)
@@ -42,7 +79,10 @@ 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)
-    return 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]]
+    return quaternion
 
 def filter_quaternion(quat, quat_prev, alpha):
     if quat_prev is None:
@@ -69,20 +109,22 @@ def box_callback(box, depth, color_info):
             # print(screw_x,screw_y)
 
             depth_array = np.array(depth_image, dtype=np.float32)
-            depth_roi = depth_array[boundingBox.ymin:boundingBox.ymax, boundingBox.xmin:boundingBox.xmax]
+            # depth_roi = depth_array[boundingBox.ymin:boundingBox.ymax, boundingBox.xmin:boundingBox.xmax]
 
-            x, y, z = computer_2d_3d(screw_x, screw_y, depth_roi, color_intrinsics)
+            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
-            p1x, p1y, p1z = computer_2d_3d(boundingBox.xmin-annulus_width, boundingBox.ymin-annulus_width, depth_roi, color_intrinsics)
+            box = [boundingBox.ymin - annulus_width, boundingBox.xmin - annulus_width, boundingBox.ymax + annulus_width, boundingBox.xmax + annulus_width]
-            p2x, p2y, p2z = computer_2d_3d(boundingBox.xmax+annulus_width, boundingBox.ymin-annulus_width, depth_roi, color_intrinsics)
+            # p1x, p1y, p1z = computer_2d_3d(boundingBox.xmin-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_roi, color_intrinsics)
+            # p2x, p2y, p2z = computer_2d_3d(boundingBox.xmax+annulus_width, boundingBox.ymin-annulus_width, depth_array, color_intrinsics)
-            p4x, p4y, p4z = computer_2d_3d(boundingBox.xmin-annulus_width, boundingBox.ymax+annulus_width, depth_roi, color_intrinsics)
+            # p3x, p3y, p3z = computer_2d_3d(boundingBox.xmax+annulus_width, boundingBox.ymax+annulus_width, depth_array, color_intrinsics)
-            p1 = [p1x, p1y, p1z]
+            # p4x, p4y, p4z = computer_2d_3d(boundingBox.xmin-annulus_width, boundingBox.ymax+annulus_width, depth_array, color_intrinsics)
-            p2 = [p2x, p2y, p2z]
+            # p1 = [p1x, p1y, p1z]
-            p3 = [p3x, p3y, p3z]
+            # p2 = [p2x, p2y, p2z]
-            p4 = [p4x, p4y, p4z]
+            # p3 = [p3x, p3y, p3z]
-            normal = compute_normal_vector(p1, p2, p3, p4)
+            # p4 = [p4x, p4y, p4z]
+            # normal_q = compute_normal_vector(p1, p2, p3, p4)
+            normal_q = compute_plane_normal(box, depth_array, color_intrinsics)
 
             # 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)
@@ -91,14 +133,14 @@ def box_callback(box, depth, color_info):
 
             # normal vector to quaternion
             screw_quat = tf.transformations.quaternion_from_euler(0, 0, 0)
-            screw_quat[0] = normal[0]
+            screw_quat[0] = normal_q[0]
-            screw_quat[1] = normal[1]
+            screw_quat[1] = normal_q[1]
-            screw_quat[2] = normal[2]
+            screw_quat[2] = normal_q[2]
-            screw_quat[3] = 0
+            screw_quat[3] = normal_q[3]
 
             # 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)
+            screw_quat_zero_z = tf.transformations.quaternion_from_euler(screw_euler[0], screw_euler[1], 0)
 
 
             # Apply low-pass filter to screw quaternion