基于 Yolov5 实现目标检测

目录

第一部分：代码准备

第二部分：基于Yolov5_Pytorch训练自己的目标检测模型

2.1 数据准备

2.2 模型训练

2.3 模型测试

---

第一部分：代码准备

Yolov5代码克隆：ultralytics/yolov5: YOLOv5 🚀 in PyTorch > ONNX > CoreML > TFLite (github.com)

 目前代码版本已经更新到了v7.0，本文使用的是v5.0的Yolov5代码版本。

第二部分：基于Yolov5_Pytorch训练自己的目标检测模型

2.1 数据准备

这里使用labelimg进行数据集的目标标注。labelimg是一款开源的图像标注工具，标签可用于分类和目标检测，并使用Qt作为其图形界面，简单易用。其注释以 PASCAL VOC格式保存为XML文件。

cmd里输入以下命令，将自行安装labelimg工具：

pip install labelimg -i https://pypi.tuna.tsinghua.edu.cn/simple

下载完成后运行labelimg：

 打开labelimg后，点击Open Dir找到自己数据集的所在目录，并根据所需进行逐个手动标注，生成所需xml格式文件。注：标注时可以用键盘上的A返回上一张，D键翻到下一张，提高效率。

 在这里可以开启自动保存。

 将数据集按照此格式划分。标签与图片区分，并按照一定比例划分训练集与验证集。

在这里下载Yolov5的预训练权重：

Releases · ultralytics/yolov5 · GitHub

 上图展示不同权重在精度与速度上的区分。本文选择的Yolov5s.pt作为预训练权重。

2.2 模型训练

在训练之前，这里需要修改两个配置文件。

首先是数据设定文件，在 data 目录里找到 voc.yaml 文件，这里建议复制一份并重命名为与自己项目相关再进行修改。 

 具体而言要进行以下4项预设的修改。将箭头指向的内容替换为自己的项目所需。

 然后需要修改模型配置文件，在models目录里找到自己选择的预训练权重的yaml文件修改。本文选择的 yolov5s.yaml。当然这里也建议复制一份并根据所需重命名再修改。

具体地，这里只需做一处修改，即将nc改为自己的类别数。

 

最后在目录里找到 train.py 文件，主要对以下训练参数进行设置，然后就可以开始训练了。其他内容根据所需修改。

运行 train.py 开始训练。

2.3 模型测试

训练结束后产生一个 runs 目录，这里面可以找到自己训练的权重文件。默认会存储精度最佳的权重best.pt与最后一个世代训练的权重文件last.pt。

在目录里找到 detect.py 文件，根据所需修改权重路径、测试数据路径以及其他相关参数。对训练的权重进行测试。

测试结果存放在 runs/detect/exp 路径中。