使用YOLOv5自己训练模型

前言：本文基于我的另一篇文章作为基础，文章戳这里，主要还是实操为主，让大家能快速上手使用。

数据集构建

1.准备工作

数据收集：

• 图片类型数据

• 视频类型数据（使用opencv进行视频抽帧，相当于转为一张一张图片的数据）

  

  具体操作：

  1、建立如下目录

  

  2、PyCharm打开该路径

  

  3、放入视频文件、新建代码文件并配置jupyter，配置jupyter操作上面有说，结果如下图。

  

  4、接下来我们逐步输入代码

  ​ ①导入必要的包

  ​

  ​ ②读取视频文件

  

  ​ ③显示读取的帧

  

  这里可以发现图片颜色不对，这是因为原始图片是RGB格式，而这里imshow显示的是BGR格式，所以需要进行如下转换。

  

  ​ ④完整代码执行

  由于视频的帧数太多，所以设置参数每30帧读取一张图片。

  最后完整的代码如下（根据自己的文件路径），执行完后可以看到左侧的数据集。

  

标注工具：

• labelimg（安装命令：pip install labelimg）

  结果如下

  

  下载好后终端输入labelimg，自动弹出该标注工具，这里如果出现了以下错误，看这篇文章，错误解决方法，设置一个系统环境变量即可。

  

再次运行，成功

LabelImg操作方法：

点击openDir出现的选择框表示选择存图片的那个文件夹，点击Change Save Dir出现的选择框是选择把图片标注结果的目的文件夹。我在这里新建了一个label文件夹用于存放标注结果，在这里有一个小问题，Python3.10安装的labelimg无法识别jpg格式的图片，就算换成png也有别的错误，测试过别的版本python就可以正常识别，所以如果你使用的是Python3.10重新换一个版本吧(我在这里换成了3.9）。

调整保存格式，labelImg打开默认保存是PascalVOC格式，我们需要点击把它改成yolo格式，如下图。

点击save保存，打开AutoSavemode可以自动保存。

此时我们可以正式开始标注图片。

右键点击CreateBox，用鼠标选择区域，并为这块区域命名，标注好后点击左侧nextImage进入下一张图片。

快捷键： a上一张、d下一张、w画框框

标注好后打开label文件夹，可以看到很多txt文件，任意打开一个，第一个表示类别0（在class文件中可以看到类别），后面四个分别表示你画的框的中心点在整体图片中的一个比例以及宽和高对于整体图片的比例。

这里有四个类别。

2.开始训练

分装图片

将images中的图片分为train和val两个文件夹放置（训练集多验证集少），如下图

同样把labels也按相同操作放置一一对应的图片到train和val文件夹中，classes不用放。

关键参数

我们可以打开train.py文件查看参数位置，可以看到data参数默认是coco128.yaml文件

让我们打开coco128.yaml看看，path指定了数据集的位置，train指定的是数据集中训练集所在的位置，val指定的是数据集中测试集所在的位置，test是可选项，下面classes就是类别的定义。

很明显，我们只需要整理好数据集改写对应的yaml文件同时更新train的，那么我们就可以训练了。

修改参数

复制一份coco128.yaml，更名为classroom

并对文件内容进行修改，修改后如下图。

train.py中修改参数

点击运行，等待一会便训练成功

模型训练时遇到的问题

修改虚拟内存的方法：

开始检测

训练完成后我们打开exp3，可以看到以下文件。

best.pt表示训练过程中最好的那个模型，last.pt表示训练过程中最后一轮的模型，labels存放整体数据比例、分布这些内容、results.csv存的就是各项指标的变化（可以用于自己画图做分析）events开头的文件保存了训练过程中各种参数的变化，可以使用tensorboard来查看，命令如下

tensorboard --logdir runs

执行结果：

按住ctrl点击网址即可查看结果。

正式检测：

输入命令，weights指定刚刚训练时最好的模型，source指定我们一开始使用的视频，–view-img参数表示可以实时用可视化界面展示训练过程。

python detect.py --weights runs/train/exp3/weights/best.pt --source datasets/classroom.mp4 --view-img

hell
tensorboard --logdir runs

执行结果：

[外链图片转存中...(img-jcqtmzr8-1700826447637)]

按住ctrl点击网址即可查看结果。

[外链图片转存中...(img-09iWVA6H-1700826447637)]

**正式检测：**

输入命令，weights指定刚刚训练时最好的模型，source指定我们一开始使用的视频，--view-img参数表示可以实时用可视化界面展示训练过程。

```shell
python detect.py --weights runs/train/exp3/weights/best.pt --source datasets/classroom.mp4 --view-img

由于原视频中有我的同学们，这里就不做展示了，最终结果是符合预期的，简单场景基本能满足需求，但是如果一直在动或者动效效果比较大的时候，最终结果会有出入，但整体来说置信度还是可以的。