使用yolov5训练自己的目标检测数据集
yolov4出来后不久,又出现了yolov5,没有论文。虽然作者没有放上和yolov4的直接测试对比,但在COCO数据集的测试效果还是很可观的。很多人考虑到YOLOv5的创新性不足,对算法是否能够进化,因此yolov5的名头仍有争议。但是既然github上有如此多的star,说明还是受到大多数人认可的一个工程。
并且yolov5是十分容易上手使用的一个目标检测开源项目,只需要改几个文件,配置好自己的数据集格式和目录结构,即可开始训练,对新手十分友好。可以先把训练跑起来,在根据需要去细扣其中的模块设计。
并且,yolov5项目一直在更新维护,除了日常处理bug,还有更优算法的版本(如激活函数的更新)更迭和新的训练可视化工具wandb的加入等,这也是十分难得的。
今天,笔者就来介绍一下自己第一次使用yolov5训练自己的数据集的过程。
工程项目准备
PyTorch等其他依赖的安装这里就不说了,网上有很多教程。
首先从yolov5 github上克隆下代码:
git clone https://github.com/ultralytics/yolov5.git
然后从github下载预训练权重(.pt文件),没有百度云版本,这里笔者提供下载好的权重文件放到云盘(密码: afp3)供不方便的朋友直接下载。下载好后放到yolov5/weights文件夹下。
数据集准备
yolov5项目所需的数据集文件路径格式如下:
---yolov5
---yolo_data
---train
---images
---labels
---val
---images
---labels
---data.yaml
数据集文件夹放在与工程文件同级的目录下。其中包含两个子文件夹和一个配置文件data.yaml,两个子文件夹分别存放训练样本和验证样本,注意images和labels文件夹的名字不要改。
先说一下data.yaml配置文件内的内容
train: ../yolo_data/train
val: ../yolo_data/val
nc: [你的检测数据集的类别数]
names: [你的检测数据集的类别名称列表]
共有四个键值对,train,val,nc,names,分别对应如下表:
键 | 值 |
---|---|
train | 你的训练数据的文件夹路径 |
val | 你的验证数据的文件夹路径 |
nc | 你的检测数据集的类别数(如COCO数据集中nc=80) |
names | 你的检测数据集的类别名称列表(如:[‘plane’, ‘car’, ‘dog’, …]) |
然后说一下数据集格式的要求:
在train/val目录下要有images和labels两个目录,分别存放图像和标签,其中的文件名要相互对应,分别为.jpg格式和.txt格式,即:
---train
---images
aaaa.jpg
...
---labels
aaaa.txt
...
其中图像不必多说,放入你的数据集图像即可,而标签label目录下的txt文件要求格式如下:
类别索引 x y w h
0 0.53857421875 0.8779296875 0.158203125 0.06184895833333333
第一列是图像中每个object的类别索引,注意这里要与data.yaml配置文件中的names相互对应。后面四个列就是该object的检测框坐标了,这里采用的是经过归一化的中心长宽xywh格式,如果你的数据集原生不是这种格式(如对角坐标xyxy)的话,请注意调整。
开始训练
直接运行yolov5/train.py即可,注意配置好参数
python train.py \
--data ../yolo_data/data.yaml \
--cfg models/yolov5s.yaml \
--weights weights/yolov5s.yaml \
--batch-size 8 \
--epochs 100
yolov5最近一次更新增加了wandb(weights and biases,其官网,直接访问不太稳定,建议科学上网访问)训练可视化工具,没有安装的话需要先安装一下,直接pip安装即可
pip install wandb
第一次使用的话会需要注册账号,访问站点https://wandb.ai/authorize,拿到API Key,复制到终端中即可,首次绑定API Key账号后再训练就不在需要输入API Key了,直接将训练可视化结果送到你的wandb账号,可在该站点登录查看。
不仅有训练过程可视化,loss曲线等
甚至还有训练过程中系统的信息的展示:
还是非常简单实用的一个训练可视化工具。
训练完成后再yolov5/runs/train/exp文件夹下可查看一些结果PR曲线,F1曲线等,还有一些样本可视化的结果。
可以看到效果还是可以的(我这个数据集只有猪仔一个类 ^^可爱小猪仔)。
测试
训练完成后可以运行yolov5.detect.py文件,指定一个文件夹,测试其中的图像。
python detect.py \
--weights runs/train/exp/weights/best.pt \
--source ../images/test \
--save-txt
如果设置打开–save-txt,则会保存结果文件。
测试结果会保存在yolov5/runs/detect内。txt结果文件则会保存在yolov5/runs/detect/labels内。
使用yolov5训练自己的数据集就先介绍到这里,如果有问题欢迎留言讨论。