改进YOLOv5s的X光安检图像违禁物品检测算法

向 娇,李国权,吴 建,林金朝

(1.重庆邮电大学 通信与信息工程学院,重庆 400065;2.光电信息感测与传输技术重庆市重点实验室,重庆 400065)

0 引 言

X光行李安检是维护公共交通安全的重要手段,目前安检员通过肉眼对X光图像进行识别的方法效率较低且易出现误检、漏检的情况,因此,需要一种更高效、更精确的违禁物品自动检测方法。

随着深度学习在各领域的快速发展,其在X光安检违禁物品识别领域也有了相应的尝试。目前,基于深度学习的违禁品自动识别可以分为违禁物品的自动分类[1-2]、违禁物品的自动检测[3-4]以及违禁物品的自动分割[5]三个方面。文献[1]最早将卷积神经网络(convolutional neural networks,CNN)通过迁移学习[6]的方式应用于X光安检违禁物品的自动分类中;文献[2]利用生成对抗网络技术扩充有限的安检数据集,从而提升安检违禁物品的识别精度;Kim等[3]通过设计一种基于U-net的O-Net结构来进行违禁物品的自动检测;Miao等[4]提出了一种类平衡层级精细化(class-balanced hierarchical refinement,CHR)模型来解决违禁品自动检测时正负样本间类不平衡的问题;Xu等[5]通过在CNN中引入注意力机制来实现安检违禁品的自动分割。虽然基于深度学习的违禁品自动识别技术已经开始被研究,但由于X光图像不同于自然光图像,再加上物品摆放的随意性,在透视性下目标特征不易被学习,导致违禁物品的检测速度达不到实际应用的要求且检测精度仍有待进一步提升。

近年来,One-Stage目标检测算法因其简单的结构和优越的性能受到了广泛的关注。其中,YOLO (you only look once)[7-9]是一系列端到端的目标检测算法的集合,具有检测速度快的特点和优势。由Ultralytics团队近期开源的YOLOv5算法[9]最大限度地兼顾了实时性和准确性,在实时违禁品检测上也有很大的应用潜力。虽然YOLOv5已被应用于工业领域,但至今仍有众多学者对其进行研究[10-12],文献[10]和文献[11]分别研究了基于YOLOv5的垃圾分类和无人机红绿蓝(red green blue,RGB)图像中玉米流苏穗的检测,文献[12]研究了基于YOLOv5的雾驾驶场景目标检测。

YOLOv5s是YOLOv5系列中最小的网络,本文以YOLOv5s作为基础模型,提出了一种改进的X光安检图像违禁物品检测方法,在满足违禁物品自动检测对实时性要求的同时,提高检测精度。在YOLOv5s的主干网络中设计并引入了重参数模块(reparameter block,Rep Block),在3×3卷积处构造一个平行的1×1卷积分支以协助主干网络在训练阶段提取更丰富的特征,推理阶段将1×1分支融入3×3分支,在不增加推理时间的同时提高检测精度。另外,在YOLOv5s颈部的路径聚合网络(path aggregation network,PAN)[13]中插入2个压缩-激励模块(squeeze and excitation block,SE Block)[14],在不增加检测时间的前提下改善了算法对违禁物品的检测效果。

1 相关工作

1.1 X光安检图像数据集

X光在安全检查任务中显示出了其强大的能力,然而可用于研究的X光违禁物品图像数据集仍然较少。GDXray[15]包含了19 407张图片,但只有少数图像(600张)包含三类违禁物品:枪、飞镖和剃须刀片,并且所有图像都是灰度图像,背景简单,与复杂的现实场景差异较大;OPIXray[16]包含了8 885张X光违禁物品图像,具有不同重叠等级和比例,但只包含一类违禁物品(不同形状的刀);SIXray[4]由8 929张具有多个类别的违禁物品图像组成,图像背景比较复杂,危险品间随意堆叠伴有遮挡,更加符合实际情况,所以本文选择SIXray作为实验数据集。图1为该数据集中一些违禁物品的示例,有些违禁物品与行李内其他的合法货物相互重叠。

图1 SIXray数据集中不同的违禁物品Fig.1 Different prohibited items in SIXray dataset

1.2 YOLOv5s算法

YOLOv5s算法由输入、主干网络、颈部和检测头4部分组成,如图2所示。输入端采用Mosaic数据增强方法[8]、自适应计算边界框和缩放图像,丰富输入数据的多样性;主干网络部分使用Focus和跨阶段局部(cross stage partial,CSP)模块,其中,CSP结构有利于提高网络特征学习能力;颈部采用特征金字塔(feature pyramid networks,FPN)加PAN结构,FPN通过向上采样增强语义传播,PAN使用向下采样增强特征定位;检测头部分使用广义交并比损失作为边界框的损失函数,使用非极大值抑制(non maximum suppression,NMS)选择边框。

图2 YOLOv5s网络结构图Fig.2 Original YOLOv5s network

本文分别对YOLOv5s算法的主干网络和颈部PAN结构进行改进,以生成新的网络,提升检测性能。

2 基于重参数的YOLOv5s违禁物品检测算法

本文通过改进YOLOv5s网络结构来提升算法对安检违禁物品的检测精度,通过设计Rep Block来提高主干网络的特征提取能力,且不增加推理时间。在颈部PAN中引入了2个SE Block,从而使网络提取到更多的特征信息。改进后的网络结构如图3所示。

图3 改进的YOLOv5s算法模型结构图Fig.3 Improved YOLOv5s network structure diagram

在违禁物品检测问题中,对YOLOv5s进行了2部分改进:①将Rep结构引入YOLOv5s算法的主干网络当中,得到了由一系列Rep Block和C3模块组成的升级主干网络。其中,Rep Block的构造和使用在2.1节中有详细介绍;②进一步调整YOLOv5的PAN结构,将2个SE Block分别插入PAN的上一检测层和下一检测层之间,得到升级的PAN网络。其中,SE Block的详细结构见2.2节。改进后的新算法不仅能丰富其主干网络的特征性、改善模型性能、提升检测效果,还能对其颈部PAN加强信息提炼,且不增加额外的推理时间,改进后的算法结构如图3所示。其中,Focus模块对图片进行切片操作,使其输入通道扩充4倍,即操作后的图片从原先的RGB三通道变成了12个通道,进一步经过卷积操作,得到没有信息丢失的二倍下采样特征图。Conv模块封装了卷积层、批归一化(batch normalization,BN)层以及SiLU激活函数。C3模块的结构和作用与BottleneckCSP基本相同,但其浮点运算数更低,运行的速度更快。空间金字塔池化(spatial pyramid pooling,SPP)模块拼接不同尺寸的最大池化结果,实现了局部特征和全局特征相融合。UpSample为上采样层,其内插值方法将图像放大到原来的2倍。通过检测头中的3个Conv[1,1]得到最终输出的特征图。

2.1 基于重参数思想设计的Rep Block

为了能提取更丰富的网络特征,提升网络检测性能,研究者们设计了许多新颖的多分支结构。新颖的组件固然能提升精度,但多分支结构带来的问题是难以应用和自定义,增加显存消耗,对推理过程不利。因此,本文利用重参数化思想设计Rep Block来提升模型精度,并通过解耦其训练阶段与测试阶段来降低对推理速度的影响。此前,已有学者对重参数进行了研究[17-19],包括构建视频去噪架构、CNN核心框架和剪枝等研究。

设本文构造的Rep Block参数如(1)式所示,即2个经过BN的平行卷积分支相加。Rep Block产生的信息流表示为y=f(x)+g(x),其中,f(x)、g(x)分别是由3×3核和1×1核实现的卷积分支,且2个卷积分支均经过了批归一化操作。

Rep(3×3)=BN(3×3)+BN(1×1)

(1)

Rep Block结构及其结构重参数化过程如图4所示,对每个进行3×3卷积,在训练阶段给构造并行的1×1卷积分支,并各自经过归一化操作后相加。在推理阶段将1×1分支融合进3×3分支,得到一个3×3卷积分支,减去另一平行分支结构,从而可以提升卷积网络的性能,不影响网络检测效率。

图4 Rep Block结构及其结构重参数化过程(分支融合)Fig.4 Rep Block structure and its structural re-parameterization process (branch fusion)

在Rep Block构造的基础上,还可以进一步基于RepVGG[19]的思想将多分支模块转为单分支。模型的转换(即多分支融合)是在训练完成后进行,包含如下2个步骤。

步骤1将每个分支中的卷积层、BN层进行融合。直接将卷积结果代入BN公式中,见图4中的蓝色箭头,则输出可表示为

M(2)=bn(W(3)*M(1),μ(3),σ(3),γ(3),β(3))+

bn(W(1)*M(1),μ(1),σ(1),γ(1),β(1))

(2)

(2)式中:W(3)∈RC1×C2×3×3和W(1)∈RC2×C1分别表示3×3和1×1卷积层的卷积核;C1,C2代表输入、输出通道数;μ(3),σ(3),γ(3),β(3)分别表示3×3卷积后BN层的累积均值、标准差、缩放因子和偏差项;μ(1),σ(1),γ(1),β(1)对应1×1卷积后BN层的累积均值、标准差、缩放因子和偏差项;输入、输出分别表示为M(1)∈RN×C1×H1×W1,M(2)∈RN×C2×H2×W2;*代表卷积运算。将参数代入(2)式,可得

(3)

(3)式中:bn是推理阶段的批量标准化函数;i∈[1,C2]。对(3)式进一步化简,可得带有偏差项的卷积层。以{W′,b′}形式表示{W,μ,σ,γ,β}转化后得到的卷积核和偏差项,则有

(4)

步骤2对3×3卷积和1×1卷积进行融合,即图4中绿色箭头的步骤。2个偏差项相加得到融合偏差项;将1×1卷积核用0填充为3×3卷积核后,再与原3×3卷积核相加,得到融合卷积核,如图5所示。

图5 普通3×3卷积核与填充卷积核相加得到融合卷积核Fig.5 Ordinary 3×3 convolution kernel and filled convolution kernel are added to obtain a fusion convolution kernel

Conv(x,W1(3))+Conv(x,W2(3))=

Conv(x,W1(3)+W2(3))

(5)

2.2 颈部PAN的改进

文献[20]针对多尺度的待检测图像,根据其骨干网络特征提取的特点,结合加权处理优化颈部网络特征融合,并在颈部网络添加中间层用于融合过渡,实现特征跨尺度连接,丰富全局信息。本文通过在YOLOv5s颈部的PAN中加载SE Blcok[14]来提升模型性能。SE Blcok通过对特征通道间的相关性进行建模,筛选出针对通道的注意力,把重要的特征进行强化来提升准确率。

SE Blcok在卷积层的使用如图6所示,SE Blcok主要包含压缩(squeeze)和激励(excitation)两部分。第1步为压缩阶段,通过全局平均池化,使输入W×H×C的特征图被压缩为1×1×C,这个压缩过的特征图具有全局感受野;第2步为激励阶段,由神经元个数分别为C×r和C的2个全连接层组成,其中,r是一个缩放参数,调节这个参数可减少通道个数从而降低计算量。

图6 SE Blcok在卷积层的使用Fig.6 Use of SE Block in convolution layer

3 实验及结果分析

3.1 数据集

将本文算法在公共数据集SIXray上进行实验,该数据集收集了8 929张带注释的图像。与其他数据集相比,SIXray类别较多、数据量相对较大。将数据集随机分为3部分,其中,验证集和测试集各占20%(1 781张),其余(5 367张)为训练集,比例约为1∶1∶4。本文在实验中去除了对剪刀类物品的检测,因为该样本的数量太少,类间数据量不平衡。数据集中各个类别的详细分布如表1所示。此外,数据集中的许多图像包含多个违禁物品,因此,物品总数远远高于图像的数量。

表1 SIXray数据集中每个类别的分布Tab.1 Distribution of each category in SIXray dataset

3.2 评价指标

检测器的检测性能评估需要同时考虑准确率(Accuracy)和召回率(Recall)。目标检测中使用IoU=0.5时的平均精度均值(mean average precision,mAP)、宏精确率(macro precision,MP)、宏召回率(macro recall,MR)和宏F1(macro-F1,MF1)等来评价网络模型的性能。准确率的定义为(6)式,召回率的定义为(7)式。其中,TTP、TTN、FFP和FFN分别代表真阳性、真阴性、假阳性和假阴性。

(6)

(7)

平均精度 (average precision,AP)是通过结合准确率和召回率而得到,用来评价模型检测单个类别的精度。mAP测量模型检测所有类别的精度,通过求所有类AAP平均值得到,其定义如(8)式。F1分数是准确率与召回率的加权平均值,定义如(9)式,其值越大表明效果越好。

(8)

(9)

同平均精度均值类似,宏精确率、宏召回率和宏F1是通过分别求所有类别准确率、召回率和F1分数的平均值而得到。此外,混淆矩阵也可用于辅助结果分析。

3.3 实验结果分析

实验的硬件配置为Core(TM) i9-10920X处理器,GeForce RTX 3090显卡,软件配置为torch 1.8.0。在训练过程中,使用带动量的随机梯度下降法(stochastic gradient descent,SGD)优化网络参数,迭代次数(epochs)的值设置为200。输入图像大小为640×640,批量大小(batch size)为64。控制其他参数一致的情况下对多种模型进行实验。

本文一共展示了4种模型的实验结果:原始YOLOv5s算法[9]、使用Rep Block的改进YOLOv5s的主干网络算法(下文称为Rep-YOLOv5s)、在颈部PAN中插入SE Block的改进算法(下文称为SE-YOLOv5s)、2种改进叠加的算法(下文称为RepSE-YOLOv5s)。

图7列出了不同算法的混淆矩阵,其中,横纵坐标的0,1,2,3,4分别对应违禁物品手枪,刀,板手,钳子,锤子。对角线值表示真阳性率(true positive rate,TPR),每列非对角线值之和表示一类的假阴性率(false negative rate,FNR)。由图7可知,RepSE-YOLOv5s算法混淆矩阵的对角线响应比其他矩阵的平均值高,混淆矩阵更好的整体分布表明本文算法可以更准确地识别违禁物品。

图7 不同算法下各类违禁物品对应的混淆矩阵Fig.7 Confusion matrix corresponding to various prohibited items with different algorithms

相比之下,算法对扳手类物品的检测效果提升最大。但总体而言,刀这类对象的真阳性率最低。原因可能是数据集中刀这类对象的具体特征并不单一,例如刀类别包含了宽厚的菜刀、细长的直刀、小型的工具刀。

表2展示了4种算法在SIXray数据集上的整体性能对比。对比原始YOLOv5s算法与Rep-YOLOv5s算法,表2中数据显示,Rep-YOLOv5s算法的mAP比原始网络提高了1.6个百分点,其余3项评价指标同样均有提升,每个类别的检测精度也均有提升,这表明Rep Block增强了主干网络的特征提取能力、对算法检测性能的提高有明显的帮助。通过对比原始YOLOv5s算法与SE-YOLOv5s算法可知,SE Block对网络加强特征提取部分的改进同样也改善了算法整体的检测性能,尤其是在MR方面提升了2.5个百分点。

表2 检测算法的性能比较Tab.2 Comparison of performance with different detection algorithms %

综合两部分改进的RepSE-YOLOv5s算法在mAP指标上分别比原始YOLOv5s、Rep-YOLOv5s、SE-YOLOv5s高2.6、1.0、1.4个百分点。在宏精确率(mAP)、宏召回率(MR)和宏F1(MF1)方面,与原始网络和另外2个网络相比也具有更显著的优势。此外,表2中还显示每个类别的检测精度均有提升,尤其是扳手类物品,检测精度由86.4%达到了91.3%,提高了4.9个百分点,这是因为新的算法提取到的小目标特征相对于原来更为丰富,因此,检测精度提升较大。这些评估表明,RepSE-YOLOv5s算法能够更准确地检测X光安检图像中所有类别的违禁物品,有进一步应用于实际场景的潜力。

此外,本文还比较了原始YOLOv5s算法、Rep-YOLOv5s算法、SE-YOLOv5s算法、RepSE-YOLOv5s算法的参数数量(×106)、模型大小(MByte)和检测单张图像所需时间(ms),如表3所示。检测单张图像所需时间是在图形处理器(graphics processing unit,GPU)上对测试集中的数据进行测试的结果,包含数据预处理、模型推理、后处理以及NMS。平均预处理时间为0.1 ms、平均NMS的时间为0.8 ms每张图。

表3 不同算法用时比较Tab.3 Comparison of time spent on different algorithms

由表3可知,RepSE-YOLOv5s算法参数数量仅仅增加了0.28%,模型大小增加了2.82%(0.4 MByte),并且检测时间几乎不变。

表4给出了本文算法与其他5种算法的对比,包括用时和精度的比较。由表4中结果可知,本文算法不仅用时远远低于其他几种算法,平均精度也均高于其他几种算法。特别地,精度比InceptionV3高出了38.7个百分点。

表4 本文算法与其他算法精度及用时的比较Tab.4 Comparison of the accuracy and time spent of the algorithm in this paper with other algorithms

图8展示了示例检测结果。其中,图8a为原始的YOLOv5s算法检测结果图,图8b为RepSE-YOLOv5s算法检测结果图。图8b中红色箭头所指处,改进的算法在复杂的背景中检测到了一把细长的刀,而原始算法却出现了漏检。由此可得,改进后的算法改善了一些背景复杂难以识别物品的检测效果。

4 结束语

针对目前X光安检图像检测精度不够高的问题,本文将模型小、检测速度快的YOLOv5s算法应用于X光安检图像的违禁物品检测,提出了一种RepSE-YOLOv5s检测算法,在提高检测精度的同时未增加推理时间。首先,利用重参数思想设计了一种Rep Block以丰富YOLOv5s算法主干网络的特征性。然后,在YOLOv5s算法颈部PAN中插入2个SE Block,改善了算法对违禁物品的检测效果。最后,在SIXray数据集上进行了实验,并对比分析了4种不同算法模型的违禁物品检测性能,结果显示,新算法的平均精度均值、宏精确率、宏召回率和宏F1分别比原始算法提升了2.6、2.0、4.0和3.0个百分点,同时检测速度保持为2.6 ms每张图,在不增加检测时间的基础上,提升了违禁物品检测的准确率。此外,本文算法还与其他5种算法进行了对比,结果显示,本文算法不仅用时远远低于另外几种算法,平均精度也远高于其他几种算法。特别地,比InceptionV3的精度高出了38.7个百分点。