深度学习目标检测：YOLOv5实现红绿灯检测(含红绿灯数据集+训练代码)

2022-12-17 11:07 648人阅读评论(0)

深度学习目标检测：YOLOv5实现红绿灯检测(含红绿灯数据集+训练代码)

1. 前言

本篇博客，我们将手把手教你搭建一个基于YOLOv5的红绿灯目标检测项目。目前，基于YOLOv5s的红绿灯检测精度平均值mAP_0.5=0.93919，mAP_0.5:0.95=0.63967，基本满足业务的性能需求。另外，为了能部署在手机Android平台上，本人对YOLOv5s进行了模型轻量化，开发了一个轻量级的版本yolov5s05_416和yolov5s05_320，在普通Android手机上可以达到实时的检测和识别效果，CPU(4线程)约30ms左右，GPU约25ms左右，基本满足业务的性能需求。

先展示一下Python版本红绿灯检测Demo效果：

【整套项目下载地址】：YOLOv5实现红绿灯检测(含红绿灯数据集+训练代码)

【尊重原创，转载请注明出处】https://blog.csdn.net/guyuealian/article/details/128240198

2. 红绿灯检测数据集说明

（1）红绿灯检测数据集

目前收集了约1W+的红绿灯（交通信号灯）检测数据集：Traffic-Lights-Dataset-Domestic+Traffic-Lights-Dataset-Foreign:

关于红绿灯检测数据集使用说明和下载，详见另一篇博客说明：《红绿灯(交通信号灯)检测数据集》

（2）自定义数据集

如果需要增/删类别数据进行训练，或者需要自定数据集进行训练，可参考如下步骤：

采集图片，建议不少于200张图片

使用Labelme等标注工具，对目标进行拉框标注：labelme工具：GitHub - wkentaro/labelme: Image Polygonal Annotation with Python (polygon, rectangle, circle, line, point and image-level flag annotation).

将标注格式转换为VOC数据格式，参考工具：labelme/labelme2voc.py at main · wkentaro/labelme · GitHub

生成训练集train.txt和验证集val.txt文件列表

修改engine/configs/voc_local.yaml的train和val的数据路径

重新开始训练

3. 基于YOLOv5的红绿灯检测模型训练

（1）YOLOv5说明

训练Pipeline采用YOLOv5: https://github.com/ultralytics/yolov5 , 原始官方代码训练需要转换为YOLO的格式，不支持VOC的数据格式。为了适配VOC数据，本人新增了LoadVOCImagesAndLabels用于解析VOC数据集，以便正常训练。另外，为了方便测试，还增加demo.py文件，可支持对图片,视频和摄像头的测试。

Python依赖环境，使用pip安装即可，项目代码都在Ubuntu系统和Windows系统验证正常运行，请放心使用；若出现异常，大概率是相关依赖包版本没有完全对应


   
    
     
      
     
     
       
      
     
    
     
      
     
     
      
       matplotlib>=
       3.2
       .2
      
     
    
     
      
     
     
      
       numpy>=
       1.18
       .5
      
     
    
     
      
     
     
      
       opencv-python>=
       4.1
       .2
      
     
    
     
      
     
     
      
       Pillow
      
     
    
     
      
     
     
      
       PyYAML>=
       5.3
       .1
      
     
    
     
      
     
     
      
       scipy>=
       1.4
       .1
      
     
    
     
      
     
     
      
       torch>=
       1.7
       .0
      
     
    
     
      
     
     
      
       torchvision>=
       0.8
       .1
      
     
    
     
      
     
     
      
       tqdm>=
       4.41
       .0
      
     
    
     
      
     
     
      
       tensorboard>=
       2.4
       .1
      
     
    
     
      
     
     
      
       seaborn>=
       0.11
       .0
      
     
    
     
      
     
     
      
       pandas
      
     
    
     
      
     
     
      
       thop  
       # FLOPs computation
      
     
    
     
      
     
     
      
       pybaseutils==
       0.6
       .5

（2）准备Train和Test数据

下载红绿灯检测数据集，总共约1W+的图片：Traffic-Lights-Dataset-Domestic+Traffic-Lights-Dataset-Foreign；

考虑到这两个数据集的标注格式一样，因此本项目打算将两个数据集一起合并训练，这样可以提高模型红绿灯检测效果。如果你需要增加自己的数据集进行训练，请根据自己环境路径，适当修改即可。

关于红绿灯检测数据集使用说明和下载，详见另一篇博客说明：《红绿灯(交通信号灯)检测数据集》

（3）配置数据文件

修改训练和测试数据的路径：engine/configs/voc_local.yaml


   
    
     
      
     
     
      
       # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
      
     
    
     
      
     
     
      
       # 数据路径
      
     
    
     
      
     
     
      
       path: 
       ""  
       # dataset root dir
      
     
    
     
      
     
     
      
       train:
      
     
    
     
      
     
     
      
         - 
       "path/to/traffic light/Traffic-Lights-Dataset-Domestic/train.txt"
      
     
    
     
      
     
     
      
         - 
       "path/to/traffic light/Traffic-Lights-Dataset-Foreign/train.txt"
      
     
    
     
      
     
     
       
      
     
    
     
      
     
     
      
       val:
      
     
    
     
      
     
     
      
         - 
       "path/to/traffic light/Traffic-Lights-Dataset-Domestic/val.txt"
      
     
    
     
      
     
     
       
      
     
    
     
      
     
     
      
       test:  
       # test images (optional)
      
     
    
     
      
     
     
      
       data_type: voc
      
     
    
     
      
     
     
       
      
     
    
     
      
     
     
      
       # 1.设置类别个数，和要训练的类别名称，ID号从0开始递增
      
     
    
     
      
     
     
      
       nc: 4  
       # number of classes
      
     
    
     
      
     
     
      
       names: { 
       'red': 0, 
       'green': 1,  
       'yellow': 2,
       'off': 3 }
      
     
    
     
      
     
     
      
       # 2.如果你想合并所有类别为一个大类，进行训练： unique表示合并所有类为单独一个类别
      
     
    
     
      
     
     
      
       #nc: 1 # number of classes
      
     
    
     
      
     
     
      
       #names: { "unique": 0 }

（4）配置模型文件

官方YOLOv5给出了YOLOv5l,YOLOv5m,YOLOv5s等模型。考虑到手机端CPU/GPU性能比较弱鸡，直接部署yolov5s运行速度十分慢。所以本人在yolov5s基础上进行模型轻量化处理，即将yolov5s的模型的channels通道数全部都减少一半，并且模型输入由原来的640×640降低到416×416或者320×320，该轻量化的模型我称之为yolov5s05。从性能来看，yolov5s05比yolov5s快5多倍，而mAP下降了16%（0.93→0.77），对于手机端，这精度勉强可以接受。

下面是yolov5s05和yolov5s的参数量和计算量对比：

模型	input-size	params(M)	GFLOPs
yolov5s	640×640	7.2	16.5
yolov5s05	416×416	1.7	1.8
yolov5s05	320×320	1.7	1.1

（5）重新聚类Anchor（可选）

官方yolov5s的Anchor是基于COCO数据集进行聚类获得（详见models/yolov5s.yaml文件）

考虑到红绿灯检测数据集，目标相对较小，并且目标框几乎都是竖状的矩形框；直接复用原始COCO的Anchor效果可能不太好；因此，这里对红绿灯数据集的标注框进行重新聚类获得新的Anchor；下表给出yolov5s，yolov5s05_416和yolov5s05_320重新聚类的Anchor结果：

yolov5s.yaml	yolov5s05_416.yaml	yolov5s05_320.yaml

一点建议：

官方yolov5s的Anchor是基于COCO数据集进行聚类获得，不同数据集需要做适当的调整，其最优Anchor建议重新进行聚类。

当然你要是觉得麻烦就跳过，不需要重新聚类Anchor，这个影响不是特别大。如果你需要重新聚类，请参考engine/kmeans_anchor/demo.py文件

（6）开始训练

整套训练代码非常简单操作，用户只需要将相同类别的数据放在同一个目录下，并填写好对应的数据路径，即可开始训练了。

修改训练超参文件： data/hyps/hyp.scratch-v1.yaml （可以修改训练学习率，数据增强等方式，使用默认即可）

编辑train.sh脚本训练，训练yolov5s或轻量化版本yolov5s05_416或者yolov5s05_320 (选择其中一个训练即可)


   
    
     
      
     
     
      
       #!/usr/bin/env bash
      
     
    
     
      
     
     
       
      
     
    
     
      
     
     
      
       #--------------训练yolov5s--------------
      
     
    
     
      
     
     
      
       # 输出项目名称路径
      
     
    
     
      
     
     
      
       project=
       "runs/yolov5s_640"
      
     
    
     
      
     
     
      
       # 训练和测试数据的路径
      
     
    
     
      
     
     
      
       data=
       "engine/configs/voc_local.yaml"
      
     
    
     
      
     
     
      
       # YOLOv5模型配置文件
      
     
    
     
      
     
     
      
       cfg=
       "yolov5s.yaml"
      
     
    
     
      
     
     
      
       # 训练超参数文件
      
     
    
     
      
     
     
      
       hyp=
       "data/hyps/hyp.scratch-v1.yaml"
      
     
    
     
      
     
     
      
       # 预训练文件
      
     
    
     
      
     
     
      
       weights=
       "engine/pretrained/yolov5s.pt"
      
     
    
     
      
     
     
      
       python train.py --data 
       $data --cfg 
       $cfg --hyp 
       $hyp --weights 
       $weights --batch-size 16 --imgsz 640 --workers 12 --project 
       $project
      
     
    
     
      
     
     
       
      
     
    
     
      
     
     
       
      
     
    
     
      
     
     
      
       #--------------训练轻量化版本yolov5s05_416--------------
      
     
    
     
      
     
     
      
       # 输出项目名称路径
      
     
    
     
      
     
     
      
       project=
       "runs/yolov5s05_416"
      
     
    
     
      
     
     
      
       # 训练和测试数据的路径
      
     
    
     
      
     
     
      
       data=
       "engine/configs/voc_local.yaml"
      
     
    
     
      
     
     
      
       # YOLOv5模型配置文件
      
     
    
     
      
     
     
      
       cfg=
       "yolov5s05_416.yaml"
      
     
    
     
      
     
     
      
       # 训练超参数文件
      
     
    
     
      
     
     
      
       hyp=
       "data/hyps/hyp.scratch-v1.yaml"
      
     
    
     
      
     
     
      
       # 预训练文件
      
     
    
     
      
     
     
      
       weights=
       "engine/pretrained/yolov5s.pt"
      
     
    
     
      
     
     
      
       python train.py --data 
       $data --cfg 
       $cfg --hyp 
       $hyp --weights 
       $weights --batch-size 64 --imgsz 416 --workers 12 --project 
       $project
      
     
    
     
      
     
     
       
      
     
    
     
      
     
     
       
      
     
    
     
      
     
     
      
       #--------------训练轻量化版本yolov5s05_320--------------
      
     
    
     
      
     
     
      
       # 输出项目名称路径
      
     
    
     
      
     
     
      
       project=
       "runs/yolov5s05_320"
      
     
    
     
      
     
     
      
       # 训练和测试数据的路径
      
     
    
     
      
     
     
      
       data=
       "engine/configs/voc_local.yaml"
      
     
    
     
      
     
     
      
       # YOLOv5模型配置文件
      
     
    
     
      
     
     
      
       cfg=
       "yolov5s05_320.yaml"
      
     
    
     
      
     
     
      
       # 训练超参数文件
      
     
    
     
      
     
     
      
       hyp=
       "data/hyps/hyp.scratch-v1.yaml"
      
     
    
     
      
     
     
      
       # 预训练文件
      
     
    
     
      
     
     
      
       weights=
       "engine/pretrained/yolov5s.pt"
      
     
    
     
      
     
     
      
       python train.py --data 
       $data --cfg 
       $cfg --hyp 
       $hyp --weights 
       $weights --batch-size 64 --imgsz 320 --workers 12 --project 
       $project

开始训练： bash train.sh

训练数据量比较大，训练时间比较长，请耐心等待哈
训练完成后，在模型输出目录中有个results.csv文件，记录每个epoch测试的结果，如loss,mAP等信息

训练模型收敛后，yolov5s红绿灯检测的mAP指标大约mAP_0.5=0.93919；而，yolov5s05_416 mAP_0.5=0.77174左右；yolov5s05_320 mAP_0.5=0.71944左右；相比而言，yolov5s05比yolov5s差了很多,mAP掉了近20个点，这主要是因为减小输入input-size，不利于小目标检测。

模型	input-size	params(M)	GFLOPs	mAP0.5
yolov5s	640×640	7.2	16.5	0.93919
yolov5s05	416×416	1.7	1.8	0.77174
yolov5s05	320×320	1.7	1.1	0.71944

（7）可视化训练过程

训练过程可视化工具是使用Tensorboard，使用方法：


   
    
     
      
     
     
      
       # 基本方法
      
     
    
     
      
     
     
      
       tensorboard --logdir=path/to/log/
      
     
    
     
      
     
     
      
       # 例如
      
     
    
     
      
     
     
      
       tensorboard --logdir ./data/model/yolov5s_640

当然，在输出目录，也保存很多性能指标的图片

这是训练epoch的可视化图，可以看到mAP随着Epoch训练，逐渐提高

这是每个类别的F1-Score分数

这是模型的PR曲线

这是混淆矩阵：

4. Python版本红绿灯检测效果

demo.py文件用于推理和测试模型的效果，填写好配置文件，模型文件以及测试图片即可运行测试了

测试图片


   
    
     
      
     
     
      
       # 测试图片
      
     
    
     
      
     
     
      
       image_dir=
       'data/test_image' 
       # 测试图片的目录
      
     
    
     
      
     
     
      
       weights=
       "data/model/yolov5s_640/weights/best.pt" 
       # 模型文件
      
     
    
     
      
     
     
      
       out_dir=
       "runs/test-result" 
       # 保存检测结果
      
     
    
     
      
     
     
      
       python demo.py --image_dir 
       $image_dir --weights 
       $weights --out_dir 
       $out_dir

测试视频文件


   
    
     
      
     
     
      
       # 测试视频文件
      
     
    
     
      
     
     
      
       video_file=
       "data/test-video.mp4" 
       # path/to/video.mp4 测试视频文件，如*.mp4,*.avi等
      
     
    
     
      
     
     
      
       weights=
       "data/model/yolov5s_640/weights/best.pt" 
       # 模型文件
      
     
    
     
      
     
     
      
       out_dir=
       "runs/test-result" 
       # 保存检测结果
      
     
    
     
      
     
     
      
       python demo.py --video_file 
       $video_file --weights 
       $weights --out_dir 
       $out_dir

测试摄像头


   
    
     
      
     
     
      
       # 测试摄像头
      
     
    
     
      
     
     
      
       video_file=0 
       # 测试摄像头ID
      
     
    
     
      
     
     
      
       weights=
       "data/model/yolov5s_640/weights/best.pt" 
       # 模型文件
      
     
    
     
      
     
     
      
       out_dir=
       "runs/car-result" 
       # 保存检测结果
      
     
    
     
      
     
     
      
       python demo.py --video_file 
       $video_file --weights 
       $weights --out_dir 
       $out_dir

如果想进一步提高模型的性能，可以尝试：

增加训练的样本数据：目前只有1W+的数据量，建议根据自己的业务场景，采集相关数据，提高模型泛化能力

使用参数量更大的模型：本教程使用的YOLOv5s，其参数量才7.2M，而YOLOv5x的参数量有86.7M，理论上其精度更高，但推理速度也较慢。

红绿灯目标检测属于小目标检测，可以尝试减小模型下采样次数，增大input-size，来提高小目标检测效果

尝试不同数据增强的组合进行训练

5. Android版本红绿灯检测效果

已经完成Android版本红绿灯检测模型算法开发，APP在普通Android手机上可以达到实时的检测和识别效果，CPU(4线程)约30ms左右，GPU约25ms左右，基本满足业务的性能需求。详细说明请查看《深度学习目标检测：Android实现红绿灯检测(含Android源码可实时运行)》https://blog.csdn.net/guyuealian/article/details/128240334

Android Demo体验：Android实现红绿灯检测APP(可实时运行))

6.项目源码下载

整套项目源码内容包含：红绿灯检测数据集 + YOLOv5训练代码和测试代码

整套项目下载地址：YOLOv5实现红绿灯检测(含红绿灯数据集+训练代码)

（1）红绿灯检测数据集：红绿灯(交通信号灯)检测数据集

Traffic-Lights-Dataset-Domestic红绿灯检测数据集

Traffic-Lights-Dataset-Foreign红绿灯检测数据集

（2）YOLOv5训练代码和测试代码（Pytorch）

整套YOLOv5项目工程的训练代码和测试代码

支持高精度版本yolov5s训练和测试

支持轻量化版本yolov5s05_320和yolov5s05_416训练和测试

根据本篇博文说明，简单配置即可开始训练

转载：https://blog.csdn.net/guyuealian/article/details/128240198

查看评论

小言_互联网的博客

小言_互联网的博客

个人资料

文章分类

文章存档

阅读排行

评论排行

推荐文章