如何使用 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，以上~~
