基于改進(jìn)YOLOv5算法的地鐵車輛轉(zhuǎn)向架螺栓缺失檢測

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

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

0 引言

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

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

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

1 YOLOv5 算法原理

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

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

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

圖1 YOLOv5網(wǎng)絡(luò)結(jié)構(gòu)示意圖

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

2 改進(jìn)的YOLOv5 算法模型

2.1 Ghost模塊

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

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

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

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

對Y′特征圖的每一個通道進(jìn)行線性運(yùn)算輸出通道數(shù)為s的特征圖，有：

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

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

圖3 Ghost BottleNeck

2.2 卷積注意力機(jī)制模塊

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

圖4 卷積注意力機(jī)制模塊(CBAM)結(jié)構(gòu)示意圖

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

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

圖5 通道注意力模塊（Channel Attention Module）結(jié)構(gòu)示意圖

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

圖6 空間注意力模塊（Spatial Attention Module）結(jié)構(gòu)示意圖

結(jié)合CBAM 模塊即插即用的優(yōu)勢，本文選擇在SPPF模塊前插入CBAM模塊，用以提高檢測精度，減少網(wǎng)絡(luò)漏檢率。

3 螺栓缺失檢測

3.1 螺栓缺失粗檢測

在對列車轉(zhuǎn)向架關(guān)鍵區(qū)域螺栓進(jìn)行準(zhǔn)確定位的基礎(chǔ)上，判斷每一個螺栓是否缺失，是列車轉(zhuǎn)向架螺栓缺失檢測的重要步驟。

通過分析列車轉(zhuǎn)向架螺栓的分布規(guī)律，可獲得正常狀態(tài)下列車轉(zhuǎn)向架的螺栓數(shù)量及分布。若將待檢圖片送入網(wǎng)絡(luò)獲得的螺栓檢測結(jié)果數(shù)量與實(shí)際的螺栓數(shù)量不一致，則該轉(zhuǎn)向架螺栓可能存在螺栓缺失，也有可能存在漏檢的情況。圖7所示為螺栓缺失的圖片樣本。

圖7 螺栓缺失樣本

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

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

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

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

式中：R（x,y）為相似度函數(shù)，其中（x,y）為所選特定圖像塊區(qū)域的位置；T（x′，y′）為所選模板T的點(diǎn)坐標(biāo)；I（x+x′，y+y′）為圖像I中點(diǎn)的坐標(biāo)。

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

圖8 螺栓缺失檢測流程圖

4 實(shí)驗(yàn)結(jié)果

4.1 實(shí)驗(yàn)環(huán)境與模型訓(xùn)練

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

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

實(shí)驗(yàn)選取了SGD（stochastic gradient descent）優(yōu)化器，初始學(xué)習(xí)率（Learning Rate）設(shè)置為0.01，權(quán)重衰減（weight decay）設(shè)置為0.000 5，總迭代次數(shù)設(shè)置為200個epochs，batch-size為32。

4.2 評估指標(biāo)與實(shí)驗(yàn)結(jié)果分析

本文將精確度（Precision）、召回率（Recall）及檢測速度（FPS）這3個常見的目標(biāo)檢測評價指標(biāo)作為本次實(shí)驗(yàn)?zāi)Ｐ偷脑u價指標(biāo)。其中，精確率能夠反映網(wǎng)絡(luò)模型對目標(biāo)的誤檢情況，召回率能夠反映網(wǎng)絡(luò)模型對目標(biāo)的漏檢情況，檢測速度能夠反映網(wǎng)絡(luò)模型能否達(dá)到實(shí)時檢測的使用要求，F(xiàn)PS值越高則檢測速度越快。具體表達(dá)如下：

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

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

表1 YOLOv5算法改進(jìn)前后性能對比

圖9（a）、圖9（c）所示為原YOLOv5網(wǎng)絡(luò)模型的檢測效果圖，圖9（b）、圖9（d）為改進(jìn)后的YOLOv5網(wǎng)絡(luò)模型的檢測效果圖。對比圖9（a）與圖9（b）、圖9（c）與圖9（d）可以直觀地看出，改進(jìn)后的YOLOv5網(wǎng)絡(luò)模型對于螺栓的定位有更高的置信度，預(yù)測框?qū)τ谀繕?biāo)的定位也更加精準(zhǔn)。不僅如此，從圖9（c）與圖9（d）的對比結(jié)果可以看出，改進(jìn)后的YOLOv5網(wǎng)絡(luò)模型對于小目標(biāo)的漏檢和誤檢情況有了一定程度的改善。

圖9 改進(jìn)前后的YOLOv5檢測效果對比圖

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

圖10 螺栓缺失檢測效果圖

本文選用Faster-RCNN算法和SSD算法對同樣的數(shù)據(jù)集進(jìn)行同步訓(xùn)練，并將三者結(jié)果對比以進(jìn)一步驗(yàn)證改進(jìn)后的YOLOv5網(wǎng)絡(luò)模型的檢測性能，訓(xùn)練的超參數(shù)也設(shè)置為相同，結(jié)果對比如表2所示。

表2 不同檢測算法的檢測能力對比結(jié)果

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

5 結(jié)語

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