yolov5+ros 识别交通信号灯实战第三节：红绿灯实时检测

操作系统：Ubuntu20.04
ROS版本：neotic
CUDA版本：12.2

闲来无事想用 yolov5 实现交通信号灯的检测，记录下过程，效果如下：

1.相机准备

1.1.相机驱动安装

相机我用是 Intel realsense D435i，使用普通的USB免驱摄像头也可以。
从本节开始，比如打开终端、切换目录、切换conda环境等基本操作的说明将简化
第一步：注册服务器公钥

sudo apt-get update && sudo apt-get upgrade && sudo apt-get dist-upgrade
sudo apt-key adv --keyserver keyserver.ubuntu.com --recv-key F6E65AC044F831AC80A06380C8B3A55A6F3EFCDE || sudo apt-key adv --keyserver hkp://keyserver.ubuntu.com:80 --recv-key F6E65AC044F831AC80A06380C8B3A55A6F3EFCDE

第二步：将服务器添加至存储库列表

sudo add-apt-repository "deb https://librealsense.intel.com/Debian/apt-repo $(lsb_release -cs) main" -u

第三步：相机SDK安装

# 必须安装
sudo apt-get install librealsense2-dkms
sudo apt-get install librealsense2-utils
# 可选安装
sudo apt-get install librealsense2-dev
sudo apt-get install librealsense2-dbg

驱动安装正确，终端输入下面的指令会打开SDK自带测试工具，如下图所示：

realsense-viewer

这里可以修改RGB摄像头的分辨率、帧率、编码格式等等，可自行研究。先保持默认配置，点击RGB Camera 后面的开关，打开RGB相机，正确打开后的效果如下。打开没效果看看界面右上角的2D | 3D 选项，需选择2D才能显示

1.2.ROS Wrapper安装

首先创建ROS工作空间，初始化ROS空间，编译，使用的命令如下：

mkdir -p catkin_ws/src
cd catkin_ws/src
catkin_init_workspace
cd .. && catkin_make

执行最后一条命令可能会报错，因为我们使用了 Anaconda ，catkin 命令找到使用了 Anaconda下的python，所示使用如下命令进行指定 python

catkin_make -DPYTHON_EXECUTABLE=/usr/bin/python3

编译成功后，下载两个源码包：

# 进入到src目录下载源码
cd ~/catkin_ws/src
git clone https://github.com/IntelRealSense/realsense-ros.git
git clone https://github.com/pal-robotics/ddynamic_reconfigure.git
# realsense_ros现在默认是ros2，我们使用的是ros neotic，所以需要签出正确的版本
cd realsense-ros
git checkout 2.3.2
# 编译
cd ~/catkin_ws && catkin_make

编译完成后，打开 catkin_ws/src/realsense-ros/realsense2_camera/launch 目录下的rs_camera.launch 文件，找到如下红色框的三个参数，并按照图片进行修改保存

输入如下命令，即可启动ROS Wrapper，彩色图发布的 Topic 名称为：/camera/color/image_raw

source devel/setup.bash
roslaunch realsense2_camera rs_camera.launch

启动好后，新开一个终端可通过 rivz 工具查看，如下图所示：

到这里，相机准备就绪。

2.创建ROS包

到这里为止，yolov5程序可以对单张图片进行红绿灯检测，相机可以实时获取图片。那么我们只需要写一个ROS包完成如下功能即可：

• 订阅相机发布彩图图片的话题
• 获取图片，调用yolov5程序
• 定义消息体，用于获取yolov5的检测结果
• 根据结果绘制检测框并发布在ROS

下面开始创建ROS包，先创建消息体

cd catkin_ws/src/
mkdir yolov5ros && cd yolov5ros
catkin_create_pkg yolov5msgs std_msgs

在yolov5msgs中创建两个消息体：Box.msg 和 Boxes.msg
Box.msg内容如下：

float64 probability
int64 xmin
int64 ymin
int64 xmax
int64 ymax
int16 num
string Class

Boxes.msg内容如下：

Header header
Header imageHeader
Box[] boxes

yolov5msgs 包的CMakeLists.txt 内容如下：

cmake_minimum_required(VERSION 3.0.2)
project(yolov5msgs)

find_package(catkin REQUIRED COMPONENTS std_msgs message_generation)
add_message_files(
        DIRECTORY msg
        FILES
        Box.msg
        Boxes.msg
)
generate_messages(
        DEPENDENCIES
        std_msgs
)
catkin_package()
include_directories(${catkin_INCLUDE_DIRS})

编译，应该没有问题！

接下来创建ROS包

cd catkin_ws/src/yolov5ros/
# 创建包
catkin_create_pkg yolov5ros rospy std_msgs
cd yolov5ros/
# 创建文件夹，分别存放：ROS包启动文件，主程序，yolov5权重文件
mkdir launch
mkdir scripts
mkdir weights

然后将我们前面用的yolov5完整程序包，复制到 catkin_ws/src/yolov5ros/yolov5ros 目录下，将上一节训练好的best.pt 文件放在ros包的 weights 目录下，完整的目录结构应该如下图所示：

在scripts目录下创建文件detect.py，文件内容如下：

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import cv2
import torch
import rospy
import numpy as np
from sensor_msgs.msg import Image
from std_msgs.msg import Header
from yolov5msgs.msg import Box, Boxes

class YoloDetect:
    def __init__(self):
        # yolov5程序路径
        yolov5Path = rospy.get_param('~yolov5Path', '')
        # 模型文件路径
        weightPath = rospy.get_param('~weightPath', '')
        # 图片发布在那个话题
        imageTopic = rospy.get_param('~imageTopic', '/camera/color/image_raw')
        self.model = torch.hub.load(yolov5Path, 'custom', path=weightPath, source='local')
        self.model.cuda()
        
        self.colorImage = Image()
        self.receiveImage = False
        self.colorSub = rospy.Subscriber(imageTopic, Image, self.imageCB,queue_size=1, buff_size=52428800)
        while(not self.receiveImage):
            rospy.loginfo('not receive image, waiting.')
            rospy.sleep(2)
        
    def imageCB(self, image):
        self.boxes = Boxes()
        self.boxes.header = image.header
        self.boxes.imageHeader = image.header
        self.receiveImage = True
        self.colorImage = np.frombuffer(image.data, dtype=np.uint8).reshape(image.height, image.width, -1)
        result = self.model(self.colorImage)
        detectBox = result.pandas().xyxy[0].values
        self.detectShow(self.colorImage, detectBox, image.height, image.width)
        cv2.waitKey(2)
        
    def detectShow(self, originImage, detectBoxs, height, width):
        # 复制图像
        image = originImage.copy()
        # 统计检测总数
        detectNum = 0
        for i in detectBoxs:
            detectNum += 1
        # 遍历所有的检测
        for item in detectBoxs:
            # 构造对象
            box = Box()
            # 获取属性
            box.xmin = np.int64(item[0])
            box.ymin = np.int64(item[1])
            box.xmax = np.int64(item[2])
            box.ymax = np.int64(item[3])
            box.num = np.int16(detectNum)
            box.Class = item[-1]
            box.probability =np.float64(item[4])
            # 定义检测框颜色
            color = np.random.randint(0, 100, 3)
            # 绘制检测框
            cv2.rectangle(image, (int(item[0]), int(item[1])),(int(item[2]), int(item[3])), 
                          (int(color[0]),int(color[1]), int(color[2])), 2)
            # 绘制标签
            if item[1] < 20:
                y = item[1] + 30
            else:
                y = item[1] - 10
            cv2.putText(image, item[-1],(int(item[0]), int(y)-10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 0), 2, cv2.LINE_AA)
            self.boxes.boxes.append(box)
        # 图像转换并显示
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        cv2.imshow('detect result', image)
            

def start():
    rospy.init_node('yolov5ros', anonymous=True)
    detect = YoloDetect()
    rospy.spin()

if __name__ == "__main__":
    start()

在launch目录下创建start.launch文件，内容如下：

<?xml version="1.0" encoding="utf-8"?>
<launch>   
  <node pkg="yolov5ros" type="detect.py" name="yolov5ros" output="screen" >
    <param name="yolov5Path"       value="$(find yolov5ros)/yolov5"/>  
    <param name="weightPath"       value="$(find yolov5ros)/weights/best.pt"/>
    <param name="imageTopic"       value="/camera/color/image_raw" />
  </node>
  <include file="$(find realsense2_camera)/launch/rs_camera.launch" />
</launch>

该包的CMakeLists.txt内容如下：

cmake_minimum_required(VERSION 3.0.2)
project(yolov5ros)
find_package(catkin REQUIRED COMPONENTS rospy std_msgs)
catkin_package()
include_directories(${catkin_INCLUDE_DIRS})

在终端激活 yolov5 环境，然后进行编译

编译过程中可能会出现如下错误，找不到 rospkg

这个是因为安装了 Anaconda 导致的，执行以下命令解决：

pip install catkin-tools rospkg pyyaml empy numpy

然后再次 catkin_make 进行编译，编译成功后执行

roslaunch yolov5ros start.launch

不出意外，你会看到如下界面

可以在电脑上打开训练时的照片，相机对着电脑屏幕验证下效果

关注作者，回复 yolov5 获取完整代码和训练数据集

文章内容仅供学习参考，如有侵权，非常抱歉，请立即联系作者删除。