基于改进YOLOv5的工业安全帽检测

YOLO系列算法是目前计算机视觉目标检测领域主流的算法模型,其中以YOLOv5为代表的算法往往具有更快的检测速度和更高的准确率。由于工业施工现场受到例如光照、遮挡等复杂因素的影响,现有的检测算法对于小目标的检测精度不佳,存在漏检、错检等问题。鉴于此,提出一种改进的YOLOv5安全帽检测算法。算法的改进主要为两方面,一方面对YOLOv5的三个预测输出层分别加入三种不同的自注意力机制,对大、中、小三个预测输出层之前分别加入SKNet模块、CA模块、ECA模块用以增强模型对于中小目标检测的鲁棒性,加入通道和空间的特征信息使得模型在预测中小目标时专注于被检测目标,同时在每一个模块引入残差连接,提高训练速度,有效解决因为引入自注意力机制造成的梯度消失问题;另一方面改进原来预测边界框的损失函数,采用DIoU损失函数加快训练的速度,提高了检测精度。在开源的数据集上进行实验验证,实验结果显示改进后的YOLOv5模型对比于改进之前的mAP值提升了1.6%。