最近更新于 2022-06-01 18:21

开始先尝试了传统视觉处理的方式进行检测,但是在条件变化后,检测效果就很差。因此决定采用神经网络,基于 YOLOv5 进行训练,考虑到 YOLOv5 是建立在 PyTorch 的基础上运行的,终端运行的时候再配置一个 PyTorch + YOLOv5 + OpenCV 就显得有点麻烦,而 OpenCV 是肯定会用到的,同时 OpenCV 有一个深度神经网络模块,可以加载训练好的模型,那么就能省掉 PyTorch 和 YOLOv5 了,只需要配置一个 OpenCV 配合训练好的模型文件就能跑起来,这样思路就确定下来了。

本文资源文件下载:https://pan.baidu.com/s/1lXIgmu1ZIcFGKZBBVG00sQ?pwd=7nl5

训练

使用的数据集来源于zhangzhengde0225/CDNet,已标注好,可以直接使用。不过要注意里面标注了两种类别的目标,第一种是人行横道,第二种是导向线。

训练方法和我这篇车牌识别-定位博文中所说的一样:https://blog.iyatt.com/?p=5536,因为部署的目标平台是 Jetson Nano 这种小型计算机板,算力较小,因此选用的模型也是最小的 n 模型。另外这个数据集的数据量比较小,不像车牌训练的那个 19 万,所以训练次数我改为了 20。

模型格式转换

完成训练得到的是 YOLOv5 格式的模型文件,可以借助 YOLOv5 里的脚本转换为 ONNX 格式,然后就能给 OpenCV 使用了。

python3 yolov5/export.py --weights yolov5/runs/train/exp/weights/best.pt --include onnx

斑马线检测

程序主体沿用基于 YOLOv5 和 OpenCV DNN 的目标检测https://blog.iyatt.com/?p=5934

测试环境:

Ubuntu 22.04 x86_64
Python 3.10.4
opencv-python 4.5.5.64
jupyter 1.0.0

使用 jupyter 运行 斑马线检测.ipynb,即可预览效果