基于改進YOLOv5算法的地鐵車輛轉向架螺栓缺失檢測

鄒一鳴，李鵬，林群煦，陳金源，鄭衡，李虎，劉凱，鐘程

（五邑大學 軌道交通學院，廣東 江門 529020）

0 引言

目前，人工巡檢依舊為列車螺栓故障排查的主要方式。由于工人在巡檢作業過程中易產生視覺疲勞，從而導致檢測的效率降低，嚴重時還會發生漏檢、誤檢等情況。若不及時發現與處理，會對列車的行車安全帶來嚴重隱患。隨著機器視覺技術的不斷發展，檢測速度快、準確度高、效率好的機器視覺檢測技術紛紛被應用到軌道車輛檢測中來，逐漸取代了人工巡檢方式。

傳統的圖像處理技術識別精度較低，且容易受到環境干擾，對于背景比較復雜的圖片特征提取較為困難。而卷積神經網絡可以獨立地學習相關性，對于不同形式出現的目標、特征具有較強的魯棒性和適應性，并且隨著硬件設備的快速迭代與更新，基于深度學習的卷積神經網絡（Convolution Neural Network，CNN）[1-3]大放異彩，越來越多的研究人員和工程師將卷積神經網絡提取的圖像特征用于目標檢測領域。目前，目標檢測方法主要分為兩大類：第一類為二階段方法，該方法是將目標的定位與目標的分類分為兩步進行，其中代表算法有R-CNN、Fast R-CNN、Faster R-CNN等[4]；第二類為一階段方法，該方法的目標定位與分類是一步進行的，代表的算法有YOLO[5]系列、SSD[6]等。二階段方法（如R-CNN）會對每個候選區域的卷積神經網絡特征進行大量計算，其識別準確度較高，但大量的計算會降低網絡的檢測速度。而一階段網絡犧牲了部分檢測精度，但提升了檢測速度。

由于螺栓多為小目標物體，受限于其尺寸，特征表達能力較弱；螺栓圖片的背景復雜，檢測過程中易受到外界干擾；螺栓缺失的樣本數據又十分稀少，導致直接檢測缺失螺栓的效果不理想。為解決上述問題，本文提出了基于深度學習的YOLOv5網絡模型檢測算法，通過引入Ghost模塊以替換傳統卷積層，并加入CBAM模塊，使檢測算法具有更高的檢測精度、更快的檢測速度，提高了網絡的適應性和魯棒性。最后通過模板匹配方法，以缺失螺栓圖像作為模板匹配異常的數據集樣本，找到螺栓缺失位置，完成缺失檢測。

1 YOLOv5 算法原理

YOLOv5是一種端到端的目標檢測算法，其骨干網絡與YOLOv4算法基本保持一致。在Neck部分，YOLOv5引入了C3結構，它是由前后各一個CBL結構、數個殘差結構（N*BottleNeck），以及一個殘差邊兩個分支堆疊（Concat）組成[7]。包含了BottleNeck1的C3結構應用于骨干網絡（Backbone）部分，增加殘差結構可以增加層與層之間的反向傳播梯度值，避免了因為網絡層數過多而帶來的梯度消失問題。包含了BottleNeck2的C3結構應用與頸部（Neck）部分[8]。在保證計算結果一致的同時，提升了檢測速度，讓模型學到更多的特征。

YOLOv5算法共有4個版本，其中包括YOLOv5s、YOLOv5m、YOLOv5l、YOLOv5x共4種。這4種模型的內容基本一致[9]，網絡的深度與寬度逐漸增加。其中YOLOv5s模型較小，對檢測硬件的要求較低，同時檢測精度和速度都比較快，因此，本文選取YOLOv5s作為本文目標檢測算法。

YOLOv5網絡模型結構如圖1所示，主要由4部分組成：輸入端（Input）、骨干網絡（Backbone）、頸部（Neck）和輸出端（Head）[10]。圖像樣本（尺寸為640×640×3）從輸入端送入網絡，并采用Mosaic數據增強、自適應圖片縮放對樣本數據進行豐富，進行自適應錨框計算等任務；骨干網絡通過對Bottleneck與CSP結構的有機結合可以加強CNN網絡的學習能力，減少性能消耗；頸部可以很好地融合不同尺寸大小的特征來適應不同大小的對象，避免小物體的特征被忽視；輸出端的作用則是對圖像特征進行預測、分類，并生成邊界框。

圖1 YOLOv5網絡結構示意圖

同時，基于原YOLOv4算法的SPP結構，YOLOv5算法中該部分換成了SPPF結構，將輸入特征串行通過CBL層和3個5×5大小的MaxPool層，并將各層結構在空間維度上進行堆疊（Concat），使主干特征的接收范圍增大，降低了網絡的過擬合風險，也提升了網絡的運行速度。

2 改進的YOLOv5 算法模型

2.1 Ghost模塊

不可否認的是，卷積神經網絡具有十分優秀的提取特征的能力，但是特征圖之間存在著大量冗余的信息，在一定程度上增加了網絡計算參數，降低了網絡計算速度。為了減少模型網絡參數，輕量化網絡模型，同時提升網絡的檢測速度和檢測精度，本文采用Ghost模塊替換了原來的卷積層。其次，使用Ghost BottleNeck模塊替換骨干網絡與頸部中的BottleNeck模塊。充分利用冗余的特征信息，通過減少計算量以生成大量特征圖，從而提高網絡特征提取能力。

Ghost模塊[11]（Ghost Module）（如圖2）在保證檢測精度的同時，有著非常優秀的壓縮效果，且可以替換現有卷積網絡中的每一個卷積層。相比于普通的卷積，Ghost模塊分兩步走。首先，Ghost模塊采用正常的卷積計算，但是在這一步驟中，使用了更少的卷積核從而得到較少通道的特征圖。第二步，對特征圖進行簡單的線性變化操作，從而進一步生成Ghost 特征圖，并將兩步結果進行堆疊（Concat）。具體過程如下：

圖2 普通卷積模塊與Ghost模塊卷積過程示意圖

式中：X為輸入特征圖；*為卷積操作；f′為當前層的卷積核；Y′為m個通道的本征特征圖。

對Y′特征圖的每一個通道進行線性運算輸出通道數為s的特征圖，有：

式中：yi為Y′的第i個特征通道；Φi,j為yi′生成的第j個特征圖的線性運算。每個yi′生成s個Ghost特征圖，因此最后得到Ghost模塊的輸出特征圖為Y=[Y11,Y12,Y13,…,Yms]。

Ghost BottleNeck（如圖3）由兩個Ghost模塊堆疊而成，類似于ResNet中的基本殘差塊。其中，第一個Ghost模塊用于擴張特征的維度，增大其比例。第二個Ghost模塊則用于減少特征維度。Ghost BottleNeck 包含stride=1 和stride=2 版本，對于stride=2，兩個Ghost 模塊之間通過一個stride=2的深度卷積進行連接。

圖3 Ghost BottleNeck

2.2 卷積注意力機制模塊

注意力在人類感知中起到重要作用。人類視覺系統的一個重要特征并不是一次處理視野中的整個場景。相反，人類的視覺結構利用一系列的部分瞥見并選擇性地關注顯著的部分，以便更好地捕捉信息[12]。卷積注意力機制模塊（Convolutional Block Attention Module，CBAM）[13]便是模仿了人類視覺系統的這一特性，其本質是一種通道和空間混合型的注意力機制，具體結構如圖4所示。

圖4 卷積注意力機制模塊(CBAM)結構示意圖

CBAM模塊集成了通道和空間兩個注意力映射過程并將其串聯組合而成，這能夠保留更多有用的特征信息[14]。對輸入的特征圖（Input Feature），該模塊會沿著兩個獨立的維度（通道注意力模塊和空間注意力模塊）依次推斷注意力圖，隨后應用相應的注意力權重對注意力圖與特征圖相乘以達到自適應修飾的效果。

如圖5所示，通道注意力模塊首先通過使用平均池化和最大池化操作來聚合特征圖的空間信息，隨后將兩種池化后的特征送入多層感知機（MutiLayer Perception，MLP）網絡計算，將通過多層感知機計算后的平均池化特征和最大池化特征相加并經過sigmoid函數激活，獲得通道注意力圖。

圖5 通道注意力模塊（Channel Attention Module）結構示意圖

如圖6所示，空間注意力模塊使用了平均池化和最大池化操作聚合特征圖的通道信息，隨后通過標準的卷積層將平均池化層和最大池化層的兩個映射疊加（concat）生成空間注意力映射，得到了空間注意力圖。

圖6 空間注意力模塊（Spatial Attention Module）結構示意圖

結合CBAM 模塊即插即用的優勢，本文選擇在SPPF模塊前插入CBAM模塊，用以提高檢測精度，減少網絡漏檢率。

3 螺栓缺失檢測

3.1 螺栓缺失粗檢測

在對列車轉向架關鍵區域螺栓進行準確定位的基礎上，判斷每一個螺栓是否缺失，是列車轉向架螺栓缺失檢測的重要步驟。

通過分析列車轉向架螺栓的分布規律，可獲得正常狀態下列車轉向架的螺栓數量及分布。若將待檢圖片送入網絡獲得的螺栓檢測結果數量與實際的螺栓數量不一致，則該轉向架螺栓可能存在螺栓缺失，也有可能存在漏檢的情況。圖7所示為螺栓缺失的圖片樣本。

圖7 螺栓缺失樣本

從圖7可知，以轉向架軸箱位置為列，該部件位置上螺栓正確數量為6個。在檢測該區域時，計算正常螺栓檢測框個數，若存在螺栓漏檢和螺栓缺失的檢測結果，則該結果與列車轉向架螺栓的正常數量不符。因此，將該結果組成缺失樣本數據集，完成粗檢測。

3.2 基于模板匹配的螺栓缺失精檢測

模板匹配是指通過模板圖像與測試圖像之間的比較，找到測試圖像上與模板圖像有最高相似度的部分，并將其選為匹配區域[15]。

基于上述螺栓缺失粗檢測，得到螺栓缺失樣本數據集。由于已知轉向架螺栓的分布與數量，提取缺失螺栓圖片作為矩形模板。為了兼顧檢測精度與檢測速度，本文使用相關性匹配方法，具體過程如下：

式中：R（x,y）為相似度函數，其中（x,y）為所選特定圖像塊區域的位置；T（x′，y′）為所選模板T的點坐標；I（x+x′，y+y′）為圖像I中點的坐標。

本文將R（x,y）閾值設定為0.7，若相似度不小于閾值則表示該螺栓缺失，若相似度小于閾值則表示該位置螺栓漏檢，具體流程如圖8所示。

圖8 螺栓缺失檢測流程圖

4 實驗結果

4.1 實驗環境與模型訓練

本文實驗的數據采集自東莞地鐵2號線現場，共有2000張圖片。由于螺栓影像特征不明顯，背景復雜，且采集時受到光照、天氣環境影響，通過對采集的數據進行翻轉、平移和角度旋轉等方法以增加數據集至3000張圖片，其中隨機抽取2250張圖片作為訓練集，750張圖片作為測試集。

實驗采用了以下的硬件，其中顯卡使用的是NVIDIA GTX1080ti（11 GB顯存），處理器為2.2 GHz E5-25650 v4，并且搭配了128 GB內存和2 TB硬盤；操作系統使用了Ubuntu 18.04版本；深度學習環境配置為：深度學習框架平臺為Pytorch 1.9.0版本，Python語言環境為3.8版本，CUDA 11.0版本，CUDNN 7.6版本。

實驗選取了SGD（stochastic gradient descent）優化器，初始學習率（Learning Rate）設置為0.01，權重衰減（weight decay）設置為0.000 5，總迭代次數設置為200個epochs，batch-size為32。

4.2 評估指標與實驗結果分析

本文將精確度（Precision）、召回率（Recall）及檢測速度（FPS）這3個常見的目標檢測評價指標作為本次實驗模型的評價指標。其中，精確率能夠反映網絡模型對目標的誤檢情況，召回率能夠反映網絡模型對目標的漏檢情況，檢測速度能夠反映網絡模型能否達到實時檢測的使用要求，FPS值越高則檢測速度越快。具體表達如下：

式中：TP、FP和FN分別表示正確檢測框、誤檢框和漏檢框的數量；FigureNumber為檢測圖片總數；TotalTime為檢測總時間。

本文改進的YOLOv5算法比原YOLOv5算法的檢測效果在檢測精度、檢測速度、召回率等方面均有所提升，如表1所示。其中，改進后的YOLOv5算法提高了3.9%的精確度，召回率提升6.1%，模型參數量也減少了3.55×106個，檢測速度從原來的52 fps提升至61 fps，增加了9 fps。該實驗結果證明了本文的改進的YOLOv5算法是有效的。

表1 YOLOv5算法改進前后性能對比

圖9（a）、圖9（c）所示為原YOLOv5網絡模型的檢測效果圖，圖9（b）、圖9（d）為改進后的YOLOv5網絡模型的檢測效果圖。對比圖9（a）與圖9（b）、圖9（c）與圖9（d）可以直觀地看出，改進后的YOLOv5網絡模型對于螺栓的定位有更高的置信度，預測框對于目標的定位也更加精準。不僅如此，從圖9（c）與圖9（d）的對比結果可以看出，改進后的YOLOv5網絡模型對于小目標的漏檢和誤檢情況有了一定程度的改善。

圖9 改進前后的YOLOv5檢測效果對比圖

如圖10（a）所示，正常的螺栓均被精準地檢測，由于該部件上螺栓的分布和數量固定，使用模板匹配方法對該圖片進行匹配，得到缺失螺栓的坐標信息，并進行標記（如圖10（b）的白色框內），缺失檢測精度達到86.2%。

圖10 螺栓缺失檢測效果圖

本文選用Faster-RCNN算法和SSD算法對同樣的數據集進行同步訓練，并將三者結果對比以進一步驗證改進后的YOLOv5網絡模型的檢測性能，訓練的超參數也設置為相同，結果對比如表2所示。

表2 不同檢測算法的檢測能力對比結果

由表2可知，改進的YOLOv5算法的檢測精度比SSD算法高出6.7%，但略低于Faster-RCNN算法；在檢測速度上，改進后的YOLOv5算法比SSD算法和Faster-RCNN算法都要快，尤其是與Faster-RCNN相比，在檢測速度方面改進后的YOLOv5有較大優勢。因此，綜合考慮本文改進的算法還是有一定效果和優勢的。

5 結語

本文在原YOLOv5網絡模型的基礎上，通過替換普通卷積為Ghost卷積、引入CBAM模塊對網絡進行改進。改進后的網絡模型克服了螺栓作為小目標物體在目標檢測過程中特征表達能力弱、目標漏檢等不足。本文以此為基礎，通過模板匹配方法對缺失螺栓進行定位，從而完成螺栓缺失的識別任務。實驗結果顯示，改進后的YOLOv5在減小模型參數的同時有更高的精度、召回率和檢測速度。不足之處在于該算法暫時無法識別螺栓松動等其他異常情況，后續會針對這些不足繼續完善算法模型。