基于Swin Transformer的YOLOv5安全帽佩戴检测方法

YOLOv5 helmet wearing detection method based on Swin Transformer

摘要:

针对目前施工现场的安全帽检测方法存在遮挡目标检测难度大、误检漏检率高的问题,提出一种改进YOLOv5的安全帽检测方法；首先,使用K-means++聚类算法重新设计匹配安全帽数据集的先验锚框尺寸；其次,使用Swin Transformer作为YOLOv5的骨干网络来提取特征,基于可移位窗口的多头自注意力机制能建模不同空间位置特征之间的依赖关系,有效地捕获全局上下文信息,具有更好的特征提取能力；再次,提出C3-Ghost模块,基于Ghost Bottleneck对YOLOv5的C3模块进行改进, 旨在通过低成本的操作生成更多有价值的冗余特征图,有效减少模型参数和计算复杂度；最后,基于双向特征金字塔网络跨尺度特征融合的结构优势提出新型跨尺度特征融合模块,更好地适应不同尺度的目标检测任务；实验结果表明,与原始YOLOv5相比,改进的YOLOv5在安全帽检测任务上的mAP@.5:.95指标提升了2.3%,满足复杂施工场景下安全帽佩戴检测的准确率要求。

Abstract:

Aiming at the problems of difficult detection of occluded objects and high false detection and missed detection rate in the current helmet detection methods on construction sites, an improved YOLOv5 helmet detection method is proposed in this paper. First, use the K-means++ clustering algorithm to redesign the prior anchor box size to match the helmet dataset. Second, Swin Transformer is used as the backbone network of YOLOv5 to extract features. The multi-head self-attention mechanism based on shiftable windows can model the dependencies between different spatial location features, effectively capture global context information, and have better Feature extraction capability. Third, the C3-Ghost module is proposed, which improves the C3 module of YOLOv5 based on Ghost Bottleneck, aiming to generate more valuable redundant feature maps through low-cost operations, effectively reducing model parameters and computational complexity. Fourth, a new feature fusion module is proposed based on the structural advantages of cross-scale feature fusion of bidirectional feature pyramid network, which can better adapt to target detection tasks of different scales. The experimental results show that compared with the original YOLOv5, the mAP@.5:.95 index of the improved YOLOv5 on the helmet detection task is improved by 2.3 percentage points, which meets the accuracy requirements of helmet wearing detection in complex construction scenarios.