基于YOLOv5和PyTorch的实时口罩检测系统开发实践-AI智能范式网

基于YOLOv5和PyTorch的实时口罩检测系统开发实践

乐正雕漆

1. 项目概述与背景

去年参与了一个公共场所防疫监控项目，需要实时检测人员是否佩戴口罩。当时尝试了多种方案，最终基于YOLOv5和PyTorch搭建的口罩检测系统在准确率和速度上达到了最佳平衡。这个经历让我深刻体会到目标检测技术在实际应用中的价值，今天就把整个搭建过程整理成技术笔记分享给大家。

YOLOv5作为目前最流行的实时目标检测框架之一，相比前代YOLOv3在精度和速度上都有显著提升。而PyTorch作为深度学习领域的主流框架，其动态计算图和丰富的模型库为快速开发提供了便利。两者结合非常适合开发轻量级的边缘端应用。

2. 环境配置与工具准备

2.1 基础环境搭建

推荐使用Anaconda创建独立的Python环境，避免包冲突。我使用的环境配置如下：

Python 3.8.5
PyTorch 1.7.1
CUDA 11.0（GPU加速）
cuDNN 8.0.5

安装PyTorch时需要注意版本匹配问题。以CUDA 11.0为例，正确的安装命令是：

bash复制pip install torch==1.7.1+cu110 torchvision==0.8.2+cu110 -f https://download.pytorch.org/whl/torch_stable.html

提示：如果使用CPU版本，去掉cu110后缀即可，但推理速度会明显下降

2.2 YOLOv5源码获取与依赖安装

官方仓库提供了完整的实现：

bash复制git clone https://github.com/ultralytics/yolov5
cd yolov5
pip install -r requirements.txt

关键依赖包括：

OpenCV：图像处理
Matplotlib：结果可视化
PyYAML：配置文件解析
TensorBoard：训练过程监控

3. 数据集准备与标注

3.1 数据收集策略

优质的数据集是模型性能的基础。我们通过多种渠道收集数据：

公开数据集：MAFA、WIDER Face等
网络爬取：注意版权问题
实际场景拍摄：确保数据分布接近真实应用

最终我们整理了约15,000张含有人脸的图片，覆盖不同：

光照条件（强光/弱光/背光）
遮挡情况（部分遮挡/完全遮挡）
口罩类型（医用/布质/N95等）
人脸角度（正面/侧面/俯仰）

3.2 数据标注规范

使用LabelImg工具进行标注，需要注意：

标注框应紧贴口罩边缘
对于部分遮挡的情况，按可见部分标注
类别分为两类：mask（正确佩戴）和no_mask（未佩戴或错误佩戴）

标注后的数据采用YOLO格式保存，每个图像对应一个.txt文件，格式为：

code复制<class_id> <x_center> <y_center> <width> <height>

其中坐标值是相对于图像宽高的归一化值。

3.3 数据增强策略

在data.yaml中配置以下增强参数：

yaml复制train: ../train/images
val: ../val/images

nc: 2  # 类别数
names: ['mask', 'no_mask']  # 类别名称

# 数据增强参数
hsv_h: 0.015  # 色调变化
hsv_s: 0.7    # 饱和度变化 
hsv_v: 0.4    # 明度变化
degrees: 10   # 旋转角度
translate: 0.1  # 平移比例
scale: 0.5    # 缩放比例
shear: 0.0    # 剪切变换

4. 模型训练与调优

4.1 模型选择与预训练

YOLOv5提供了多种规模的模型：

YOLOv5s：最小最快，适合移动端
YOLOv5m：平衡型
YOLOv5l：高精度
YOLOv5x：最大最精确

考虑到实时性要求，我们选择YOLOv5m作为基础模型，使用官方预训练权重初始化：

bash复制python train.py --img 640 --batch 16 --epochs 100 --data data.yaml --cfg models/yolov5m.yaml --weights yolov5m.pt

4.2 关键训练参数

输入尺寸：640x640（平衡精度和速度）
Batch size：根据GPU显存调整（11GB显存可设16）
学习率：0.01（使用余弦退火调度）
优化器：SGD（动量0.937，权重衰减5e-4）
损失函数：GIoU + obj + cls

4.3 训练监控与调优

使用TensorBoard监控训练过程：

bash复制tensorboard --logdir runs/train

重点关注以下指标：

损失曲线（train/val应同步下降）
mAP@0.5（主要精度指标）
召回率（避免漏检）

如果出现过拟合，可以：

增加数据增强强度
添加Dropout层
提前停止训练

5. 模型测试与部署

5.1 性能评估

使用验证集测试模型表现：

bash复制python val.py --weights runs/train/exp/weights/best.pt --data data.yaml --img 640

我们的最佳模型达到了：

mAP@0.5：98.7%
推理速度：45 FPS（RTX 2080Ti）
模型大小：85MB

5.2 模型导出

部署前需要将模型转换为ONNX或TorchScript格式：

bash复制python export.py --weights best.pt --img 640 --batch 1 --include onnx

5.3 部署方案

根据场景选择不同部署方式：

桌面应用：

python复制import torch
model = torch.hub.load('ultralytics/yolov5', 'custom', path='best.pt')
results = model(img)

Web服务（Flask示例）：

python复制from flask import Flask, request
app = Flask(__name__)

@app.route('/detect', methods=['POST'])
def detect():
    file = request.files['image']
    img = Image.open(file.stream)
    results = model(img)
    return results.pandas().xyxy[0].to_json()

移动端：使用TensorRT或MNN加速

6. 实际应用中的优化技巧

6.1 小目标检测优化

对于远距离人脸检测：

增加输入分辨率（--img 1280）
使用更密集的anchor box
添加注意力机制

6.2 误检过滤策略

通过后处理减少误检：

置信度阈值：0.5
NMS阈值：0.4
人脸比例过滤（排除过大/过小检测框）

6.3 边缘设备优化

在树莓派等设备上的优化方法：

使用YOLOv5s模型
量化为INT8精度
启用OpenVINO加速

7. 常见问题与解决方案

CUDA内存不足：
- 减小batch size
- 使用--img 320降低分辨率
- 尝试--device cpu
标注不准确导致性能差：
- 检查标注框是否准确
- 确保没有漏标目标
- 平衡各类别样本数量
推理速度慢：
- 导出时设置--half使用FP16
- 使用TensorRT加速
- 优化预处理流水线

这个项目让我深刻体会到，一个好的AI系统不仅需要优秀的算法，更需要：

高质量的数据
合理的工程实现
持续的优化迭代

在实际部署中，我们还加入了温度检测和人数统计功能，形成了完整的防疫解决方案。代码和配置文件已整理到GitHub，欢迎交流讨论。