ROS节点检测ArUco标签并记录空间坐标

#!/usr/bin/env python
# -*- coding: utf-8 -*-
import rospy
import csv
from sensor_msgs.msg import Image, CameraInfo
from geometry_msgs.msg import PointStamped
from cv_bridge import CvBridge
import cv2
import os
from datetime import datetime
aruco_pub = None
camera_info = None
depth_image_msg = None
csv_data = []
def detect_aruco(image):
    global csv_data
    aruco_dict = cv2.aruco.Dictionary_get(cv2.aruco.DICT_ARUCO_ORIGINAL)
    cv_image = CvBridge().imgmsg_to_cv2(image, "bgr8")
    gray = cv2.cvtColor(cv_image, cv2.COLOR_BGR2GRAY)
    parameters = cv2.aruco.DetectorParameters_create()
    corners, ids, _ = cv2.aruco.detectMarkers(gray, aruco_dict, parameters=parameters)
    if ids is None:
        rospy.loginfo("No ArUco marker found in the image.")
    else:
        for i in range(len(ids)):
            aruco_id = ids[i][0]
            aruco_center_x = int((corners[i][0][0][0] + corners[i][0][1][0] + corners[i][0][2][0] + corners[i][0][3][0]) / 4)
            aruco_center_y = int((corners[i][0][0][1] + corners[i][0][1][1] + corners[i][0][2][1] + corners[i][0][3][1]) / 4)
            try:
                depth_image_msg = rospy.wait_for_message("/camera/depth/image_rect_raw", Image, timeout=1.0)
            except rospy.ROSException:
                rospy.logwarn("Failed to get depth image.")
                return
            depth_image = CvBridge().imgmsg_to_cv2(depth_image_msg, "passthrough")
            aruco_center_x = min(max(0, aruco_center_x), depth_image.shape[1] - 1)
            aruco_center_y = min(max(0, aruco_center_y), depth_image.shape[0] - 1)
            depth = depth_image[aruco_center_y, aruco_center_x]
            if depth != 0:
                if camera_info is not None:
                    K = camera_info.K
                    Z = depth / 1000.0
                    X = (aruco_center_x - K[2]) * Z / K[0]
                    Y = (aruco_center_y - K[5]) * Z / K[4]
                    rospy.loginfo("Detected ArUco marker with ID: {}. 3D Coordinates: ({}, {}, {})".format(aruco_id, X, Y, Z))
                    # 记录数据到csv_data列表
                    csv_data.append([rospy.Time.now().to_sec(), X, Y, Z])
def image_callback(data):
    global depth_image_msg
    detect_aruco(data)
def camera_info_callback(data):
    global camera_info
    camera_info = data
def main():
    global aruco_pub, csv_data
    rospy.init_node('aruco_detection_node', anonymous=True)
    image_sub = rospy.Subscriber("/camera/color/image_raw", Image, image_callback)
    camera_info_sub = rospy.Subscriber("/camera/color/camera_info", CameraInfo, camera_info_callback)
    # 不再发布3D坐标信息到话题/aruco_info，因此这里不需要创建Publisher了
    rospy.spin()
    # 保存数据到csv文件
    if csv_data:
        folder_path = "aruco_tracks"
        if not os.path.exists(folder_path):
            os.makedirs(folder_path)
        file_name = datetime.now().strftime("%Y%m%d_%H%M%S") + ".csv"
        file_path = os.path.join(folder_path, file_name)
        with open(file_path, mode='w') as file:
            writer = csv.writer(file)
            writer.writerow(['Time', 'X', 'Y', 'Z'])
            writer.writerows(csv_data)
if __name__ == "__main__":
    try:
        main()
    except rospy.ROSInterruptException:
        pass