https://xugaoxiang.com/2020/07/02/yolov5-training/
请参考之前的文章,YOLOv5目标检测
YOLOv5
的预训练模型是基于 COCO 数据集,如果自己想去复现下训练过程,可以依照下面的命令
$ python train.py --data coco.yaml --cfg yolov5s.yaml --weights '' --batch-size 64
yolov5m 48
yolov5l 32
yolov5x 16
COCO
的数据集可以通过data
文件夹下get_coco2017.sh
脚本进行下载,包含图片和lable
文件。COCO
的数据集实在是太大了,整个压缩包有18G,考虑到自己到的网速还有机器的算力,还是洗洗睡吧。。。
如果没有对应目标的公开数据集,那就只有自己出手收集了,图片到手后,接下来就是艰辛的打标签工作了,这里使用工具LabelImg
,下载地址是
https://github.com/tzutalin/labelImg/releases/tag/v1.8.1
LabelImg
使用Qt
做了图形化的界面,操作还是很方便的,这也是选择它的理由,它提供了默认的class
,如果你不需要这些类型的话,可以将其删除
接下来就可以打开exe
文件,点击Open
导入图片,按下快捷键w
,选定目标后,会弹出输入框,写上class
名称,就可以了,如果有多个目标,那就继续标
[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-ZT8OXIjR-1593679533741)(https://image.xugaoxiang.com/imgs/2020/07/a909c99b3162803d.png)]
labelImg
还支持文件夹的导入,在标完一张后,在左侧选择Next Image
就可以切换到下一张继续了。输出格式部分,目前labelImg
支持YOLO
和PascalOVC
2种格式,前者标签信息是存储在txt
文件中,而后者是存储在xml
中
打完标签后,就可以进行保存了,图片和标签文件我们分开存放,但是文件名是对应的,只是扩展名不同
最后来看看标签文件的内容
[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-QTIqrfvn-1593679533765)(https://image.xugaoxiang.com/imgs/2020/07/e2862fb5ae1d188a.png)]
一行代表一个目标,格式是
class x_center y_center width height
第一列是class
的索引,计数从0开始,比如这里的0代表的是basketball
,1代表的是face
;后4列是x_center/image_width
、y_center/image_height
、width/image_width
、height/image_height
,取值范围是0 ~ 1
ROBOFLOW提供了一些公开的数据集,我们下载其中的口罩数据集进行训练,链接是https://public.roboflow.ai/object-detection/mask-wearing,下载下来是一个压缩包,解压后,文件夹内的文件结构是这样的
其中文件夹train
包含了参加训练的图片以及对应的label
文件,两者只有扩展名不同而已,目前图片只有105张。我们将包含数据集的文件夹重命名为mask
,存储在yolov5
工程的同级目录下
[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-vAsT7OIP-1593679533790)(https://image.xugaoxiang.com/imgs/2020/07/fd568bfb729b399d.png)]
接着修改mask/data.yaml
文件内容为
(base) xugaoxiang@1070Ti:~/Works/github/mask$ cat data.yaml
train: ../mask/train/images
val: ../mask/valid/images
nc: 2
names: ['mask', 'no-mask']
最后修改yolov5/models/yolov5s.yaml
,将nc = 80
修改为nc = 2
,因为数据集中只有mask
和no-mask
2个类别
接下来执行训练命令
cd yolov5
python train.py --img 640 --batch 16 --epochs 300 --data ../mask/data.yaml --cfg models/yolov5s.yaml --weights ''
训练结束后,在weights
文件夹下就生成了best.pt
和last.pt
,到mask/test/images
找些图片测试一下
python detect.py --weight weights/best.pt --source ../mask/test/images/1224331650_g_400-w_g_jpg.rf.b816f49e2d84044fc997a8cbd55c347d.jpg
[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-R5TocMf6-1593679533800)(https://image.xugaoxiang.com/imgs/2020/07/bdfad59d7f768358.jpg)]
[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-aYtc5ZtT-1593679533806)(https://image.xugaoxiang.com/imgs/2020/07/16982ad827626e3e.jpg)]
效果还算ok
。感兴趣的话,自己动手试试吧