浙江龙网

  1. 首页 > 生活百科 > >

如何使用 Yolov4 训练人脸口罩检测模型

生活百科

前言疫情当下 ， 出入医院等公共场所都被要求佩戴口罩 。这篇博客将会介绍如何使用 Yolov4，训练一个人脸口罩检测模型（使用 Yolov4 的原因是目前只复现到了 v4 ） ， 代码地址为 https://github.com/zhiyiYo/yolov4 。
Yolov4Yolov4 的神经网络结构相比 Yolov3 变化不是很大，主要更换了激活函数为 Mish，增加了 SPP 块和 PAN 结构（图源 《yolo系列学习笔记----yolov4（SPP原理）》） 。

如何使用 Yolov4 训练人脸口罩检测模型

文章插图
感觉 Yolov4 最大的特点就是使用了一大堆的 Trick，比如数据增强方面使用了马赛克数据增强、Mixup 数据增强，将定位损失函数更换为 CIOU 损失 。论文中提到了很多的 Trick，我的代码中没有全部复现，不过在 VOC2012 数据集训练了 160 个 epoch 之后 mAP 也能达到 83%，效果还是不错的 。
可以在终端使用下述指令下载 Yolov4 的代码：
git clone https://github.com/zhiyiYo/yolov4.git人脸口罩数据集网上可以找到很多人脸口罩数据集，这里使用的是 AIZOOTech 提供的数据集 。由于这个数据集的结构和 Pascal VOC 数据集不一样，所以重新组织一下数据集 ， 并且修复和移除了数据集中的非法标签 ， 可以在 Kaggle 上下载此数据集 。目前这个数据集包含 6130 张训练图像，1839 张测试图像，对于 Yolov4 的训练来说应该是绰绰有余的 。下载完数据集将其解压到 data 文件夹下 。
在训练之前，我们需要使用 K-means 聚类算法对训练集中的边界框进行聚类，对于 416×416 的输入图像，聚类结果如下：
anchors = [[[100, 146], [147, 203], [208, 260]],[[26, 43], [44, 65], [65, 105]],[[4, 8], [8, 15], [15, 27]]]训练神经网络训练目标检测模型一般都需要加载预训练的主干网络的权重，可以从谷歌云盘下载预训练好的权重 CSPDarknet53.pth 并将其放在 model 文件夹下 。这里给出训练所用的代码 train.py，使用 python train.py 就能开始训练 。模型会先冻结训练上 50 个 epoch，接着解冻训练 110 个 epoch：
# coding:utf-8from net import TrainPipeline, VOCDatasetfrom utils.augmentation_utils import YoloAugmentation, ColorAugmentation# 训练配置config = {"n_classes": len(VOCDataset.classes),"image_size": 416,"anchors": [[[100, 146], [147, 203], [208, 260]],[[26, 43], [44, 65], [65, 105]],[[4, 8], [8, 15], [15, 27]]],"darknet_path": "model/CSPdarknet53.pth","lr": 1e-2,"batch_size": 8,"freeze_batch_size": 16,"freeze": True,"freeze_epoch": 50,"max_epoch": 160,"start_epoch": 0,"num_workers": 4,"save_frequency": 10,"no_aug_ratio": 0}# 加载数据集root = 'data/FaceMaskDataset/train'dataset = VOCDataset(root,'all',transformer=YoloAugmentation(config['image_size']),color_transformer=ColorAugmentation(config['image_size']),use_mosaic=True,use_mixup=True,image_size=config["image_size"])if __name__ == '__main__':train_pipeline = TrainPipeline(dataset=dataset, **config)train_pipeline.train()测试神经网络训练完使用 python evals.py 可以测试所有保存的模型 ， evals.py 代码如下：
# coding:utf-8import jsonfrom pathlib import Pathimport matplotlib as mplimport matplotlib.pyplot as pltfrom net import EvalPipeline, VOCDatasetmpl.rc_file('resource/theme/matlab.mplstyle')# 载入数据集root = 'data/FaceMaskDataset/val'dataset = VOCDataset(root, 'all')anchors = [[[100, 146], [147, 203], [208, 260]],[[26, 43], [44, 65], [65, 105]],[[4, 8], [8, 15], [15, 27]]]# 列出所有模型，记得修改 Yolo 模型文件夹的路径model_dir = Path('model/2022-10-05_22-59-44')model_paths = [i for i in model_dir.glob('Yolo_*')]model_paths.sort(key=lambda i: int(i.stem.split("_")[1]))# 测试所有模型mAPs = []iterations = []for model_path in model_paths:iterations.append(int(model_path.stem[5:]))ep = EvalPipeline(model_path, dataset, anchors=anchors, conf_thresh=0.001)mAPs.append(ep.eval()*100)# 保存数据with open('eval/mAPs.json', 'w', encoding='utf-8') as f:json.dump(mAPs, f)# 绘制 mAP 曲线fig, ax = plt.subplots(1, 1, num='mAP 曲线')ax.plot(iterations, mAPs)ax.set(xlabel='iteration', ylabel='mAP', title="mAP curve")plt.show()得到的 mAP 曲线如下图所示，在第 120 个 epoch 达到最大值 94.14%：
如何使用 Yolov4 训练人脸口罩检测模型

文章插图
下面使用一张真实图像看看训练效果如何，运行 demo.py：
# coding:utf-8from net import VOCDatasetfrom utils.detection_utils import image_detect# 模型文件和图片路径model_path = 'model/Yolo_120.pth'image_path = 'resource/image/三上老师.jpg'# 检测目标anchors = [[[100, 146], [147, 203], [208, 260]],[[26, 43], [44, 65], [65, 105]],[[4, 8], [8, 15], [15, 27]]]image = image_detect(model_path, image_path, VOCDataset.classes, anchors=anchors, conf_thresh=0.5)image.show()

推荐阅读


上一篇：用 VS Code 搞Qt6：使用 PySide 6

下一篇：一篇文章带你掌握主流服务层框架——SpringMVC