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 ROSImage


class YoloTargetDetectionNode(Node):
    def __init__(self):
        super().__init__('yolo_target_detection_node')
        
        # Initialize the YOLOv8 model
        self.model = YOLO("yolov8n.pt")  # 选择你训练的模型
        self.bridge = CvBridge()

        # Create a subscriber for RGB image
        self.image_sub = self.create_subscription(
            Image,
            '/rgb_image',  # 修改为你订阅的topic
            self.image_callback,
            10
        )

        # Create a publisher for output image with bounding boxes
        self.obb_pub = self.create_publisher(
            ROSImage,
            '/obb_image',
            10
        )

    def image_callback(self, msg):
        try:
            # Convert ROS Image message to OpenCV image
            cv_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 YOLOv8
        results = self.model(cv_image)
        
        # YOLOv8 returns a list of results, each result is a Results object
        result = results[0]  # Get the first result (assuming single image inference)
        
        # Get bounding boxes, class IDs, and confidences
        boxes = result.boxes.xywh.cpu().numpy()  # Bounding boxes (x_center, y_center, width, height)
        confidences = result.boxes.conf.cpu().numpy()  # Confidence scores
        class_ids = result.boxes.cls.cpu().numpy()  # Class IDs
        labels = result.names  # Class names

        # Draw bounding boxes on the image
        for 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 ID
            label = labels[class_id]  # Get the class label
            
            # Convert center to top-left coordinates for cv2
            x1 = 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 image
            cv2.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 message
        try:
            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 boxes
            self.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:
        pass
    finally:
        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