YOLOv7模型训练 | 附数据集
AI算法与图像处理
共 1332字,需浏览 3分钟
·
2022-08-08 22:32
环境
ubuntu 18.04 64bit torch 1.7.1+cu101
数据集准备
这里以前面我们进行 YOLOv5 模型训练 时用到的口罩数据集为例,这个数据集来自网站 roboflow.com
,再次安利一下这个站点,真的非常棒,不止有详尽的博客教程,还有很多的开放数据集,而且支持的数据格式也很丰富,绝对值得经常去逛一逛。
口罩数据集下载地址:https://public.roboflow.com/object-detection/mask-wearing/4,这里也在百度网盘上存一份,需要的自取
链接:https://pan.baidu.com/s/1JvniT205zX79wASqiKtt5Q
提取码:9feh
数据集下载后解压,将文件夹重命名为 mask
,并放到 yolov7
的根目录下(这里可以随意,只要前后的路径匹配上就可以了),完整的目录结构如下
可以发现,其实 yolov7
和 yolov5
数据集格式是一模一样的,通过标注工具 labelimg
也可以得出同样的结论
训练
模型训练开始之前,我们需要修改部分配置文件
首先是数据集中的 data.yaml
train: mask/train/images
val: mask/valid/images
nc: 2
names: ['mask', 'no-mask']
其次,是修改 yolov7
中的配置文件 cfg/training/yolov7.yaml
,主要是 nc
这个字段
nc: 2 # number of classes
然后,就可以进行训练了,执行
python train.py --data mask/data.yaml --cfg cfg/training/yolov7.yaml --weights '' --name yolov7 --hyp data/hyp.scratch.p5.yaml
模型测试
最后,我们来测试下模型的效果
python detect.py --source mask/test/images/shutterstock_1627199179_jpg.rf.350e69105dd1458572a590c3e3ef2538.jpg --weight runs/train/yolov7/weights/best.pt
python detect.py --source mask/test/images/the-first-day-of-wuhan-s-closure-some-people-fled-some-panicked_jpg.rf.51ed69bf8d327d93b429a08581f6dea0.jpg --weight runs/train/yolov7/weights/best.pt
评论