diff --git a/src/maintain/scripts/test copy.py b/src/maintain/scripts/test copy.py new file mode 100755 index 0000000..59582a1 --- /dev/null +++ b/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() \ No newline at end of file diff --git a/src/maintain/scripts/test.py b/src/maintain/scripts/test.py index 59582a1..457f63f 100755 --- a/src/maintain/scripts/test.py +++ b/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 -pkg_path = rospkg.RosPack().get_path('maintain') -rospy.loginfo(pkg_path) -img_template = cv.imread(pkg_path + '/scripts/tt.png',0) +import os +import sys +from rostopic import get_topic_type +from detection_msgs.msg import BoundingBox, BoundingBoxes -def callback(rgb, depth): - rospy.loginfo("callback") - bridge = CvBridge() - # rospy.loginfo(rgb.header.stamp) - # rospy.loginfo(depth.header.stamp) +bridge = CvBridge() +def box_callback(box, depth, color_info): 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) - cv.imshow("img_matcher", img_matcher) - cv.waitKey(1000) - - except CvBridgeError as e: - print(e) + depth_array = np.array(depth_image, dtype=np.float32) + depth_roi = depth_array[boundingBox.ymin:boundingBox.ymax, boundingBox.xmin:boundingBox.xmax] -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: - # Initiate SIFT detector - sift = cv.SIFT_create() + 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] + # print(normal) - # find the keypoints and descriptors with SIFT - kp1, des1 = sift.detectAndCompute(img_template,None) - kp2, des2 = sift.detectAndCompute(img,None) + # publish screw pose + screw_pose = PoseStamped() + 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 - index_params = dict(algorithm = FLANN_INDEX_KDTREE, trees = 5) - search_params = dict(checks = 50) + pose_pub.publish(screw_pose) - 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) - ts.registerCallback(callback) + pose_pub = rospy.Publisher('/screw_pose', PoseStamped, queue_size=10) + + ts = message_filters.TimeSynchronizer([box_sub, depth_sub, color_info], 1) + ts.registerCallback(box_callback) rospy.spin() \ No newline at end of file diff --git a/src/yolov5_ros/launch/yolov5.launch b/src/yolov5_ros/launch/yolov5.launch index 87ff11a..5cade72 100644 --- a/src/yolov5_ros/launch/yolov5.launch +++ b/src/yolov5_ros/launch/yolov5.launch @@ -16,7 +16,7 @@ - + diff --git a/src/yolov5_ros/src/yolov5/best.pt b/src/yolov5_ros/src/yolov5/best.pt index ffdd929..52d5f43 100644 Binary files a/src/yolov5_ros/src/yolov5/best.pt and b/src/yolov5_ros/src/yolov5/best.pt differ