基于YOLOv5s的改進實時紅外小目標檢測

谷 雨,張宏宇,彭冬亮

(1.杭州電子科技大學自動化學院,浙江 杭州 310018;2.杭州電子科技大學圣光機聯合學院,浙江 杭州 310018)

1 引 言

紅外成像系統具有全天候、抗干擾能力強、探測距離遠等優勢,因此基于紅外成像的目標檢測技術在軍事偵查、紅外制導、自動駕駛等領域得到了廣泛應用[1]。與可見光圖像不同,紅外圖像分辨率低、背景復雜,目標多呈現為弱小目標狀態,嚴重影響了檢測精度,因此如何提高紅外小目標檢測性能成為亟待解決的問題。

傳統的紅外小目標檢測方法主要有三種[2],基于濾波器的紅外小目標檢測算法思路簡單、計算量小,但其對于復雜背景的抑制較差,檢測精度低;基于人眼視覺系統的檢測方法易于實現,但其檢測精度依賴于分割閾值,有一定局限性;基于矩陣分解的方法對于復雜背景有較高的可靠性,但由于計算復雜,檢測實時性較差。

隨著深度學習理論的發展,基于深度學習的目標檢測取得了遠超傳統方法的性能?；谏疃葘W習的通用目標檢測算法可以分為基于候選框的兩階段檢測算法和基于回歸的單階段檢測算法[3]。直接將上述通用目標檢測算法用于紅外小目標檢測時,由于紅外圖像分辨率低、目標尺寸小、缺乏細節紋理特征的特性,增加了紅外目標的檢測難度,檢測性能會降低,因此研究學者針對深度學習紅外圖像弱小目標檢測進行了一系列優化。Wu等人[4]基于YOLOv3[5]算法,將網絡預測層擴展到4個特征尺度,通過引入GIoU[6]改進了損失函數,提高了檢測性能,在FLIR紅外數據集上平均準確率提高了3.4 %。Zheng等人[7]針對紅外小型無人機目標檢測,設計了一個特征增強模塊以增強“目標特征”,同時將自適應特征融合方法引入特征融合中,以解決跨尺度特征融合中特征表達弱化的問題。Zhao等人[8]結合DenseNet[9]和YOLOv5s[10],將YOLOv5s的部分C3模塊替換為DenseBlock模塊,并且在主干網絡中加入SENet[11]模塊,提高了特征提取能力同時降低參數量,并且使用簡化的BiFPN取代了原始PANet[12]結構,增強了網絡提取不同尺度特征的能力,在夜間道路場景下對行人和車輛檢測的平均準確率提高了3.49 %。MFSSD[13]重新設計了特征融合網絡,加強了不同層次網絡之間的信息交互,實現了深層特征和淺層特征的有效融合?，F有的卷積神經網絡受感受野限制,無法獲取大范圍內目標和背景的差異性,后續的研究學者開始嘗試將Transformer[14]用于目標檢測,TPH-YOLOv5[15]通過探索自注意力機制使用Transformer預測頭,提升了密集場景和遮擋情況下小目標的檢測性能。Xin等人[16]使用SwinTransformer替換YOLOv5s中的部分C3模塊,在FLIR數據集上平均準確率較初始YOLOv5s提高了5.6 %。Liu等人[17]為了獲取紅外圖像中的全局依賴,提出了一種基于Transformer的紅外弱小目標檢測方法,利用Transformer的自注意力機制,在全局范圍內學習目標特征。同時為了避免目標丟失,使用了類似U-Net[18]的網絡結構來融合不同尺度的特征,在兩個公共數據集上取得了更好的檢測結果。

結合紅外圖像的特性和YOLO系列算法的優勢,本文提出了一種基于改進YOLOv5s的實時紅外小目標檢測模型,主要的創新點如下:

(1)在特征提取階段,使用SPD-Conv[19]進行下采樣,避免小目標特征丟失,同時通過串聯多個不同空洞率的空洞卷積增強多尺度特征提取能力。在主干網絡中加入了CBAM[20]空間和通道混合注意力模塊,以提高模型的表征能力,提升小目標的檢測性能。

(2)在特征融合階段,引入由深到淺的注意力模塊,將深層語義特征嵌入到淺層空間特征中,提高淺層特征的表達能力。

(3)在預測階段,裁剪網絡中針對大目標檢測的預測層及相關特征提取和特征融合層,降低了模型大小,提高了檢測實時性。

(4)最后采用Infrared-PV和SIRST[21]數據集驗證了提出算法的有效性。

2 基于改進YOLOv5s的紅外小目標檢測

2.1 YOLOv5網絡結構

根據網絡深度和寬度不同,YOLOv5模型由小到大可分為YOLOv5s、YOLOv5m、YOLOv5l和YOLOv5x。由于紅外圖像分辨率較低,樣本數量少,復雜的網絡會導致過擬合,因此本文選擇YOLOv5s作為紅外小目標檢測基準模型。YOLOv5主要分為輸入、特征提取、特征融合和預測輸出四個部分。輸入模塊使用Mosaic進行數據增強以增加樣本數量。特征提取模塊主要由CBS、C3和SPPF模塊組成,CBS采用步長為2的卷積對特征圖進行下采樣。C3模塊借鑒了CSPNet[22](Cross Stage Partial Network)的設計,將CSPBottleNeck中的卷積減少到3個,在不降低檢測精度的前提下減少了模型參數,提高了實時性。SPPF模塊在空間金字塔池化[23](Spatial Pyramid Pooling,SPP)的基礎上使用多個小尺寸池化核級聯代替SPP模塊中單個大尺寸池化核,進一步提高了檢測速度。在特征融合階段,YOLOv5s采用特征金字塔網絡[24](Feature Pyramid Network,FPN)和PANet(Path Aggregation Network)的多尺度特征融合策略,增強多尺度特征的融合能力。預測模塊主要用于檢測目標,當輸入圖像分辨率為640×640時,分別輸出20×20、40×40和80×80大小的特征圖,對應大、中、小目標檢測層。

2.2 基于改進YOLOv5s的紅外小目標檢測模型

盡管YOLOv5s性能優異,但其在紅外場景下的檢測精度仍有待提高,故本文從特征提取、特征融合、預測輸出三個方面改進YOLOv5s,提出了一個實時紅外小目標檢測模型Infrared-YOLOv5s,以提高紅外小目標檢測精度,其結構如圖1所示,圖中改進模塊用不同顏色標識。

圖1 Infrared-YOLOv5s網絡結構

2.2.1 基于SPD-Conv和IASPP的改進特征提取網絡

現有卷積神經網絡通常使用步長為2的卷積或最大池化進行下采樣,由于紅外圖像分辨率低、目標細節特征缺失,這種下采樣方式會導致細節信息丟失。因此,本文引入了SPD-Conv來替換YOLOv5s中的下采樣模塊。

SPD-Conv由空間到深度轉換層(Space-to-depth,SPD)和一個1×1卷積的CBS模塊組成。SPD層對特征圖進行下采樣時保留了通道維度中的所有信息,因此沒有信息丟失。在SPD層后添加1×1卷積降低通道數。如圖2所示,給定S×S×C1的特征圖,將其切片為四個S/2×S/2×C1的特征子圖,將這些子圖按通道拼接,得到S/2×S/2×4C1的特征圖,最后使用1×1卷積調整通道數。使用SPD-Conv進行下采樣可以最大程度保留小目標的細節特征,有利于后續的特征提取操作。

圖2 SPD-Conv結構圖

針對YOLOv5采樣過程中小目標容易丟失的問題,如圖1所示,本文在特征提取階段加入CBAM[20]注意力模塊,使網絡更專注于對小目標的檢測。在目標檢測任務中,較大的感受野可以獲得更為全局、語義層次更高的特征,但下采樣操作增大感受野的同時會帶來分辨率的降低,導致小目標丟失。為了解決這個矛盾,引入空洞卷積[25](AtrousConvolution),在減少分辨率損失的前提下,增大感受野?？斩纯臻g金字塔池化[26](Atrous SpatialPyramidPooling,ASPP)將不同感受野特征圖融合,使得像素點分類更準確。然而,隨著采樣率的增加,空洞卷積的效果會變差。為了在融合多尺度特征信息的同時獲得更大的感受野,本文重新設計了ASPP模塊,提出了改進空洞空間金字塔池化(Improved AtrousSpatialPyramidPooling,IASPP)模塊。如圖3所示,ISAPP包含三個分支,輸入特征圖經過1×1卷積得到輸出out;經過自適應全局平均池化得到輸出pool;在空洞卷積分支中,經過3×3的普通卷積得到輸出out1,然后將out1輸入采樣率為2的空洞卷積得到out2,并將其與out1拼接得到add1,輸入到采樣率為3的空洞卷積得到輸出out3,將out3與add1拼接得到add2,串聯的空洞卷積結構可以在不同采樣率的特征圖間共享特征,從而改善ASPP因采樣率變大導致效果變差的問題,增大感受野的同時又能獲取多尺度信息。ISAPP的最終輸出為Cat(pool,out,add2)。

圖3 IASPP模塊結構圖

2.2.2 基于由深到淺注意力的多尺度特征融合

淺層特征感受野小,分辨率高,包含更多細節信息,對于目標定位較為重要;深層特征可以提供更好的語義信息和對場景上下文的理解,有助于解決目標和背景干擾物之間的歧義,但隨著分辨率的降低細節信息丟失嚴重。因此實現淺層特征和深層特征的有效融合,可以提高檢測性能。如圖4(a)所示,YOLOv5通過Concat操作將淺層特征和深層特征直接按通道拼接,不能反映不同尺度特征的重要性。在多尺度特征融合階段,引入由深到淺的注意力模塊(Deep-to-Shallow Attention Module,DSAM)如圖4(b)所示,將深層語義特征嵌入到淺層空間特征,可以幫助處理歧義,提高分類準確率。

圖4 深層特征與淺層特征融合策略

對深層特征H首先通過全局平均池化(Global Average Pooling,GAP)編碼得到全局特征U,然后利用1×1卷積學習各通道之間的關系,再經過Sigmoid激活函數將輸出限制在0～1之間,得到權重w:

w=σ(SiLU(Conv(U)))

(1)

式中,σ表示Sigmoid激活函數,Conv為1×1卷積,SiLU為卷積層激活函數;將權重w與原始深層特征H相乘即可得到加權后的特征H′。最后,將加權后的深層特征H′和淺層特征L拼接,得到融合特征圖。該模塊可以將深層特征更有效地傳遞給淺層特征,提高了特征融合效果。

2.2.3 適用于紅外小目標檢測的預測層結構

本文檢測對象為紅外圖像中的弱小目標,在YOLOv5中,大目標檢測層的特征圖是對輸入圖像進行32倍下采樣得到的,當目標尺寸小于32×32像素時,會出現目標采樣不到的現象。因此,對于檢測紅外小目標,YOLOv5中的大目標檢測層屬于冗余層,會增加模型大小但對于小目標檢測沒有幫助?；谏鲜鼋Y論,如圖1所示,本文裁減了YOLOv5網絡中的大目標檢測層及其相應特征提取和特征融合層,只進行4次下采樣,僅保留8倍和16倍下采樣的特征圖進行紅外小目標檢測,改進后的網絡結構去除了大量冗余計算,在保證檢測精度的前提下,降低了模型大小,防止出現過擬合,提高了檢測實時性。

3 實驗及結果分析

3.1 紅外小目標檢測數據集

本文使用SIRST紅外數據集進行實驗,該數據集是南京航空航天大學發布的一個不同場景下的單幀紅外小目標數據集[21]。共有427張紅外圖像,包含500多個目標。圖5為SIRST數據集中的部分紅外圖像及標注信息。目標所處的環境復雜多變,且目標尺寸多樣且亮度差異較大。數據集標注信息使用SIRST數據集的分割真值圖像利用最小包圍盒算法得到,標注為VOC格式,保存為XML文件。其中訓練集256張圖片,驗證集85張圖片,測試集86張圖片。

圖5 SIRST數據集示例圖像及標注信息

3.2 訓練環境和配置

本文模型實現采用Pytorch1.7.1,實驗所用的計算機配置如下:CPU為i7-8700k,主頻為3.70 GHz,GPU為1080Ti,內存為32 G,操作系統為Windows10。實驗代碼基于YOLOv5-6.1版本改進,訓練次數(epoch)為100次,批大小為16,初始學習率為0.01,采用SGD梯度下降優化器,采用one-cycle學習率衰減,輸入的紅外圖像分辨率為640×640,其他為默認參數設置。

3.3 評價指標

為準確評估模型在紅外圖像上的檢測性能,本文采用平均精度值(mean Average precision,mAP)和F1值(F1-Score)作為評價指標。數據集中每個類別根據準確率(Precision,P)和召回率(Recall,R)可繪制一條PR曲線,曲線與坐標軸圍成的面積即為AP值。其中準確率和召回率計算如式(2),其TP為真正例,FP為假正例,FN為假反例:

(2)

當檢測框與真值框的交并比(Intersection over Union,IoU)大于0.5時認為目標被準確預測,在IoU取0.5時計算每個類別的平均精度和總平均精度,記為mAP0.5。

F1值是分類問題的一個評價指標,同時兼顧了分類模型的精確率和召回率,可以看作是模型精確率和召回率的一種調和平均值,計算方法如式(3):

(3)

3.4 消融實驗

為驗證各模塊的有效性,本文以YOLOv5s為基準,設計了如下消融實驗:A為采用2.2.1節的基于SPD-Conv和IASPP的改進特征提取網絡,B為采用2.2.2節的基于由深到淺注意力的多尺度特征融合,C為采用2.2.3節的適用于紅外小目標檢測的預測層結構。實驗結果如表1所示,其中實驗1為YOLOv5s基準模型實驗結果。

表1 不同模塊消融實驗結果

(1)實驗2和6證明,在不同的預測層結構下,采用基于SPD-Conv和IASPP的改進特征提取網絡,mAP0.5分別提高了1.4 %和1.7 %。SPD-Conv下采樣模塊最大程度地保留了小目標特征,改進空洞空間金字塔池化模塊可以融合多尺度特征,增大了感受野,以適應不同目標尺寸大小的變化,CBAM注意力模塊同時利用了空間注意力和通道注意力,使網絡更關注目標,進一步提高了檢測精度。

(2)實驗3和7證明,在特征融合階段引入由深到淺的注意力模塊,在不增加額外的內存占用的情況下,mAP0.5分別提高了0.4 %和0.6 %。DSAM注意力模塊將深層特征豐富的語義信息嵌入到淺層特征,提高了淺層特征的表達能力,相比原始YOLOv5直接通道相加的方法,可以獲得更豐富的語義和空間特征,因此可以提高檢測精度。

(3)實驗4證明,與基準YOLOv5s比較,使用適用于紅外小目標檢測的預測層結構,mAP0.5僅降低0.4 %,模型大小僅為原始的四分之一左右,提升了檢測實時性。

綜上所述,使用所有改進策略的實驗8,提出的Infrared-YOLOv5s較基準YOLOv5s,mAP0.5提高了2.3 %,且模型大小僅為原始的27.1 %,驗證了改進算法在紅外小目標圖像數據集上的有效性。

3.5 SIRST數據集算法驗證

為驗證本文算法的有效性,本文以YOLOv5s模型為基準,并與文獻[17]和[21]提出的算法進行了對比。實驗結果如表2所示,在SIRST數據集上,較基準模型YOLOv5s,改進模型mAP0.5提高了2.3 %,F1分數提高了3.18,驗證了改進算法在紅外數據集上的有效性。雖然F1分數比文獻[17]提出的模型低,但由于文獻[17]使用了Transformer結構,使模型參數增加,訓練和檢測速度較慢,本文算法檢測實時性更好,檢測時間僅為文獻[17]的十分之一,實現了檢測性能和檢測速度的平衡。由圖6的檢測結果圖像可知,Infrared-YOLOv5s模型在低對比度和復雜多目標場景下的紅外小目標檢出率優于YOLOv5s,虛警率更低。

表2 不同算法在SIRST數據集上的檢測性能

圖6 不同算法在SIRST數據集的檢測結果圖像

3.6 基于Infrared-PV數據集的遷移實驗

為充分驗證本文算法的有效性和魯棒性,在自建的Infrared-PV數據集上進行了遷移實驗。Infrared-PV數據集包括行人(Person)和車輛(Vehicle)兩類目標,總計包2138張圖片,其中白熱圖1000張,黑熱圖838張,熱力圖300張,采用VOC格式進行標注,保存為XML文件。平均單張圖片包含20個左右目標,單張圖片最多目標數超過100個,目標比較密集,且目標占整幅圖像的10 %以下,以中小目標為主,適合遷移驗證本文算法的有效性。圖7為Infrared-PV數據集的示例圖像。

圖7 Infrared-PV數據集示例圖像

圖8 不同算法在Infrared-PV數據集的檢測結果圖像

實驗結果如表3所示,YOLOv7采用高效的ELAN主干網絡并結合多種訓練優化策略,檢測精度比YOLOv5s提高了2.3 %。本文提出的模型針對紅外小目標檢測任務改進特征提取網絡并結合基于注意力的特征融合,較YOLOv5s基準模型,檢測精度提高了2.8 %,達到84.5 %,優于YOLOv7算法和兩階段的CascadeRCNN算法。由于采用了適用于紅外小目標檢測的預測層結構,在PC機上推理速度可達172.5 f/s,實時性更好。由圖8中檢測結果圖像可知,改進模型在密集和遮擋場景下表現優于YOLOv5s模型。實驗表明,本文算法對于尺度差異較大、重疊目標和密集目標實現了較好的魯棒性。

表3 不同算法在Infrared-PV數據集上的檢測性能

3.7 Nvidia Xavier設備部署實驗

為驗證本文算法在移動設備上的性能,在Nvidia Xavier設備上進行了部署實驗。Xavier是一款高性能AI邊緣設備,擁有一顆8核心ARM架構CPU,16GB、256位LPDDR4x內存,其GPU含有8個流式多核處理器,擁有512個CUDA核、64個張量核(Tensor-Core)、兩個深度學習加速器(Deep Learning Accelerator,DLA)和其他硬件資源,最高算力可達32萬億次每秒(TeraOperationsPerSecond,TOPS),功耗在10W到30W之間,擁有強大的計算能力且功耗較低,其軟硬件資源如表4所示。

表4 NvidiaXavier軟硬件資源

圖9為NvidiaXavier設備部署實物圖,界面使用QT搭建,集成了模型訓練、圖像及視頻檢測、性能測試等功能。使用PC機訓練得到的權重文件在Xavier設備上進行測試,推理速度可達28 f/s,達到邊緣設備部署的實時性要求。

圖9 NvidiaXavier設備部署實物圖

4 結 論

本文研究了YOLOv5網絡結構及其各個模塊的作用,通過分析紅外小目標圖像的特性,提出了一種基于YOLOv5s的改進實時紅外小目標檢測模型Infrared-YOLOv5s。首先在特征提取階段采用SPD-Conv下采樣避免了紅外小目標特征丟失,設計了改進空洞空間金字塔池化模塊,增強多尺度特征提取能力,以適應目標尺寸變化;其次在特征融合階段引入由深到淺的注意力模塊,將深層語義特征嵌入到淺層空間特征中,提高淺層特征的表達能力;預測階段裁剪了針對大目標檢測的特征提取、融合以及預測層,降低了模型大小,提升了檢測實時性。最后基于SIRST數據集對各個改進模塊設計了消融實驗和不同算法之間的對比實驗。實驗結果表明,改進后的算法在SIRST數據集上mAP0.5提高了2.3 %,保證檢測精度的同時,在NvidiaXavier設備上推理速度達到28 f/s,能夠滿足實際部署需求。在Infrared-PV數據集上的遷移實驗表明,改進模型較YOLOv5s基準模型,mAP0.5提高2.8 %,進一步驗證了改進算法的有效性和魯棒性。