1 环境和数据集配置

在上一篇文章中，详细介绍如何从零开始上手YOLOv5，包括环境搭建、模型使用、模型训练等，这个过程按照我写的教程一步步来，会非常简单。

但是从官网中下载过来的代码，直接跑通，训练的是公共数据集，比如COCO、VOC等等。

但是在现实中，我们往往需要训练自己的数据集。比如检测道路上的车辆、行人以及交通标志等；安防领域检测无人机、道路障碍等。此时我们需要将YOLOv5部署在自己的数据集上，再训练模型，测试模型。

其实要做到这一点并不困难，我们会采用最简单的方法，遵循两点：

• 将代码自带的数据集路径换成我们自己的。

• 将我们数据集的结构保持和自带的数据集一致。

做到了以上两点，模型自然就在我们自己的数据集上训练了。

1.1 COCO数据集的结构

我们直接通过coco.yaml文件来看看COCO的数据集结构。

我们在用yolov5训练COCO数据集时运行的命令是：

python train.py --data coco.yaml --epochs 300 --weights '' --cfg yolov5n.yaml --batch-size 128

这其中包含coco.yaml，就表明我们要训练的数据是COCO数据集，我们看看yaml文件里是什么内容：

最上面有一个parent，包含yolov5和datasets，其中yolov5是代码的文件夹，datasets文件夹中放的就是具体的数据集。 再往下看。

• path： 这个path是train.py所在的路径，一般就在yolov5文件夹中。

path: ../datasets/coco说明数据存放在datasets中，给出了数据集存放的路径，当然物品们也可以修改。

• train、val、test： 下面有3个txt，给出了训练集、验证集以及测试集数据的具体路径。

• name： 这个数据集中包含哪些种类，当然总共有80种，这里我没有全部截取出来。

1.2 数据集制作方法

知道了COCO数据集的结构，接下来，我们就按照上面这种格式去制作自己的数据集就可以。

做检测的数据集标签包含2种数据：检测的物体的种类以及边框四个点的坐标，比如我们来看一个标签：

前面的数字就表示物体的种类，后面几个0-1之间的小数就是物体的坐标。当然一张图片里不一定只有一个需要检测的物体，因此这个标签里有多个这样的数据。

那么我们如何打标签呢？自己一个个手写太麻烦了，我们采用工具labelimg。

激活虚拟环境，然后直接执行安装命令：

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

安装好后，在命令行中直接输入labelimg就有弹窗了：

labelimg如何使用呢？

首先我们需要创建一个datasets的文件夹，再在datasets下创建yolo_data,在yolo_data下创建两个文件夹：images和labels。

然后在images文件夹下分别再创建train和val，分别存放训练集的图片以及验证集的图片；

在labels文件夹下分别再创建train和val，分别存放训练集的图片以及验证集的标签；

如下图所示：

其中图片需要我们自己分类好放进去，标签的txt就是由labelimg生成的。

我们首先来制作训练集的标签，通过labelimg打开yolo_data/images/train:

打开后：

此时我们按"w"键，并且鼠标在图片上框出想要检测的物体，随后会自弹出一个窗口让你表明框出的物体所属类别(也就是前面的yaml文件中描述的)，比如我们用1表示这是一个人,再点击ok。

此时你会发现，在yolo_data/labels/train中保存的不是txt文件，而是xml文件，这是由于保存成了VOC格式，如果想要保存成yolo格式，只需要切换，点击下面这个箭头，再点击PascalVOC：

就切换为了yolo格式：

窗口中的Next Image以及Prev Image表示下一张图片和上一张图片，我们将文件夹yolo_data/images/train中的图片标签保存至yolo_data/labels/train；将文件夹yolo_data/images/val中的图片标签保存至yolo_data/labels/val，就完成了数据集的制作。

2 编辑yaml文件

数据集制作好后我们的工作已经完成了一大半，接下来我们编辑yaml文件：

我们将自己的yaml文件命名为yolo_data.yaml，你想叫什么名字都可以：

• 其中第一行是表明我们数据集的路径，在代码train.py路径的上一层

• train: ../images/train指定了训练集的图像路径，这个文件是相对于path的，即实际路径为datasets/yolo_data/images/train

• val: ../images/val指定了验证集的图像路径，这个文件是相对于path的，即实际路径为datasets/yolo_data/images/val

• names表明每个数字表示的种类 有人可能会说标签文件怎么获取呢？

在yolov5中，有直接的代码处理逻辑，刚才我们在用labelimg打标签时，生成的标签文件名和对应的图片一样，yolov5在训练时，会将datasets/yolo_data/images/val直接替换为datasets/yolo_data/labels/val，这样就可以找到标签了。

3 训练自己的数据集

数据集也制作好了，yaml文件也制作好了，我们将数据集放在和代码同级目录下，并将yaml文件放到data目录下，直接执行:

python train.py --data yolo_data.yaml --epochs 300 --weights '' --cfg yolov5n.yaml --batch-size 128

代码就起飞了：

最后我们训练得到的模型会存储到yolov5-master/runs/train/exp下面，如果我们要测试，执行命令：

python detect.py --weights runs/train/exp9/weights/best.pt --source test.jpg

生成的结果保存在 runs/detect/exp下。

欢迎大家关注我的gzh：阿龙AI日记