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

技术分享 2年前 (2022-10-07) 0 999+

前言

疫情当下，出入医院等公共场所都被要求佩戴口罩。这篇博客将会介绍如何使用 Yolov4，训练一个人脸口罩检测模型（使用 Yolov4 的原因是目前只复现到了 v4 😇），代码地址为 https://github.com/zhiyiYo/yolov4。

Yolov4

Yolov4 的神经网络结构相比 Yolov3 变化不是很大，主要更换了激活函数为 Mish，增加了 SPP 块和 PAN 结构（图源《yolo系列学习笔记----yolov4（SPP原理）》）。

感觉 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-8 from net import TrainPipeline, VOCDataset from 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-8 import json from pathlib import Path  import matplotlib as mpl import matplotlib.pyplot as plt  from net import EvalPipeline, VOCDataset  mpl.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%：

下面使用一张真实图像看看训练效果如何，运行 demo.py：

# coding:utf-8 from net import VOCDataset from 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()

不错，效果非常好 😊：