做个 ROS 2 视觉检测开源库-加载模型并进行检测

大家好，我是小鱼。

接着上一节（做个 ROS 2 视觉检测开源库-功能包框架搭建）继续完善模型工具加载和话题订阅发布功能，完善后的代码如下：在 __init__方法中创建订阅和发布者

...
from ament_index_python.packages import get_package_share_directory
from vision_msgs.msg import Detection2DArray, ObjectHypothesisWithPose, Detection2D
from sensor_msgs.msg import Image
from cv_bridge import CvBridge
import cv2
import yolov5


class YOLOv5Ros2(Node):
    def __init__(self):
        ...
        # 加载模型和工具
        if self.model_path == "":
            package_share_directory = get_package_share_directory('yolov5_ros2')
            self.model_path = package_share_directory + "/config/yolov5n.pt"
        self.yolov5_model = yolov5.load(self.model_path, self.device)
        self.bridge = CvBridge()
        # 创建订阅发布者
        self.yolo_result_pub = self.create_publisher(
            Detection2DArray, "yolo_result", 10)
        if self.pub_result_image:
            self.result_img_pub = self.create_publisher(
                Image, "result_image", 10)
        self.image_sub = self.create_subscription(
            Image, self.image_topic, self.image_callback, 10)

上面的代码中，首先导入相关库和接口类，其中 CvBridge 用于将图像消息转换成 OpenCV 格式。在 init 函数中，首先判断模型文件路径是否为空，若空则加载功能包内部的模型。然后创建了图像话题订阅者和结果的发布者实例。有了订阅者就可以在回调函数中拿到图像数据，接着来完善回调函数：

def image_callback(self, msg):
    """
    处理输入图像的回调函数


    Args:
        msg: 输入图像消息
    """
    image = self.bridge.imgmsg_to_cv2(msg,desired_encoding="rgb8") # 转RBG格式图像
    result = self.yolov5_model(image) # 进行检测
    self.get_logger().info(str(result))
    self.pub_result(result, msg.header) # 发布结果
    if self.pub_result_image:   # 根据配置决定是否发布结果图像
        self.pub_result_with_image(result, image, msg.header)

在回调函数 image_callback 中，首先将收到的消息转换成 OpenCV 格式，接着调用模型进行检测并打印检测结果，最后调用 pub_result 发布结果，是否发布绘制好的结果图像则要根据参数 pub_result_image 决定。需要注意图像发布时传入 msg.header 是为了让发布数据的 header 和原始图像保持一致，方便用户根据 header 筛选出结果对应的原始图像。完成模型加载和检测后，接着我们来编写两个发布函数。