基于YOLOv5 算法的鋼軌扣件分類檢測研究

陳修忻，楊光

（上海地鐵維護(hù)保障有限公司工務(wù)分公司，上海，200233）

0 引言

鋼軌扣件是連接鋼軌和軌枕使之形成軌排的部件，在保證軌道穩(wěn)定性、可靠性方面起著重要作用[1]。鋼軌扣件在軌道線路存在的狀態(tài)可分為正常、斷裂、缺失、移位四大類，我國長期以來對鋼軌扣件狀態(tài)檢查主要還是以人工和靜態(tài)檢測為主。近年來，隨著我國鐵路與城市軌道交通的迅猛發(fā)展，信息化、智能化處置軌道線路故障是當(dāng)前的一大趨勢。

在鋼軌扣件檢測的研究領(lǐng)域，基于圖像處理與機(jī)器視覺技術(shù)的檢測方法不斷提出，白堂博等[2]提出一種基于改進(jìn)Faster R-CNN 的深度學(xué)習(xí)理論方法進(jìn)行扣件定位；齊勝等[3]提出基于BEMD-IPSO-SVM 的檢測算法對扣件完損狀態(tài)進(jìn)行檢測；許貴陽等[4]提出基于K 均值聚類算法改進(jìn)掩膜區(qū)域卷積神經(jīng)網(wǎng)絡(luò)(Mask R-CNN)實例分割算法中的區(qū)域建議網(wǎng)絡(luò)方法對普速鐵路有砟軌道2 個扣件數(shù)據(jù)集和高速鐵路無砟軌道1 個扣件數(shù)據(jù)集上進(jìn)行軌道扣件狀態(tài)檢測；Giben X 等[5]提出基于深度學(xué)習(xí)卷積神經(jīng)網(wǎng)絡(luò)的語義分割對扣件進(jìn)行檢測；以上方法對軌道扣件檢測均能獲得較高精度，但存在內(nèi)存消耗大、檢測速度慢的缺點(diǎn)，這難以實現(xiàn)快速、高效、高精度 的在線檢測目標(biāo)。針對這種情況，采用基于YOLOv5 算法的鋼軌扣件檢測方法，分析不同模型在不同參數(shù)下的檢測效果，最終得到高效率、高精度的鋼軌扣件檢測算法模型。

1 YOLOv5 算法

■1.1 YOLOv5 算法模型

在YOLO 算法提出之前，目標(biāo)檢測主流算法框架為R-CNN 系列的two-stage 算法[6]。隨著近幾年目標(biāo)檢測算法的快速發(fā)展，以YOLO 系列為代表的one-stage 算法被提出，其核心要義是將目標(biāo)檢測定義為一個線性回歸問題[7]。YOLO 算法是將單個卷積神經(jīng)網(wǎng)絡(luò)應(yīng)用于整個圖像，將圖像分割成若干個網(wǎng)格并預(yù)測每個網(wǎng)格的類概率和邊界框，以此實現(xiàn)目標(biāo)檢測速度快、精度高的目的。

YOLOv5 算法是在YOLOv4 的基礎(chǔ)上改進(jìn)而來，其最主要的區(qū)別在于對anchors 的處理機(jī)制，YOLOv5 作為YOLO 系列最新成員，其強(qiáng)大的目標(biāo)檢測算法以及輕量級的模型能夠訓(xùn)練出快速而準(zhǔn)確的權(quán)重文件。YOLOv5 模型沿用YOLOv3、YOLOv4 的整體布局，主要包含圖像輸入、Backbone、Neck、Head 四個部分[8]，其網(wǎng)絡(luò)結(jié)構(gòu)如圖1所示。

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

■1.2 圖像輸入

YOLOv5 采用Mosaic 數(shù)據(jù)增強(qiáng)操作提升模型的網(wǎng)絡(luò)精度與訓(xùn)練速度，同時提出一種自適應(yīng)錨框計算與自適應(yīng)圖片縮放方法。

（1）Mosaic 數(shù)據(jù)增強(qiáng)

該操作是通過利用四張隨機(jī)圖像，對其進(jìn)行隨機(jī)縮放、裁剪、排布后進(jìn)行拼接且每張圖像都有自己對應(yīng)的框，最后形成一張新的圖像且該張圖像又具有其對應(yīng)的框，接著將該張新圖像傳入神經(jīng)網(wǎng)絡(luò)中學(xué)習(xí)，該方法極大豐富了被檢物體的背景[9]。

（2）自適應(yīng)錨框計算

針對不同數(shù)據(jù)集，均需設(shè)定特定長寬錨點(diǎn)框。在網(wǎng)絡(luò)訓(xùn)練階段，模型在初始錨點(diǎn)的基礎(chǔ)上輸出對應(yīng)預(yù)測框，隨后計算其與GT 框的差距并執(zhí)行反向更新操作，進(jìn)而更新整個網(wǎng)絡(luò)參數(shù)，YOLOv5 算法可根據(jù)數(shù)據(jù)集的名稱自適應(yīng)計算出最佳錨框點(diǎn)[10]。

（3）自適應(yīng)圖片縮放

在傳統(tǒng)目標(biāo)檢測算法中，不同圖像其長寬不盡相同，因此，在檢測前需將圖像統(tǒng)一縮放到某一標(biāo)準(zhǔn)尺寸，這種方式可能造成大量的信息冗余進(jìn)而影響整個算法檢測速度。YOLOv5 算法提出一種能夠自適應(yīng)添加最少黑邊到縮放圖像的方法，這樣可以極大提升算法檢測速度。

■1.3 Backbone 網(wǎng)絡(luò)

（1）Focus 模塊

在YOLOv5 網(wǎng)絡(luò)中，第一層即為Focus 模塊，其主要目的是對圖像進(jìn)行切片操作，具體操作是在一張圖像中每隔一個像素提取一個數(shù)值，最后一張圖像形成互補(bǔ)的四張圖像。相較于YOLOv3網(wǎng)絡(luò)，F(xiàn)ocus可以起到減少層數(shù)、參數(shù)量、計算量、cuda 內(nèi)存占用的作用，同時在mAP 影響很小的情況下，提升推理速度和梯度反傳速度。

（2）SPPF 結(jié)構(gòu)

SPPF 結(jié)構(gòu)是輸入多個串行的5×5 大小的MaxPool 層，其結(jié)構(gòu)圖如圖2 所示。

圖2 SPPF 結(jié)構(gòu)圖

■1.4 Neck 網(wǎng)絡(luò)

Neck 網(wǎng)絡(luò)在YOLOv5 中主要用于生成特征金字塔，其目的是增強(qiáng)模型對不同縮放尺度對象的檢測，以便能很好地實現(xiàn)識別不同大小、不同尺度的同一類物體。在YOLOv5 中，Neck 網(wǎng)絡(luò)采用了FPN+PAN 結(jié)構(gòu)。其中，F(xiàn)PN 通過融合高低層特征可很好的提升小目標(biāo)檢測效果，PAN 則在FPN 的基礎(chǔ)上增加了一個自底向上方向的增強(qiáng)，進(jìn)而提升大目標(biāo)物體的檢測效果。

■1.5 評價指標(biāo)

目標(biāo)檢測首先考慮的就是目標(biāo)識別效果的判別，本文采用測試精度P、召回率R、mAP@.5 和mAP@.5:.95作為評價鋼軌扣件檢測模型的性能。其中，P、R 分別指的是一個檢測模型預(yù)測的正樣本的確是正樣本的概率與實際正樣本的個數(shù)；AP 為P 對R的積分，而mAP@.5 是將IoU 閾值（IoU 表示檢測框與真實標(biāo)簽的交集和并集的比值）設(shè)為0.5 時，計算每一類的所有圖片的AP，然后所有類別求平均，mAP@.5:.95是不同IoU 閾值（從0.5 到0.95，步長0.05）上的平均mAP，mAP 主要用來評價檢測模型的整體性能。

2 試驗驗證

■2.1 圖像采集

本文鋼軌扣件圖像是由線陣相機(jī)現(xiàn)場采集而來，通過對某些扣件拆檢以及安裝部分?jǐn)嗔芽奂确绞綄崿F(xiàn)制作彈條缺失、彈條斷裂、彈條移位和螺栓缺失四種狀態(tài)扣件樣本圖片共計706 張（圖片中彈條斷裂120 處、彈條缺失234 處、彈條移位419 處、螺栓缺失79 處），將其按8：2 分為測試集圖像和訓(xùn)練集圖像。采集圖像（部分截圖）如圖3 所示。

圖3 現(xiàn)場采集扣件圖

■2.2 圖像訓(xùn)練及結(jié)果輸出

將選好的706 張圖像數(shù)據(jù)集中每個鋼軌故障扣件進(jìn)行分類標(biāo)注最終生成一一對應(yīng)的*.txt 的labels 文件，隨后將其輸入到 YOLOv5 網(wǎng)絡(luò)中，經(jīng)過算法的訓(xùn)練得到相應(yīng)訓(xùn)練模型并最終將訓(xùn)練完成的模型輸入到算法網(wǎng)絡(luò)中對未知扣件進(jìn)行預(yù)測分析，由此實現(xiàn)對故障扣件的檢測，方法框架如圖4 所示。

圖4 方法框架

本文主要分析了YOLOv5s、YOLOv5m 兩種模型下不同訓(xùn)練次數(shù)（epochs 取50、100 次做對比）對各類故障扣件的識別效果，具體情況如表1 所示。表1 顯示：YOLOv5s模型下訓(xùn)練得到的鋼軌故障扣件檢測模型大小是YOLOv5m模型下訓(xùn)練得到的模型大小的三分之一左右；YOLOv5s、YOLOv5s 模型下訓(xùn)練鋼軌故障扣件時，epochs 取100 時訓(xùn)練出的檢測模型要優(yōu)于epochs 取50 時的檢測模型；由于螺栓缺失類別數(shù)量最少，在相同模型、相同參數(shù)下其訓(xùn)練出檢測模型質(zhì)量最差；YOLOv5m 模型下訓(xùn)練鋼軌故障扣件時，epochs 取100 時各類別鋼軌故障扣件識別精度P、召回率R 和mAP@.5 均超過到0.9 且所有類別故障扣件mAP@.5：.95 也達(dá)到0.822。

表1 不同參數(shù)下各類扣件狀態(tài)識別效果

圖5 所示為epochs 取100、YOLOv5m 模型下訓(xùn)練得到的鋼軌故障扣件檢測模型推理結(jié)果，圖5(a)顯示的是鋼軌彈條缺失QS 與鋼軌彈條移位YW，圖5(b)顯示的是鋼軌彈條斷裂DL，圖5(c)顯示的是鋼軌螺栓缺失LSQS，圖中顯示的數(shù)值為在該檢測模型下預(yù)測對應(yīng)鋼軌故障扣件類型的置信度大小，從數(shù)值可以看出，該模型對各類狀態(tài)鋼軌扣件具有很高的識別精度。

圖5 預(yù)測結(jié)果

最后，將識別到的鋼軌故障扣件截圖分類保存至對應(yīng)類別文件夾，結(jié)果如圖6 所示。圖6(a)為彈條斷裂截圖，圖6(b)為彈條缺失截圖，圖6(c)為彈條移位截圖，圖6(d)為螺栓缺失截圖。

圖6 故障扣件截圖

3 結(jié)束語

本文采用YOLOv5 算法，分析了YOLOv5s、YOLOv5m兩種模型對不同狀態(tài)鋼軌扣件樣本訓(xùn)練50、100 次后的識別效果，結(jié)果顯示：針對輕量級鋼軌扣件圖像集分類檢測，采用YOLOv5m 模型、對全部樣本訓(xùn)練100 次的YOLOv5算法對鋼軌彈條斷裂、缺失、移位、螺栓缺失四種狀態(tài)扣件的識別精度P、召回率R、mAP@.5、mAP@.5：.95 分別為0.988、0.967、0.987、0.822，該方法對鋼軌故障扣件分類檢測具有很好的工程應(yīng)用價值。