如何使用 Yolov4 訓練人臉口罩檢測模型

• 2022 年 10 月 6 日
• 筆記
• deep-learning, Python

前言

疫情當下，出入醫院等公共場所都被要求佩戴口罩。這篇博客將會介紹如何使用 Yolov4，訓練一個人臉口罩檢測模型（使用 Yolov4 的原因是目前只復現到了 v4 😇），代碼地址為 //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 //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 人臉口罩檢測模型的過程，代碼放在了 //github.com/zhiyiYo/yolov4，以上~~