视频讲解
ROS2下编写package利用orbbec相机进行yolov8实时目标检测
在《ROS2下编写orbbec相机C++ package并Rviz显示》的基础上,继续添加对获取的图像使用YOLO进行目标检测
首先安装YOLO以及相关库
pip3 install ultralytics
使用如下指令测试下yolo安装情况
yolo task=detect mode=predict model=yolov8n.pt source='https://ultralytics.com/images/bus.jpg'
成功会在当前位置下生成runs,如下图为检测加上标签的图片,确认yolo调用成功
接下来使用orbbec发布的图像,进行YOLO实时识别
编写yolo识别package
ros2 pkg create --build-type ament_python yolo_target_detection --dependencies rclpy sensor_msgs cv_bridge
指定依赖项 rclpy(ROS 2 Python 客户端库)、sensor_msgs(用于处理图像消息)和 cv_bridge(用于在 ROS 图像消息和 OpenCV 图像之间进行转换)
在src/yolo_target_detection/yolo_target_detection目录下创建一个 Python 脚本,例如yolo_target_detection.py,并添加以下代码:
import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Image
from cv_bridge import CvBridge
import cv2
from ultralytics import YOLO
import numpy as np
from std_msgs.msg import Header
from sensor_msgs.msg import Image as ROSImageclass YoloTargetDetectionNode(Node):def __init__(self):super().__init__('yolo_target_detection_node')# Initialize the YOLOv8 modelself.model = YOLO("yolov8n.pt") # 选择你训练的模型self.bridge = CvBridge()# Create a subscriber for RGB imageself.image_sub = self.create_subscription(Image,'/rgb_image', # 修改为你订阅的topicself.image_callback,10)# Create a publisher for output image with bounding boxesself.obb_pub = self.create_publisher(ROSImage,'/obb_image',10)def image_callback(self, msg):try:# Convert ROS Image message to OpenCV imagecv_image = self.bridge.imgmsg_to_cv2(msg, "bgr8")except Exception as e:self.get_logger().error(f"Error converting image: {e}")return# Perform object detection using YOLOv8results = self.model(cv_image)# YOLOv8 returns a list of results, each result is a Results objectresult = results[0] # Get the first result (assuming single image inference)# Get bounding boxes, class IDs, and confidencesboxes = result.boxes.xywh.cpu().numpy() # Bounding boxes (x_center, y_center, width, height)confidences = result.boxes.conf.cpu().numpy() # Confidence scoresclass_ids = result.boxes.cls.cpu().numpy() # Class IDslabels = result.names # Class names# Draw bounding boxes on the imagefor i, box in enumerate(boxes):x_center, y_center, width, height = box[:4]confidence = confidences[i]class_id = int(class_ids[i]) # Get the class IDlabel = labels[class_id] # Get the class label# Convert center to top-left coordinates for cv2x1 = int((x_center - width / 2))y1 = int((y_center - height / 2))x2 = int((x_center + width / 2))y2 = int((y_center + height / 2))# Draw bounding box and label on the imagecv2.rectangle(cv_image, (x1, y1), (x2, y2), (0, 255, 0), 2)cv2.putText(cv_image, f"{label} {confidence:.2f}", (x1, y1 - 10),cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 0, 0), 2)# Convert the image with bounding boxes back to ROS messagetry:obb_image_msg = self.bridge.cv2_to_imgmsg(cv_image, encoding="bgr8")obb_image_msg.header = Header()obb_image_msg.header.stamp = self.get_clock().now().to_msg()obb_image_msg.header.frame_id = "camera_frame" # 根据你的相机frame进行调整# Publish the image with bounding boxesself.obb_pub.publish(obb_image_msg)self.get_logger().info("Published object-bound box image.")except Exception as e:self.get_logger().error(f"Error publishing image: {e}")def main(args=None):rclpy.init(args=args)node = YoloTargetDetectionNode()try:rclpy.spin(node)except KeyboardInterrupt:passfinally:node.destroy_node()rclpy.shutdown()if __name__ == '__main__':main()
打开src/yolo_target_detection/setup.py文件,添加以下内容:
[develop]
script_dir=$base/lib/yolo_target_detection
[install]
install_scripts=$base/lib/yolo_target_detection
在终端中执行以下命令构建和安装包:
colcon build --packages-select yolo_target_detection
source install/setup.bash
ros2 run yolo_target_detection yolo_target_detection
打开Rviz及Orbbec节点发布rgb_image消息即可,同时配置Rviz增加新的image显示,订阅消息为obb_image