caculate the x y z of the object in the camera frame

copied:     src/maintain/scripts/test.py -> src/maintain/scripts/test copy.py
	modified:   src/maintain/scripts/test.py
	modified:   src/yolov5_ros/launch/yolov5.launch
	modified:   src/yolov5_ros/src/yolov5/best.pt

src/maintain/scripts/test copy.py

@@ -0,0 +1,93 @@
+#! /home/wxchen/.conda/envs/gsmini/bin/python
+
+import numpy as np
+import cv2 as cv
+from matplotlib import pyplot as plt
+import rospy
+from sensor_msgs.msg import Image
+import message_filters
+from cv_bridge import CvBridge, CvBridgeError
+import rospkg
+
+MIN_MATCH_COUNT = 10
+pkg_path = rospkg.RosPack().get_path('maintain')
+rospy.loginfo(pkg_path)
+img_template = cv.imread(pkg_path + '/scripts/tt.png',0)
+
+def callback(rgb, depth):
+    rospy.loginfo("callback")
+    bridge = CvBridge()
+    # rospy.loginfo(rgb.header.stamp)
+    # rospy.loginfo(depth.header.stamp)
+    try:
+        rgb_image = bridge.imgmsg_to_cv2(rgb, 'bgr8')
+        depth_image = bridge.imgmsg_to_cv2(depth, '16UC1')
+
+        img_matcher = matcher(rgb_image)
+        cv.imshow("img_matcher", img_matcher)
+        cv.waitKey(1000)
+    
+    except CvBridgeError as e:
+        print(e)
+
+def matcher(img):
+
+    try:
+        # Initiate SIFT detector
+        sift = cv.SIFT_create()
+
+        # find the keypoints and descriptors with SIFT
+        kp1, des1 = sift.detectAndCompute(img_template,None)
+        kp2, des2 = sift.detectAndCompute(img,None)
+
+        FLANN_INDEX_KDTREE = 1
+        index_params = dict(algorithm = FLANN_INDEX_KDTREE, trees = 5)
+        search_params = dict(checks = 50)
+
+        flann = cv.FlannBasedMatcher(index_params, search_params)
+        matches = flann.knnMatch(des1,des2,k=2)
+
+        # store all the good matches as per Lowe's ratio test.
+        good = []
+        for m,n in matches:
+            if m.distance < 0.7*n.distance:
+                good.append(m)
+
+        if len(good)>MIN_MATCH_COUNT:
+            src_pts = np.float32([ kp1[m.queryIdx].pt for m in good ]).reshape(-1,1,2)
+            dst_pts = np.float32([ kp2[m.trainIdx].pt for m in good ]).reshape(-1,1,2)
+
+            M, mask = cv.findHomography(src_pts, dst_pts, cv.RANSAC,5.0)
+            matchesMask = mask.ravel().tolist()
+
+            h,w = img_template.shape
+            pts = np.float32([ [0,0],[0,h-1],[w-1,h-1],[w-1,0] ]).reshape(-1,1,2)
+            dst = cv.perspectiveTransform(pts,M)
+
+            roi = img[np.int32(dst)[0][0][1]:np.int32(dst)[2][0][1], np.int32(dst)[0][0][0]:np.int32(dst)[2][0][0]]
+            # roi = detect_black(roi)
+            
+            # img2 = cv.polylines(img2,[np.int32(dst)],True,255,3, cv.LINE_AA)
+        else:
+            print( "Not enough matches are found - {}/{}".format(len(good), MIN_MATCH_COUNT) )
+
+        return roi
+    except Exception as e:
+        print(e)
+
+
+
+
+if __name__ == "__main__":
+
+    rospy.init_node("maintain")  
+    rospy.loginfo("maintain task start ......") 
+
+    rgb_sub = message_filters.Subscriber("/camera/color/image_raw", Image)
+    depth_sub = message_filters.Subscriber("/camera/aligned_depth_to_color/image_raw", Image)
+
+    ts = message_filters.TimeSynchronizer([rgb_sub, depth_sub], 1)
+    ts.registerCallback(callback)
+
+
+    rospy.spin()

src/maintain/scripts/test.py

@@ -4,90 +4,82 @@ import numpy as np
 import cv2 as cv
 from matplotlib import pyplot as plt
 import rospy
-from sensor_msgs.msg import Image
+from sensor_msgs.msg import Image , CameraInfo
+from geometry_msgs.msg import PoseStamped
 import message_filters
 from cv_bridge import CvBridge, CvBridgeError
 import rospkg
 
-MIN_MATCH_COUNT = 10
+import os
-pkg_path = rospkg.RosPack().get_path('maintain')
+import sys
-rospy.loginfo(pkg_path)
+from rostopic import get_topic_type
-img_template = cv.imread(pkg_path + '/scripts/tt.png',0)
+from detection_msgs.msg import BoundingBox, BoundingBoxes
 
-def callback(rgb, depth):
+bridge = CvBridge()
-    rospy.loginfo("callback")
+def box_callback(box, depth, color_info):
-    bridge = CvBridge()
-    # rospy.loginfo(rgb.header.stamp)
-    # rospy.loginfo(depth.header.stamp)
     try:
-        rgb_image = bridge.imgmsg_to_cv2(rgb, 'bgr8')
+        color_intrinsics = color_info.K
         depth_image = bridge.imgmsg_to_cv2(depth, '16UC1')
+        # get the center of screw 
+        boundingBox = box.bounding_boxes[0]
+        screw_x = (boundingBox.xmax + boundingBox.xmin) / 2
+        screw_y = (boundingBox.ymax + boundingBox.ymin) / 2
+        # print(screw_x,screw_y)
 
-        img_matcher = matcher(rgb_image)
+        depth_array = np.array(depth_image, dtype=np.float32)
-        cv.imshow("img_matcher", img_matcher)
+        depth_roi = depth_array[boundingBox.ymin:boundingBox.ymax, boundingBox.xmin:boundingBox.xmax]
-        cv.waitKey(1000)
-    
-    except CvBridgeError as e:
-        print(e)
 
-def matcher(img):
+        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]
+        # print(x,y,z)
+        # calculate normal direction of screw area
 
-    try:
+        X,Y = np.meshgrid(np.arange(depth_roi.shape[1]), np.arange(depth_roi.shape[0]))
-        # Initiate SIFT detector
+        X = X.reshape(-1)
-        sift = cv.SIFT_create()
+        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)
 
-        # find the keypoints and descriptors with SIFT
+        # publish screw pose
-        kp1, des1 = sift.detectAndCompute(img_template,None)
+        screw_pose = PoseStamped()
-        kp2, des2 = sift.detectAndCompute(img,None)
+        screw_pose.header.stamp = rospy.Time.now()
+        screw_pose.header.frame_id = "camera_color_optical_frame"
+        screw_pose.pose.position.x = x
+        screw_pose.pose.position.y = y
+        screw_pose.pose.position.z = z
+        screw_pose.pose.orientation.x = 0
+        screw_pose.pose.orientation.y = 0
+        screw_pose.pose.orientation.z = 0
+        screw_pose.pose.orientation.w = 1
 
-        FLANN_INDEX_KDTREE = 1
+        pose_pub.publish(screw_pose)
-        index_params = dict(algorithm = FLANN_INDEX_KDTREE, trees = 5)
-        search_params = dict(checks = 50)
 
-        flann = cv.FlannBasedMatcher(index_params, search_params)
-        matches = flann.knnMatch(des1,des2,k=2)
 
-        # store all the good matches as per Lowe's ratio test.
-        good = []
-        for m,n in matches:
-            if m.distance < 0.7*n.distance:
-                good.append(m)
 
-        if len(good)>MIN_MATCH_COUNT:
-            src_pts = np.float32([ kp1[m.queryIdx].pt for m in good ]).reshape(-1,1,2)
-            dst_pts = np.float32([ kp2[m.trainIdx].pt for m in good ]).reshape(-1,1,2)
-
-            M, mask = cv.findHomography(src_pts, dst_pts, cv.RANSAC,5.0)
-            matchesMask = mask.ravel().tolist()
-
-            h,w = img_template.shape
-            pts = np.float32([ [0,0],[0,h-1],[w-1,h-1],[w-1,0] ]).reshape(-1,1,2)
-            dst = cv.perspectiveTransform(pts,M)
-
-            roi = img[np.int32(dst)[0][0][1]:np.int32(dst)[2][0][1], np.int32(dst)[0][0][0]:np.int32(dst)[2][0][0]]
-            # roi = detect_black(roi)
-            
-            # img2 = cv.polylines(img2,[np.int32(dst)],True,255,3, cv.LINE_AA)
-        else:
-            print( "Not enough matches are found - {}/{}".format(len(good), MIN_MATCH_COUNT) )
-
-        return roi
     except Exception as e:
         print(e)
 
 
 
-
 if __name__ == "__main__":
 
     rospy.init_node("maintain")  
     rospy.loginfo("maintain task start ......") 
 
-    rgb_sub = message_filters.Subscriber("/camera/color/image_raw", Image)
+    box_sub = message_filters.Subscriber("/yolov5/detections", BoundingBoxes)
     depth_sub = message_filters.Subscriber("/camera/aligned_depth_to_color/image_raw", Image)
+    color_info = message_filters.Subscriber("/camera/color/camera_info", CameraInfo)
 
-    ts = message_filters.TimeSynchronizer([rgb_sub, depth_sub], 1)
+    pose_pub = rospy.Publisher('/screw_pose', PoseStamped, queue_size=10)
-    ts.registerCallback(callback)
+
+    ts = message_filters.TimeSynchronizer([box_sub, depth_sub, color_info], 1)
+    ts.registerCallback(box_callback)
 
 
     rospy.spin()

src/yolov5_ros/launch/yolov5.launch

@@ -16,7 +16,7 @@
     <arg name="inference_size_w" default="640"/>
 
     <!-- Visualize using OpenCV window -->
-    <arg name="view_image" default="true"/>
+    <arg name="view_image" default="false"/>
 
     <!-- ROS topics -->
     <arg name="input_image_topic" default="/camera/color/image_raw"/>