融合注意力机制的安全帽检测算法

张帅帅,田 锦,刘 静

(1.安徽理工大学计算机科学与工程学院,安徽 淮南 232001;2.金陵科技学院网络与通信工程学院,江苏 南京 211169)

工业施工场所安全帽的佩戴问题,一直是人们比较关注的焦点。秉承生命至上的原则,须对安全帽的佩戴情况加大监督力度,但管理人员精力有限,不可能时刻监督未戴安全帽的人员。为了消除工地上此类安全隐患,需要一个可以检测工作人员是否佩戴安全帽的智能装备[1]。

安全帽检测是一种目标检测类问题。目标检测常常应用于人脸识别、实时监控、无人驾驶等方面[2-3],检测算法分为传统的目标检测算法和基于深度学习的目标检测算法。传统的目标检测算法,前期先选出候选框,再经过手工设计的提取器提取特征,最后送入分类器中进行检测。随着卷积神经网络(convolutional neural networks,CNN)的快速发展,利用卷积神经网络可以使网络模型适应现实中的复杂场景,而且利用卷积神经网络对大规模的数据集进行训练,还可以提高模型的泛化能力。随着深度学习的发展,基于深度学习[4]的目标检测算法的检测速度和精确度,都有了很大的提升。目前基于深度学习的安全帽佩戴检测算法一般分为两类,一类是以R-CNN[5]和Fast R-CNN[6]系列算法为代表的两阶段算法,另一类是以SSD[7]和YOLO系列[8-10]为代表的单阶段算法。在单阶段算法中,李明山等[11]通过在 SSD 算法中引入特征融合,提高了检测的准确性,但忽略了远距离的目标。方明等[12]在 YOLOv2结构中采用MobileNet中的轻量化网络,实现了对网络一定程度的压缩,减少了模型的计算量,提升了模型的收敛性能,但对物体的检测准确率较低。丁文龙等[13]在YOLOv3模型中采用K-means++聚类,引入残差模块,获得安全帽检测的先验尺寸,提高了准确率,但对于远距离目标的检测仍有明显不足,存在漏检、误检的情况。王英郦[14]在YOLOv5模型中引入轻量化的CBAM注意力模块,提高了检测精度,但存在误检情况。

基于深度学习的安全帽检测算法,检测准确度和速度优于传统方法,但在检测远距离目标方面存在不足。由于远距离目标特征信息少,因而在网络中捕捉到的特征信息也会减弱,在实际检测中会出现漏检、误检的情况。针对此问题,本文对容易部署的YOLOv5s算法模型进行改进,通过在主干网络Backbone层添加坐标注意力(coordinate attention,CA)模块,增强复杂背景中远距离目标特征的提取能力,从而提高检测能力;再优化损失函数,使用EIoU损失函数提高回归精度和收敛速度。

1 基础网络模型

YOLOv5是由Uitralytics LLC公司发布的一种单阶段目标检测算法,它提供了YOLOv5m、YOLOv5s、YOLOv5x、YOLOv5l四种不同的模型算法,其中YOLOv5s模型在深度和宽度上都是最小的,且计算量较少,因此选用YOLOv5s作为基础网络模型。网络模型包括4部分:输入层、主干网络层、特征融合层和预测层。YOLOv5s的结构如图1所示。

第1部分为输入层(Input),采用Mosaic数据增强的方法,对输入的图片进行数据增强、尺寸大小调整处理,该方法是将4张图片随机地拼凑到一起,提高数据样本的多样性;第2部分为主干网络层(Backbone),采用Focus、C3、SPP[15]等多种特征提取模块,对安全帽佩戴区域进行特征提取,生成目标语义信息;第3部分为特征融合层(Neck),使用FPN[16]+PAN结构,采用自下而上的方法传递位置信息,聚合主干层中不同检测层的参数,在降低网络深度的同时加强安全帽特征信息的融合能力;第4部分为预测层(Prediction),采用3种不同尺度的特征图来检测大目标、中目标和小目标对象,生成边界框用来预测检测的类别,提高模型对安全帽检测的准确度。

2 改进网络模型

2.1 引入注意力模块

由于安全帽检测图片背景复杂,可能会造成特征信息的丢失,为了使模型能更好地关注所需要的目标信息,在主干网络Backbone层整体引入CA模块,将其添加在主干网络层的最后一层,强化重要的特征信息和对特征信息的提取能力。CA模块[17]的结构如图2所示。

图2 CA模块结构图

CA模块能够将横向和纵向的位置信息编码到通道中,使得移动网络能够关注大范围的位置信息但又不会带来过多的计算量,不仅获得了通道间的信息,还考虑了方向相换的位置信息,有助于模型更好地定位和识别目标。它足够灵活和轻量,能够直接插入移动网络的核心结构。CA模块并没有引入2维的全局池化,而是分解成了两个1维的特征编码,捕获2个方向的聚合特征,避免位置信息的损失。这样可以沿着一个空间方向捕捉远程依赖的关系,同时沿着另一个空间方向保留精确的位置信息。CA模块将生成的特征图分别编码为对方向感知和对位置敏感的注意力特征图以及可以应用于输入的特征图,增加对安全帽的关注度,提高泛化能力。

具体地,给定输入X,使用尺寸为(H,1)或(1,W)的池化核,分别沿水平坐标和垂直坐标对每个通道进行编码。因此,高度为h的第c通道的输出可以表示为:

(1)

宽度为w的第c通道的输出可以表示为:

(2)

上述式(1)和式(2)变换后沿水平和垂直方向聚合信息,得到方向感知特征图。和以往只能产生单一特征向量的通道注意力不同,这种转换使注意力模块可以捕获到一个方向的持久信息,也可以捕捉到另一个方向的位置信息,从而帮助网络精准确定所需要的目标。

通过坐标信息的嵌入将获得的两个方向的特征图进行连接操作,并使用1×1卷积变换函数进行变换:

f=δ(F1([Zh,Zw]))

(3)

式中,[Zh,Zw]为沿着空间维数的连接操作,δ为非线性激活函数,f为空间信息在两个方向上的特征映射。用1×1卷积变换函数F1分别将fh和fw变换为具有相同通道数的张量并输入x,得到:

gh=δ(Fh(fh))

(4)

gw=δ(Fw(fw))

(5)

最后输出Y:

(6)

CA模块沿着垂直和水平两个方向同时进行编码,这种编码过程可以使坐标注意力机制更加准确地定位ROI,从而帮助整个模型更好地进行检测。

将CA模块加入C3后面的一层,加入CA模块后的Backbone网络结构如图3所示。在主干网络中添加CA模块后,网络更加轻便,实现了轻量化,计算量也相应地减少,而且可以在众多的信息中注意当前所需要的关键信息,降低对其他信息的关注。这样不仅可以使网络能够分配更多的注意力给安全帽,能够更加准确地对安全帽进行检测,还可以提升对远距离物体的检测能力。

图3 添加CA模块后的Backbone网络结构

2.2 优化损失函数

原网络模型使用的损失函数为CIoU[18],CIoU损失函数反映的并不是宽高与其置信度真实值的差异,因此会影响模型的优化。EIoU损失函数[19]将纵横比的影响因子拆开,分别计算目标框和锚框的宽高,考虑了重叠面积、中心点距离以及宽高边长真实差异问题。EIoU损失函数对宽高预测结果进行惩罚的损失函数为:

(7)

该损失函数包含3个部分:重叠损失LIoU、中心距离损失Ldis和宽高损失Lasp。重叠损失LIoU和中心距离损失Ldis这两部分是CIoU中的方法,宽高损失Lasp使得目标盒的宽度与高度之差相比锚盒要小,使收敛速度变快。式(7)中IoU为边界框与真实边界框的交并比,(b,bgt)为预测边界框与真实边界框的中心点,ρ(·)为欧氏距离计算函数,Cw和Ch分别为覆盖两个Box的最小外接框的宽度和高度。

为显示损失函数的优点,选取IoU、GIoU、CIoU和EIoU 4类损失函数进行对比实验,结果如表1所示。由表1可以看出,EIoU损失函数检测的准确率最高。因此,本文将原有的CIoU更换为EIoU。对于边界框回归任务,EIoU 选择计算宽高的差异值来取代纵横比。使用EIoU损失函数可以加速收敛,解决样本回归中产生的不平衡问题,提高回归精度,进而提高检测的准确率。

表1 损失函数对比结果

3 实验结果与分析

3.1 实验环境

实验环境配置:微处理器CPU Inter Xeon E5-2 620v4@2.10 GHz、GPU NVIDIA Titan Xp@Super 12 GB、Cuda 11.0、PyTorch 1.10.0、Python 3.8。

输入网络训练模型的图片大小为640像素×640像素,batch_size设置为8,利用SGD函数优化参数,动量因子为0.937,权重衰减系数为0.000 5,所用的参照模型全部按照50个epoch训练。测试阶段的交并比设置为0.5,置信度阈值设置为0.5。

3.2 数据集

实验中需要训练的图片来自公开数据集SHWD(safety helmet wearing-dataset)。数据集SHWD共有7 581张图片,包含了不同施工场景和不同拍摄距离的图像。进行数据标注时,采用开源的labelimg工具,对图片中出现的佩戴安全帽的人员标注为hat,未佩戴安全帽的人员标注为 person,标注后保存的格式为PASCAL VOC格式。实验前将PASCAL VOC格式的xml文件转换为txt文件,将图片大小统一调整为640像素×640像素,为后期实验做准备。

3.3 评价指标

使用平均精度(average precision,AP)和均值平均精度(mean average precision,mAP)检测模型的准确性能,使用目标检测的精确率(precision,P)和召回率(recall,R)来计算平均精度,各评价指标的计算如式(8)—式(11)所示:

(8)

(9)

(10)

(11)

3.4 消融实验

为验证本文改进的网络模型的有效性,以YOLOv5s模型为基础进行消融对比实验,结果如表2所示。实验结果显示,2.1节和2.2节所述的改进方法(添加CA模块和使用EIoU损失函数)均对mAP有所提高。结合表2消融实验结果可知,本文改进的网络模型与原模型相比,精确率P提高了1.1个百分点,mAP@0.5提高了1.9个百分点。

表2 消融实验结果 单位:%

3.5 结果分析

为进一步验证改进模型的有效性,将本文算法、YOLOv3、Faster RCNN、引入特征融合的SSD算法[11]、融合轻量化CBAM注意力机制的YOLOv5算法[14]和原YOLOv5s进行实验测试,所得安全帽佩戴检测结果如表3所示。表3显示,改进后的网络模型有更高的检测精度,mAP值达到了94.5%,相较于对比算法提高了1.1～12.3个百分点,hat和person的AP值相较于其他算法均有一定的提升。其中文献[14]person的AP值比原YOLOv5s的低,但其hat的AP值相比原YOLOv5s提高了2.1个百分点,因此mAP值也提高了。对于本文改进的网络模型,所有评估指标都比原YOLOv5s高,训练结果比较理想,验证了本文模型的可行性。

表3 实验结果对比 单位:%

为了更加直观地显示改进后的算法对安全帽佩戴的检测效果,选取部分检测图像进行对比,如图4—图6所示。图4(a)、图5(a)和图6(a)为原YOLOv5s算法的检测结果,图4(b)、图5(b)和图6(b)为改进后算法的检测结果。原YOLOv5s算法结果图中,将不属于hat目标的物体检测成hat的属于误检,未检测出图片中hat目标的属于漏检。

(b)改进的YOLOv5s

(a)YOLOv5s

(b)改进的YOLOv5s

(a)YOLOv5s

(b)改进的YOLOv5s

原YOLOv5s算法存在安全帽漏检、误检的情况,改进的算法在相同的环境下,可以更准确地检测出目标。在图4中,YOLOv5s算法没有检测出右上方远处模糊的安全帽hat目标,而改进的算法可以正确检测出,说明改进的算法对远距离的目标有更好的检测效果。在图5中,YOLOv5s算法没能检测出左上角被遮挡的安全帽目标和旗子旁边佩戴安全帽的工作人员,而改进的算法可以检测出,说明改进的算法对遮挡目标有更好的检测效果。在图6中,YOLOv5s算法把左下角的铁管目标误检成安全帽hat目标,且没有检测出右侧被铁管和铁丝遮挡的安全帽目标,而改进的算法可以准确检测出。以上结果表明,改进的算法能有效改善漏检、误检以及目标检测度不高的情况。

4 结 语

本文提出了一种融合注意力机制的安全帽检测算法,该算法以YOLOv5s模型为基础,通过在主干网络中添加CA模块,加强安全帽特征信息的提取能力,减少其他特征对模型的干扰和影响;同时,使用EIoU损失函数融合更多的边界框信息,优化边界框的回归精度和收敛速度。实验结果表明,改进的算法平均检测精度达到了94.5%,相较于原算法提高了1.9个百分点,有效改善了安全帽检测中的漏检、误检情况,提高了模型的整体检测性能。