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

前言

疫情当下，出入医院等公共场所都被要求佩戴口罩。这篇博客将会介绍如何使用 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()

不错，效果非常好 😊：

后记

至此，介绍完了训练 Yolov4 人脸口罩检测模型的过程，代码放在了 https://github.com/zhiyiYo/yolov4，以上~~