尧图网站设计 尧图网站设计YAOTU DESIGN
ARTICLE DETAIL

资讯详情

深耕网站设计与一线实操的经验洞察。

PyTorch实现YOLOv4口罩识别:小目标检测实战解析

PyTorch实现YOLOv4口罩识别:小目标检测实战解析 简介一套基于YOLOv4与PyTorch的人脸口罩识别项目源码主要面向计算机视觉方向的在校学生、毕业设计开发者也适合希望系统学习目标检测实战的初学者。代码已经过完整测试并成功运行除提供Python训练与推理脚本外还附带详细的文档说明、口罩数据集标注、预训练模型权重以及多种场景的测试视频能够完整覆盖从数据准备、模型训练、效果评估到最终推理的各个环节。压缩包内共包含三百六十七个文件以XML标注文件为主体同时包括Python源码、文本配置、模型文件、效果图片与AVI视频等整体大小仅九点三一兆字节下载与部署都很轻便。目前已有五十二人学习下载。包内额外提供运行成功后的演示视频和结果截图便于读者直观对照输出效果若遇到环境配置或运行问题作者也可通过私聊远程教学辅助解决适合直接用于课程设计、毕业设计、项目初期演示是一份门槛低、可复现性强的实战入门资料。1. 口罩识别不是图像分类而是小目标检测拿到这套 YOLOv4 PyTorch 的人脸口罩识别源码时我第一反应是先翻它的文件结构output202212031535.avi、output202212041735.avi是现成的测试视频result.jpg是检测结果帧video_tmp.jpg大概率是中间调试产物README.md是入口文档。值得注意的还有.gitignore、yolov4-pytorch-master.iml和VSWorkspaceState.json前者说明作者把数据集和权重排除在版本控制外后两者意味着这个项目既被 PyCharm 打开过也被 Visual Studio 加载过工程迁移没有遇到环境障碍。口罩识别的难点不在“识别口罩”本身而在口罩区域在画面中占比极小、人脸角度多变、遮挡物干扰强把它做成纯分类模型大概率翻车标准解法是目标检测框架而这套代码选择的正是 YOLOv4。适合的人群很明确计算机视觉方向的毕设学生、刚接触 PyTorch 目标检测的工程师以及需要在本地快速复现一个可用于课程答辩或立项演示的完整项目。它不只是一个模型文件而是从数据标注格式到视频推理的完整链路。2. YOLOv4 在 PyTorch 中的核心结构与推理链路2.1 主干网络 CSPDarknet53 与 Mish 激活的设计意图YOLOv4 的主干特征提取网络是 CSPDarknet53它在 Darknet53 的基础上引入了 CSPCross Stage Partial结构。CSP 的核心操作是把输入特征图沿通道维度分成两部分一部分直接连接到模块尾部另一部分经过残差块堆叠后再拼接这样做的直接收益是减少了重复梯度计算在推理速度基本不变的前提下保持了较高的特征复用率。PyTorch 实现里通常用nn.ModuleList组织 5 个 CSP 阶段每个阶段内部包含若干个BasicBlock每个BasicBlock是 1x1 卷积 3x3 卷积的残差对。关键细节是激活函数。YOLOv4 在主干中大量使用 Mish 而非 ReLUMish 的公式是x * tanh(softplus(x))它没有 ReLU 在负数区间的硬零截断梯度在负半轴依然平滑训练深层网络时更容易收敛。代价是计算量增加PyTorch 中可以直接用nn.Mish()或者用x * torch.tanh(F.softplus(x))手动实现后者在导出 ONNX 或 TorchScript 时反而更可控。import torch import torch.nn as nn class Mish(nn.Module): def forward(self, x): return x * torch.tanh(F.softplus(x)) class BasicBlock(nn.Module): def __init__(self, in_channels, hidden_channelsNone): super().__init__() hidden_channels hidden_channels or in_channels self.conv1 nn.Conv2d(in_channels, hidden_channels, 1, biasFalse) self.bn1 nn.BatchNorm2d(hidden_channels) self.conv2 nn.Conv2d(hidden_channels, in_channels, 3, padding1, biasFalse) self.bn2 nn.BatchNorm2d(in_channels) def forward(self, x): identity x out Mish()(self.bn1(self.conv1(x))) out Mish()(self.bn2(self.conv2(out))) return out identity这段代码里biasFalse不是随手写的因为后面紧跟着 BatchNorm 层卷积的偏置会被 BN 的均值偏移抵消保留 bias 只会增加冗余计算。残差连接out identity要求输入输出通道数一致所以conv1先降通道再conv2升回来这种瓶颈设计与 ResNet 的 Bottleneck 思路一致。2.2 SPP 与 PANet 的 PyTorch 实现CSPDarknet53 输出的三个有效特征层会进入 SPP 和 PANet 结构。SPPSpatial Pyramid Pooling的作用是扩大感受野把同一特征图分别做 5x5、9x9、13x13 的最大池化再与原图在通道维拼接。为什么要这样做因为口罩在画面中的尺寸变化非常大戴口罩的人可能站在 0.5 米处占据半屏也可能站在 5 米外只有几十个像素多尺度池化让模型对这类尺度抖动不敏感。class SPP(nn.Module): def __init__(self, pool_sizes(5, 9, 13)): super().__init__() self.pools nn.ModuleList( [nn.MaxPool2d(kernel_sizes, stride1, paddings // 2) for s in pool_sizes] ) def forward(self, x): features [x] [pool(x) for pool in self.pools] return torch.cat(features, dim1)SPP 后面跟 PANet。PANet 是一条自底向上的路径增强结构它把高层的语义信息向下传递同时把浅层的空间位置信息向上融合。人脸口罩任务里浅层特征负责找“脸在哪”深层特征负责判断“口罩戴没戴”两条路径拼接后模型才能同时保证定位精度和分类精度。2.3 YOLO Head 解码从网格到真实坐标YOLOv4 预测头输出的原始张量形状是(batch, anchors, grid_h, grid_w, 5 num_classes)最后一个维度的 5 对应(tx, ty, tw, th, obj_conf)。解码时要把偏移量换算成真实框坐标核心公式是中心点坐标bx sigmoid(tx) grid_xby sigmoid(ty) grid_y宽高bw anchor_w * exp(tw)bh anchor_h * exp(th)中心点用 sigmoid 把输出限制在 0 到 1 之间确保预测中心不会偏离当前网格太远。宽高用 exp 是因为锚框的宽高比例是对数空间学习的exp 还原回线性空间。def decode_prediction(pred, anchors, num_classes): batch, num_anchors, grid_h, grid_w, _ pred.shape pred pred.view(batch, num_anchors, grid_h, grid_w, 5 num_classes) pred[..., 0] torch.sigmoid(pred[..., 0]) # tx pred[..., 1] torch.sigmoid(pred[..., 1]) # ty pred[..., 4] torch.sigmoid(pred[..., 4]) # obj_conf # 后续通过 meshgrid 生成 grid_x / grid_y 并叠加 return predobj_conf用 sigmoid 压缩到 [0, 1]它表示当前网格内是否有目标的置信度。人脸口罩场景里这张图通常被分成 13x13、26x26、52x52 三个尺度52x52 的浅层网格负责捕捉小尺寸人脸13x13 的深层网格负责大尺寸人脸多尺度预测是 YOLOv4 能同时处理近景大头和远景小人脸的关键。3. 数据准备、训练参数与模型复现3.1 数据格式VOC XML 还是 YOLO Txt打开源码包的 data 目录常见的组织方式有两种VOC 格式JPEGImages Annotations ImageSets/Main或 YOLO 格式images labels。这个项目的 README 里如果标注了 voc 或 yolo 关键字直接对应找即可。YOLO 格式的标签文件是个 txt每行内容为class_id x_center y_center width height四个坐标值都是相对图像宽高的归一化比例取值范围 0~1。VOC 格式则是 XML记录绝对像素坐标。如果数据是 VOC 格式训练前需要转换成 YOLO 格式。常见做法是直接用开源脚本批量转换也可以手写一个解析器import xml.etree.ElementTree as ET import os def voc_to_yolo(xml_path, out_dir, class_names): tree ET.parse(xml_path) root tree.getroot() img_w int(root.find(size/width).text) img_h int(root.find(size/height).text) yolo_lines [] for obj in root.iter(object): name obj.find(name).text if name not in class_names: continue box obj.find(bndbox) xmin float(box.find(xmin).text) ymin float(box.find(ymin).text) xmax float(box.find(xmax).text) ymax float(box.find(ymax).text) x_center (xmin xmax) / 2 / img_w y_center (ymin ymax) / 2 / img_h width (xmax - xmin) / img_w height (ymax - ymin) / img_h yolo_lines.append(f{class_names.index(name)} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}) out_path os.path.join(out_dir, os.path.splitext(os.path.basename(xml_path))[0] .txt) with open(out_path, w) as f: f.write(\n.join(yolo_lines))转换逻辑里容易忽略的坑是坐标归一化用宽高做分母不是用最大边。如果图片是矩形而非正方形直接用绝对像素会导致锚框比例失真。另外class_names的索引顺序必须和训练配置里的classes_path保持一致否则类别会错位。3.2 从测试视频抽帧扩充数据集源码包里的output202212031535.avi、output202212041735.avi这类测试视频本身就可以用来抽帧补充训练集。用 OpenCV 按帧读取间隔固定帧数保存再手动筛选清晰含人脸的画面。口罩识别的数据难点是正负样本不均衡不带口罩的人脸样本往往远多于带口罩的且戴口罩的样本里口罩颜色、样式、遮挡程度差异极大。抽帧后应尽量保证两类样本比例控制在 1:1 到 2:1 之间超过这个范围时模型会偏向多数类。3.3 训练参数设定参考YOLOv4 的 PyTorch 训练脚本通常在train.py里通过 argparse 接收参数这套源码的训练入口一般支持--model_path、--epochs、--batch_size、--lr等。以本人复现经验来看口罩识别场景的参数初始化可以参考下表参数推荐值说明input_size416x416显存有限时的首选速度与精度较平衡batch_size8~16取决于显存8GB 显存建议 8epochs50~100数据量小则 50数据量大可到 100initial_lr1e-3采用余弦退火或 step 衰减optimizerAdam前期收敛快后期可切 SGDmosaic4开启马赛克增强小目标收益明显label_smoothing0.1降低过拟合类别数少时尤其有效3.4 训练循环与保存策略for epoch in range(start_epoch, epochs): model.train() for batch_idx, (images, targets) in enumerate(train_loader): images images.to(device) targets [t.to(device) for t in targets] loss_dict model(images, targets) loss sum(loss_dict.values()) optimizer.zero_grad() loss.backward() clip_grad_norm_(model.parameters(), max_norm10.0) optimizer.step() if (epoch 1) % 10 0: torch.save(model.state_dict(), fcheckpoints/epoch_{epoch1:03d}.pth)clip_grad_norm_这一步很多人会漏掉。口罩识别任务里标签框有时标注得不准边界样本的梯度会异常大不做梯度裁剪轻则 loss 震荡重则训练直接发散。max_norm10.0是经验值。4. 测试视频推理实战与评估方法4.1 从视频帧到检测结果的完整流程源码包里提供了现成测试视频output202212031535.avi和output202212041735.avi运行推理脚本即可输出带检测框的视频。核心逻辑分为三步读取视频帧、缩放至模型输入尺寸通常是 416x416、执行前向推理并解码结果。加载权重时需要注意torch.load的映射位置import cv2 import torch import numpy as np from models.yolo import YOLOv4 device torch.device(cuda if torch.cuda.is_available() else cpu) model YOLOv4(num_classes2, anchorsANCHORS).to(device) model.load_state_dict(torch.load(weights/best.pth, map_locationdevice)) model.eval() cap cv2.VideoCapture(output202212031535.avi) writer None while True: ret, frame cap.read() if not ret: break img, ratio, (pad_w, pad_h) letterbox(frame, (416, 416)) img_tensor torch.from_numpy(img.transpose(2, 0, 1)).float().unsqueeze(0) / 255.0 with torch.no_grad(): preds model(img_tensor.to(device)) boxes, scores, classes post_process(preds, conf_thres0.45, iou_thres0.45) # 坐标还原到原图并绘制 for box, score, cls in zip(boxes, scores, classes): x1, y1, x2, y2 restore_box(box, ratio, pad_w, pad_h) label fmask:{score:.2f} if cls 0 else fno-mask:{score:.2f} cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2) cv2.putText(frame, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2) if writer is None: writer cv2.VideoWriter(output_mask.avi, cv2.VideoWriter_fourcc(*XVID), 30, frame.shape[:2][::-1]) writer.write(frame) cv2.imshow(mask detection, frame) if cv2.waitKey(1) 0xFF ord(q): break cap.release() writer.release() cv2.destroyAllWindows()这段代码里letterbox是关键预处理步骤。它把原始视频帧等比缩放后填充到 416x416而不是直接拉伸原因是直接拉伸会改变人脸宽高比导致锚框回归不准确。restore_box时要减去 padding 再除以缩放比例否则画框位置会偏移。conf_thres0.45表示置信度低于 45% 的检测框直接丢弃iou_thres0.45是 NMS 去重阈值两个参数直接影响误检率和漏检率。4.2 评估指标mAP、Precision、Recall 与 FPS模型训练完成后不能只看 loss 曲线下降就认为收敛。口罩识别场景建议同时关注三个指标mAP0.5IoU 阈值 0.5 下的平均精度低于 85% 时说明模型仍存在明显的定位偏差Recall漏检率直接反映在“人没戴口罩但没被框出来”这类安全场景优先级最高FPS视频推理场景通常要求不低于 25低于此值实时性不可用源码包的utils/utils.py通常内置了 mAP 计算函数或者可以使用独立评估脚本迭代验证集。评估时注意conf_thres设为 0.001这是 COCO 评估惯例目的是测试模型在极端低阈值下的潜力而不是只看某个阈值下的表现。4.3 类别不均衡与误检的常见调参路径测试视频里如果出现“口罩框跳变”即同一张脸的检测框在相邻帧之间大幅移动通常不是模型问题而是conf_thres设置过低导致低置信度框参与了 NMS。另一个常见问题是把“手遮脸”或“围巾遮脸”误判成口罩。这类误检在数据层面需要补充负样本在推理层面可以抬高conf_thres到 0.5 以上。如果只做演示直接调阈值最快如果要提升模型上限需要回训练集增补场景。5. 进阶技巧自适应阈值联动与 TorchScript 导出5.1 动态置信度阈值而非固定值固定conf_thres在单一视频上表现良好但换一个光线条件差或人多的视频就会失效。更稳妥的做法是做阈值联动核心思路是当检测框数量为 0 且画面中确实有人认为存在时逐步降低置信度阈值当误检框数量过多时自动抬升阈值。这种策略在监控场景中尤其实用因为监控摄像头视角固定背景先验稳定各帧之间阈值变化不会太剧烈。实现上只需维护一个长度为 30 的滑动窗口记录最近帧的框数量和平均置信度动态调整下一帧的阈值即可。5.2 导出 TorchScript 提升推理效率PyTorch 模型直接推理时Python 解释器开销占比不低。将模型导出为 TorchScript 后可以脱离 Python 依赖C 端直接加载同时享受图优化带来的加速model.eval() example_input torch.randn(1, 3, 416, 416) traced_model torch.jit.trace(model, example_input) traced_model.save(mask_detector_ts.pt)导出后建议用torch.jit.load验证一遍输出与原始 PyTorch 模型的差值两者输出张量的最大绝对误差应小于 1e-4。超过这个量级说明模型中存在动态控制流或不可追踪操作需要改用torch.jit.script的方式手动标记。此外导出后的模型还能配合 ONNX Runtime 或 TensorRT 做进一步部署优化但那是另一个话题了。对于毕设阶段TorchScript 导出一次再配合 OpenCV 做视频读写整个推理链路就能脱离训练环境独立运行这也是把项目交付出去最稳妥的形态。本文还有配套的精品资源点击获取
返回列表