做个 ROS 2 视觉检测开源库-功能包框架搭建

大家好，我是小鱼。今天继续来学做开源库。

《做个 ROS 2 视觉检测开源库-功能包框架搭建》

新建目录 chapt9/chapt9_ws/src，接着在目录下新建 yolov5_ros2 功能包，并添加相关依赖，完整命令如下：

ros2 pkg create yolov5_ros2 --build-type ament_python --dependencies rclpy yolov5 cv_bridge sensor_msgs vision_msgs cv2 --license Apache-2.0

接着在 chapt9_ws/src/yolov5_ros2/yolov5_ros2 下新建 yolov5_ros2.py ，编写如下代码：

import rclpy
from rclpy.node import Node
from rcl_interfaces.msg import ParameterDescriptor


class YOLOv5Ros2(Node):
    def __init__(self):
        super().__init__('yolov5_ros2')


        # 声明ROS参数，用于配置YOLOv5节点的行为
        self.declare_parameter("device", "cpu", ParameterDescriptor(
            name="device", description="计算设备选择，默认：cpu，可选：cuda:0"))
        self.declare_parameter("model_path", "", ParameterDescriptor(
            name="model_path", description="YOLOv5模型路径，默认为空"))
        self.declare_parameter("image_topic", "/image", ParameterDescriptor(
            name="image_topic", description="输入图像话题，默认：/image_raw"))
        self.declare_parameter("pub_result_image", False, ParameterDescriptor(
            name="pub_result_img", description="是否发布识别结果图像，默认：False"))


        # 获取ROS参数的值
        self.device = self.get_parameter('device').value
        self.model_path = self.get_parameter('model_path').value
        self.image_topic = self.get_parameter('image_topic').value
        self.pub_result_image = self.get_parameter('pub_result_image').value


    def pub_result(self, result, header):
        """
        发布YOLOv5的识别结果


        Args:
            result: YOLOv5的识别结果数据
            header: ROS消息头信息
        """
        # TODO: 实现识别结果的发布逻辑


    def pub_result_with_image(self, result, image, header):
        """
        发布包含识别结果的图像


        Args:
            result: YOLOv5的识别结果数据
            image: 包含识别结果的图像数据
            header: ROS消息头信息
        """
        # TODO: 实现包含识别结果的图像发布逻辑


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


        Args:
            msg: 输入图像消息
        """
        # TODO: 实现输入图像的处理逻辑


def main():
    rclpy.init()
    rclpy.spin(YOLOv5Ros2())
    rclpy.shutdown()


if __name__ == "__main__":
    main()

在上面的代码中，我们创建了一个名为 YOLOv5Ros2 的节点类，并在 init 函数中声明计算设备 device、模型路径 model_path、图像话题 image_topic 和是否发布结果图像 pub_result_image 四个参数。接着又在函数体中定义了两个结果发布函数和一个图像回调函数以备使用。需要考虑当用户并没有提供预训练模型时，我们可以使用一个默认的模型替代，所以我们需要放置一个预训练模型到我们的功能包中，在 chapt9_ws/src/yolov5_ros2/ 下新建 config 目录，然后将 yolov5n.pt 复制到该目录下。接着我们可以修改 setup.py 对该节点进行注册并拷贝模型文件到 install 目录，主要修改代码如下：

from setuptools import find_packages, setup
from glob import glob
import os


package_name = 'yolov5_ros2'


setup(
    ...
    data_files=[
        ('share/ament_index/resource_index/packages',
            ['resource/' + package_name]),
        (os.path.join('share', package_name, 'config'), glob('config/**')),
        ('share/' + package_name, ['package.xml']),
    ],
    ...
    entry_points={
        'console_scripts': [
            "yolov5_ros2=yolov5_ros2.yolov5_ros2:main"
        ],
    },
)

完成上面的工作后，我们的代码框架就算有了，接着就可以继续完善代码了。