基于改進YOLOv5s的X光圖像危險品檢測

張康佳, 張鵬偉, 陳景霞, 龍閔翔, 林文濤

(陜西科技大學 電子信息與人工智能學院, 陜西 西安 710021)

0 引言

公共安全一直是人類關注的熱點問題,為了預防危及公共安全的事件發生,機場、火車站等各個重要場所都配有X光安檢機和專門的安檢人員,但是由于X光圖像背景復雜,人工安檢不可避免會發生錯檢、漏檢的情況.

深度學習在近年來取得了巨大的進步,尤其是在圖像識別和目標檢測方面.目前常用的目標檢測方法分為兩大類:一是以R-CNN系列[1-4]為代表的兩階段方法,將定位任務和分類任務分成兩個階段去處理;二是以YOLO(You Only Look Once)系列[5-8]和SSD[9](Single shot multibox detector)為代表的一階段方法,直接得到分類預測和位置坐標信息.而X光圖像危險品檢測屬于目標檢測的下游分支,同樣可以應用兩階段方法和一階段方法展開研究.

在兩階段法目標檢測方面,Akcay等[10]探討了傳統的基于滑動窗口的卷積神經網絡(Convolutional Neural Networks, CNN)和基于區域的全卷積網絡(Region-based Fully Convolutional Networks,R-FCN)在X射線圖像檢測的適用性,并將R-FCN和殘差網絡相結合在ImageNet數據集上進行兩類槍支檢測,均值平均精度(mean Average Precision,mAP)達到了93.6%的精度.Steitz等[11]針對多視角X射線圖像數據,引入一種基于Faster R-CNN的多視圖檢測模型,利用聚合的多視圖特征,在自定義3D數據集上進行最終分類,取得了不錯的效果.Gaus等[12]研究比較了Faster R-CNN、Mask R-CNN和RetinaNet模型在對不同幾何形狀、圖像分辨率和顏色分布的X射線圖像檢測的可遷移性,在二分類和三分類X光圖像數據集上進行目標檢測的mAP性能分別達到了88%和86%.Liu等[13]根據被檢測物體與背景顏色差異,從X射線圖像中分離出被檢測物體,然后將其送到R-CNN網絡中學習,最終在自建的六分類X光危險品圖像數據集上進行目標檢測的mAP達到77%.Zhang等[14]提出了一種基于X射線的分類區域神經網絡(XMC R-CNN),通過使用X射線物質分類器算法、有機分離算法和無機剝離算法解決X射線圖像中的行李重疊問題,最終對槍、刀檢測的查全率分別達到了96.5%和95.8%.Bhowmik等[15]使用內容感知重組特征(CARAFE),對6類X射線危險品圖像進行檢測,取得了70%的平均精度.上述基于兩階段的X光圖像危險品檢測方法,雖然對X光圖像取得了較高的檢測精度,但是所使用的危險品數據集類別較少,檢測速度低下.無法滿足實際情況下的X光圖像危險品檢測多樣性鑒別和實時性的檢測要求.

在一階段法目標檢測方面,Liu等[16]采用YOLO9000網絡針對X光圖像中剪刀、氣溶膠特征進行了深入研究,最終對這兩類的危險品進行檢測的平均查準率和查全率分別達到了94.5%和92.6%.Galvez等[17]比較了YOLOv3模型在IEDXray數據集上使用遷移學習和從頭訓練方法對X射線圖像中的簡易爆炸裝置進行檢測的區別,發現在多尺度X光圖像檢測中,從頭訓練的目標檢測精度能夠達到52.40%,而用遷移學習的方法只能達到29.17%,說明從0開始訓練的YOLOv3在X光圖像危險品檢測方面優于遷移學習.Wei等[18]在SSD網絡上通過添加額外的卷積層,并采用多任務遷移學習方法,在GDXray數據集上進行了3類目標檢測,取得了91.5%的平均精度.Qiao等[19]為了提高小規模違禁品的檢測精度,在SSD網絡上增加了特征融合模塊和非對稱卷積模塊,將SSD模型進行目標檢測的mAP提高了2.48%.Zhou等[20]在YOLOv4網絡上引入可變性卷積,并使用GHM損失優化損失函數,在Sixray數據集上6類目標檢測的平均精度達到了91.4%.上述基于一階段的目標檢測方法,雖然在X光危險品圖像檢測中取得了較高的精度,但是模型相對較大,無法做到實時部署,并且所使用X光圖像數據集中危險品類別較少,仍然無法滿足實際應用需求.

針對上述X射線圖像危險品檢測實際存在的問題,本文在YOLOv5s模型的基礎上,引入了卷積注意力機制[21](Convolutional Block Attention Module,CBAM),并將注意力機制與特征金字塔[22](Feature Pyramid Networks for Object Detection,FPN)結構聯合使用進行X射線圖像危險品檢測,使模型在特征學習與融合的過程中選擇性強調危險品相關特征,并抑制背景干擾特征.同時在模型的檢測頭部引入解耦頭進行優化,使分類和定位任務分開處理,并使用了更適合分類任務的全連接層來處理分類任務,而定位任務仍然使用卷積層來完成.通過這種方式來提高YOLOv5s模型在危險品檢測中的精度.

1 YOLOv5s模型

YOLOv5s是一種用于單階段目標檢測的模型,它由輸入端、基準網絡、Neck網絡和head輸出端組成.

(1)輸入端:輸入YOLOv5s模型的圖像大小一般為618*618或者416*416,本文輸入的X射線危險品圖像大小設置為416*416,通道數為3.該部分對輸入的危險品圖像進行一系列預處理操作,首先將危險品圖像隨機縮放到網絡指定的輸入大小,再進行歸一化操作,通過對其進行隨機排布,以進一步增強圖像的特征表達能力.

(2)主干網絡: Focus和CSP是主干網絡的兩個重要組成部分.其中,Focus部分通過slice切片操作,對輸入的危險品圖像進行裁剪,將其寬高縮小一半,變為了208*208.同時,將危險品圖像的通道數乘4,變為12,得到面積減少為原來的四分之一,通道數變為原來的四倍的特征映射.YOLOv5s的CSP結構包括三個部分:第一部分是由卷積、歸一化和激活函數組成的CBS(Conv+BN+SiLU)模塊;第二個部分是由多個殘差組件組成的殘差模塊;第三部分是單個卷積層.其作用是通過將特征圖按照通道維度均分為兩部分,其中一部分經過殘差模塊之后,再通過跨層連接將兩部分再次合并,以達到減少計算量的同時,提高準確率.

(3)Neck網絡:Neck位于主干網絡和head輸出端之間,YOLOv5s在Neck部分采用FPN+PAN(Path Aggregation Network)結構.FPN部分使網絡輸出特征圖按從小到大結構排列,進而將更多的語義信息傳遞下來.而PAN結構使網絡輸出特征圖按從大到小的結構排列,進而將更多的定位信息傳遞上去.通過將FPN+PAN相結合,將低層的強定位信息與高層的強語義信息進行融合,得到更加有用的特征信息.

(4)Head輸出端:該部分從三個不同尺度對目標進行預測,最后對預測的結果進行非極大值抑制,進一步精簡預測結果.

2 對YOLOv5s模型的改進

針對YOLOv5s模型特征融合時背景信息的干擾和檢測頭部耦合問題,本文在YOLOv5s模型的基礎上所作改進如下,在FPN特征金字塔處,引入注意力機制模塊;在檢測頭部分,對模型的預測頭部進行解耦.將其改進后的模型命名為att_decouple_YOLOv5s,其結構如圖1所示.

2.1 注意力機制的引入

在模型對危險品檢測中,雖然FPN的自頂向下的特征金字塔結構,能夠更好地將強語義信息傳遞下去,與淺層網絡的所傳遞的特征信息進行特征融合.但是淺層網絡所傳遞的特征信息并非都是有用的,其中包含許多背景干擾信息.所以為了減少背景干擾信息的傳遞,本文在淺層網絡與FPN 信息融合處,結合注意力機制進行研究.

注意力機制在計算機視覺中的應用主要是三種,分別是通道域、空間域和混合域.本文所使用的卷積注意力模塊(CBMA)則是混合域上的應用,通過使用注意力機制抑制淺層網絡所傳遞的特征信息中的干擾信息,如圖2所示:輸入特征圖X(N×H×W),經過通道注意力機制生成的通道注意力圖Mc(N×1×1),然后將二者相乘,得到特征圖X′(N×H×W),再將生成的特征圖X′輸入到空間注意力模塊,得到空間注意力圖Ms(N×1×1),將X′與Ms相乘,得到最終輸出.

圖2 注意力機制

在圖3中,可以看到通道注意力部分對輸入特征圖的操作過程.輸入特征圖X經過最大池化和平均池化,分別生成兩張N×1×1特征圖:平均池化圖FA和最大池化圖FM.將FA和FM分別送入共享全連接層(MLP),然后對兩種特征進行加和操作,并使用Sigmoid激活函數來強化不同通道的權重.最終得到通道注意力圖,計算過程如下:

圖3 通道注意力機制

Mc(X)=S(MLP(AP(X);MP(X)))=
S(W1(W0(FA))+W1(W0(FM)))

(1)

式(1)中:S為Sigmoid激活函數,AP為平均池化,MP為最大池化,W0和W1分別為共享全連接層的第1層和第2層,c為通道數.

空間注意力模塊專注于位置信息,與通道注意力模塊相輔相成,如圖4所示.

圖4 空間注意力機制

在該模塊,對輸入特征圖X′在通道維度上對特征點分別求平均值和最大值,分別得到關于平均值的空間特征圖SAP和關于最大值的空間特征圖SMP.然后通過拼接操作cat將兩種特征圖進行拼接,再利用7×7的卷積F生成空間注意力圖,最后使用Sigmoid激活函數對不同位置的特征點進行強化或抑制.其計算過程如下:

MS=S(F(cat(AP(X′),MP(X′))))=
S(F(cat(SMP,SAP)))

(2)

在圖1所示的注意力機制部分,本文通過卷積注意力模塊,將淺層網絡傳遞的特征信息分別從空間域和通道域進行不同區域的關注,提高有用信息的比重,并減少背景信息的影響.如圖1的注意力機制部分和圖2所示,輸出的底層特征圖被注意力模塊(CbamBlock)給予不同的權重.通過將其與高層特征相結合,提取出模型的有用特征,從而增強網絡學習的目標特征信息,此外,這種方法基本上不會增加模型的大小,也不會增加訓練和推理成本.

2.2 檢測頭部解耦

在目標檢測中,Song等[23]通過對空間敏感性熱圖進行可視化,發現分類和定位所關注的感興趣區域的不同,分類更關注于顯著性區域信息,定位更關注邊緣區域的信息,由此可以看出分類和定位任務之間的耦合一直存在沖突問題.YOLO系列之前是將分類和定位信息耦合在一起,Ge等[24]提出的YOLOX模型在YOLOv3的基礎上采用無錨anchor-free并對檢測頭部進行解耦.該模型在檢測頭部首先使用1×1的卷積來減少通道維數,然后使用兩個平行的分支,每個分支都包含兩個3×3的conv層,分別來處理分類和定位任務,以達到提升模型性能的效果.但由于本文選用的YOLOv5s進行危險品檢測,其在有錨框的檢測,仍舊使用的是定位和分類耦合在一起的檢測頭,因此本文將YOLOX中的解耦頭應用在YOLOv5s下的有錨檢測中,并且在其基礎上進行了進一步的改進,其具體處理過程如圖5所示.

將模型提取的特征圖通過1×1的卷積降維,可以得到256×H×W的特征圖Y,將得到的特征圖分別輸入三個平行的分支,三個分支分別處理分類任務、置信度得分和定位任務.

在分類任務分支中,對特征圖進行R1(reshape)操作,使其變成H×W×256的特征圖,將得到的特征圖依次經過兩個全連接層f1和f2,對其分類特征進行整合計算,再經過全連接層f3,得到H×W×36的分類特征圖.最后再對其進行R2(reshape)操作,得到特征圖Y1(36×H×W),計算公式如下:

Y1=R2(f3(f2(f1(R1(Y)))))

(3)

在置信度得分的分支中,將其看作是一個軟標簽分類任務,對特征圖Y進行R1(reshape)操作,再經過全連接層f4對置信度特征進行整合計算,最后經過全連接層f5,得到H×W×3的置信度特征圖,對其進行R2(reshape)操作得到特征圖Y2(3×H×W),其計算公式如下:

Y2=R2(f5(f4(R1(Y))))

(4)

在定位任務的分支中,使用卷積C1和卷積C2依次對其定位特征進行整合計算;再經過卷積層C3得到定位特征圖Y3(256×H×W),計算公式如下:

Y3=C3(C2(C1(Y)))

(5)

最后是將三個分支得到的結果按照通道維度進行cat連接操作得到最終的輸出特征圖Y′.其計算公式如下:

Y′=cat(C3,C2,C1)

(6)

本文通過上述方法對YOLOv5s的檢測頭進行解耦,解耦部分如圖1 head(解耦頭)部分和圖5所示,將提取后的特征分別進行定位任務與分類任務處理.根據定位任務和分類任務的特性,為定位任務選擇卷積層進行處理,為分類任務選擇全連接層進行處理,使YOLOv5s模型在X光圖像危險品檢測中的檢測精度有了較大的提升.

3 實驗結果及分析

本文基于NVIDA GeForce GTX 3090顯卡、32 GB內存的 centos和 pytorch框架展開實驗.

3.1 數據集及評價指標

本文實驗所使用的數據集是Wang等[25]在2021年公開的大規模違禁品檢測數據集pidray,其包含了47,677張X射線圖像下違禁品圖像,種類為12類,分別是槍、刀、扳手、鉗子、剪刀、錘子、手銬、警棍、噴霧器、充電寶、打火機和子彈,每個種類在數據集中所占的比例如圖6所示.

圖6 數據集中各類危險品數量

本文使用多種關于目標檢測模型性能評估指標來衡量危險品檢測的準確性,包括預測精度(Precision)、召回率(Recall)、均值平均精確度(mean average precision,mAP).

精度定義為:

(7)

召回率定義為:

(8)

均值平均精確度定義為:

(9)

其中,TP表示正確檢測出危險品的數量,FN代表被判定為背景的危險品數量,FP代表將背景區域判定為危險的數量.

3.2 模型評估

為了驗證本文所選模型算法的優越性,首先用本文改進的att_decouple_YOLOv5s模型在pidray數據集上進行訓練,并將訓練結果與RetinaNet、Faster R-CNN、Mask R-CNN、SSD512和Cascade R-CNN等模型在pidray數據集上的訓練結果進行對比,對比結果如表1所示.在pidray數據集上,本文改進的att_decouple_YOLOv5s模型,相比于RetinaNet、Faster R-CNN 、Mask R-CNN,mAP、SSD512和Cascade R-CNN,其mAP性能分別提高了25.1%、22.4%、20.8%、19%、15.6% .從實驗結果可以看出,相較于目前主流的目標檢測模型,本文改進的att_decouple_YOLOv5s模型更加適用于X光圖像危險品檢測.

表1 各個模型mAP對比

3.3 消融實驗

為了進一步驗證本文所提算法的有效性,將實驗時所有對比模型的輸入圖片大小都統一設置為416*416,batch-size大小設置為32,并通過多組對比實驗驗證本文所采用的各個方法的有效性.

為了驗證本文在YOLOv5s網絡的FPN處引入注意力機制的有效性.首先在YOLOv5s網絡模型的FPN結構中加入注意力機制,將加入注意力機制后的模型命名為att_YOLOv5s,在模型的各個參數與YOLOv5s模型各個參數設置相同的情況下在pidray數據集進行實驗.實驗結果如表2所示,改進的att_YOLOv5s模型相比于YOLOv5s的精度提升了3.4%,而召回率只下降了0.1%,在IoU閾值為0.5的情況下,模型的mAP提升了0.7%,在IoU閾值為0.5到0.95的情況下,mAP只下降了0.1%,表明注意力機制在FPN出的引入有效性.而對比模型的參數量,如表3所示,YOLOv5s參數量為7042489,att_YOLOv5s參數量為7075275,比YOLOv5s模型的參數量只增加了0.46%,表明引入注意力機制基本上不會增加模型體量.由此可見,在YOLOv5s模型的FPN處引入注意力機制對特征提取的有效性.

表3 模型參數量對比

為了驗證本文在YOLOv5s中引入解耦頭以及對解耦頭改進的有效性,本文在YOLOv5s中head預測部分,引入了解耦頭,將引入解耦頭后的模型命名為decouple_YOLOv5s,并且對引入后的解耦頭用全連接層替換定位分支和置信度分支的卷積層,將重新設計解耦頭之后的模型命名為new_decouple_YOLOv5s.仍舊是將改進后模型各個參數與YOLOv5s算法模型各個參數設置相同的情況下進行實驗.實驗結果如表2所示,YOLOv5s加入解耦頭后的decouple_YOLOv5s模型相對于YOLOv5s模型,模型的精度提升了1.1%,召回率提升了1.1%;在IoU閾值為0.5的情況下,mAP上升了1.5%;在IoU閾值為0.5到0.95的情況下,mAP上升了2.1%.而對于參數量來說,如表3所示,decouple_YOLOv5s模型參數量相比于YOLOv5s模型的參數量,只增加1.03倍.證明了在YOLOv5s中引入解耦頭的有效性.而new_decouple_YOLOv5s相對于YOLOv5s,精度提升了2.9%;相對于decouple_YOLOv5s,精度提升了1.8%;new_decouple_YOLOv5s在IoU閾值為0.5,相對于YOLOv5s,mAP提升了1.9%;相對于decouple_YOLOv5s,mAP提升了0.4%;在IoU閾值為0.5到0.95,相對于YOLOv5s,mAP提升了3.3%;相對于decouple_YOLOv5s,mAP提升了1.2%.而對于參數量來說,如表3所示,new_decouple_YOLOv5s參數量相比decouple_YOLOv5s模型的參數量少0.21倍.以上充分說明了對頭部解耦,以及將分類問題用全接層處理的有效性.

為了驗證YOLOv5s同時在FPN處加注意力機制和對頭部解耦的有效性,本文在att_YOLOv5s模型中加入改進的解耦頭,并將其命名為att_decouple_YOLOv5s,參數設置不變,對比實驗結果如表2所示.通過對比綜合評價指標發現att_decouple_YOLOv5s,無論是IoU閾值為0.5,還是在IoU閾值為0.5到0.95,其mAP都高于單獨在FPN處引入注意力機制的att_YOLOv5,或者單獨對頭部解耦的new_decouple_YOLOv5s,進一步驗證了同時在頭部解耦和FPN處引入注意力機制能夠有效提升模型的精度.

在通過觀察最終改進后的模型訓練時mAP的變化曲線,發現在epoch=200時候,模型的檢測

性能仍舊有很大的提升空間.為了尋找最優epoch,本文又分別在其他參數不變的情況下,將模型分別訓練了300 epoch和400 epoch,其模型mAP變化曲線如圖7所示.從圖7能夠看出,200 epoch到300 epoch之間,模型mAP性能mAP有明顯的提升,300 epoch到400 epoch之間,模型mAP性能基本保持不變,可見epoch=300的時模型性能達到最優.

圖7 不同epoch下的mAP的變化曲線

最后,本文從測試圖像中隨機選取一部分危險品圖像進行檢測,檢測結果如圖8所示.

圖8 模型改進前后對比圖

圖8(a)顯示了在 X光圖像背景下對鉗子(Pliers)進行檢測,原YOLOv5s模型所識別出鉗子的置信度低于本文提出的att_decouple_YOLOv5s模型.圖8(b)顯示了在 X光圖像背景下對充電寶(Powerbank)檢測,原YOLOv5s模型將背景區域誤檢為剪刀(Scissors).圖8(c)顯示了在 X光圖像背景下在對槍(gun)進行檢測,原YOLOv5s將槍錯檢為剪刀,而本文提出的att_decouple_YOLOv5s模型不僅檢測出了gun,而且取得了0.88的置信度.雖然本文提出的att_decouple_YOLOv5s模型將與錘子(Hammer)特征相似的背景區域錯誤檢出,但是同時也給出了低置信度的打分.測試結果再一次驗證了本文所提att_decouple_YOLOv5s模型的有效性和先進性.

4 結論

針對X光危險品檢測任務中,其物品擺放復雜,背景復雜,危險品種類繁多,難以識別等問題,本文對經典的目標檢測模型YOLOv5s進行改進,得到了一種新的att_decouple_YOLOv5s模型.首先通過在YOLOv5s網絡的backbone和Neck的特征融合部分引入注意力機制,使淺層特征中有用特征的權重進一步加強,背景特征權重得到減弱,從而提升了特征融合的有效性.同時,在原YOLOv5s的檢測頭部分,通過對檢測頭部解耦,將分類任務和定位任務分開處理,減少了分類和定位因為耦合在一起所產生的沖突問題.在pidray數據集上進行了多組X光圖像危險品檢測的對比實驗.

實驗結果表明,所提的att_decouple_YOLOv5s模型無論是從模型復雜度,還是從模型的檢測精度來說,其各項性能評價指標都比其他同類目標檢測方法有明顯的提升.在實際應用中,X射線圖像中危險品種類繁多,遠遠不止12類.此外,X射線下的圖像重疊嚴重,背景復雜,對X射線圖像危險品的準確檢測有著很大的影響.在未來研究工作中,本課題組將使用多個X射線圖像危險品數據集,對更多類型的危險品進行檢測.同時還將進一步研究更有效的X射線圖像特征學習的模型和方法,緩解物品重疊和復雜背景的干擾,從而進一步提升X光圖像中危險品檢測的性能.