基于YOLOv5 的鐵軌異物入侵檢測(cè)

蔡鑫楠，丁學(xué)文，3，張子怡，劉文艷，宋文文，董國軍

（1 天津職業(yè)技術(shù)師范大學(xué) 電子工程學(xué)院，天津 300222；2 天津市高速鐵路無線通信企業(yè)重點(diǎn)實(shí)驗(yàn)室，天津 300350；3 天津云智通科技有限公司，天津 300350）

0 引言

近年來，人工智能浪潮方興未艾，機(jī)器視覺［1］技術(shù)得到重視。根據(jù)所執(zhí)行任務(wù)的不同，機(jī)器視覺技術(shù)可分為目標(biāo)檢測(cè)和語義分割兩大類。在人工智能技術(shù)發(fā)展過程中，目標(biāo)檢測(cè)作為有關(guān)的基礎(chǔ)研究環(huán)節(jié)，使得研究時(shí)常常見到的良好的深度學(xué)習(xí)網(wǎng)絡(luò)結(jié)構(gòu)也已逐漸受到了學(xué)術(shù)界的廣泛關(guān)注。

目前研究可知，用于軌道異物的目標(biāo)檢測(cè)模型主要分為2 類。一類是RCNN［2-3］系列，該系列目標(biāo)檢測(cè)又分為2 步：第一步是確定候選目標(biāo)的位置，第二步是對(duì)候選目標(biāo)進(jìn)行分類和細(xì)化。RCNN 網(wǎng)絡(luò)模型檢測(cè)實(shí)時(shí)性較低，而軌道異物檢測(cè)對(duì)實(shí)時(shí)性要求較高。另一類是單目標(biāo)識(shí)別模型，主要包括YOLO［4］系列和SSD［5-6］系列。該方法不需要識(shí)別候選區(qū)域，直接通過回歸坐標(biāo)確定目標(biāo)類別的概率和位置。只是已有研究指出，SSD 網(wǎng)絡(luò)對(duì)于小目標(biāo)和集中的異物檢測(cè)效果欠佳，而軌道異物的種類也雜亂無章，大小目標(biāo)兼存。綜上所述可知，基于RCNN 網(wǎng)絡(luò)模型與基于SSD 網(wǎng)絡(luò)模型都具有一定的局限性。本文將基于YOLOv5 的深度學(xué)習(xí)算法模型應(yīng)用于最優(yōu)軌道異物入侵檢測(cè)模型的設(shè)計(jì)，在保證精確度和實(shí)效性的前提下，可對(duì)軌道異物入侵進(jìn)行更好的檢測(cè)，并能起到及時(shí)預(yù)警的作用。

1 YOLOv5 算法原理

YOLOv5 模型在2020 年6 月由Jocher［7］提出。Jocher 在研究了YOLOv3 的基礎(chǔ)上更新了YOLOv5，初版本的YOLOv5 非常迅速、性能高效且使用便捷。盡管YOLOv5 并未在以往的YOLO 模型基礎(chǔ)上提出更為新穎的模型體系或結(jié)構(gòu)改進(jìn)，但是YOLOv5 也仍然提高了目標(biāo)檢測(cè)方法的最新水平。且本文采用的PyTorch 框架即為YOLOv5 引入的新型訓(xùn)練部署框架，這樣也就使得自定義模型訓(xùn)練能夠得到更高的時(shí)效性。

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

研究可知，YOLOv5 網(wǎng)絡(luò)結(jié)構(gòu)由3 個(gè)主要組件組成，對(duì)此可做分別表述如下。

（1）Backbone：在不同圖像細(xì)粒度上聚合并形成圖像特征的卷積神經(jīng)網(wǎng)絡(luò)。

（2）Neck：一系列混合和組合圖像特征的網(wǎng)絡(luò)層，并將圖像特征傳遞到預(yù)測(cè)層。

（3）output：對(duì)圖像特征進(jìn)行預(yù)測(cè)，生成邊界框合并預(yù)測(cè)類別。

對(duì)于YOLOv5，無論是V5s、V5m、V5l、還是V5x的Backbone、Neck 和output 皆一致。唯一的區(qū)別則在于模型的深度和寬度設(shè)置。在coco 數(shù)據(jù)集的測(cè)試結(jié)果見圖1［8］。圖1中，灰色折線為EfficientDet模型，其余4 種為YOLOv5 系列的不同網(wǎng)絡(luò)模型。由圖1 可得，YOLOv5s 模型的精度小、模型小、易于移植，YOLOv5x 模型的精度高、模型大、較為臃腫。本文選用的YOLOv5s 網(wǎng)絡(luò)結(jié)構(gòu)如圖2 所示。

圖1 YOLOv5 權(quán)重文件的測(cè)試Fig.1 Test of YOLOv5 weight file

圖2 YOLOv5s 結(jié)構(gòu)示意圖Fig.2 Schematic diagram of YOLOv5s structure

1.2 YOLOv5 損失函數(shù)

YOLOv5 代碼選用IoU指標(biāo)來評(píng)價(jià)目標(biāo)框和預(yù)測(cè)框的位置損失，用nn.BCEWithLogitsLoss或FocalLoss來評(píng)價(jià)目標(biāo)框和預(yù)測(cè)框的類損失和置信度損失。YOLOv5 代碼用寬高比選擇對(duì)應(yīng)真實(shí)框的預(yù)測(cè)框，且每一個(gè)真實(shí)框?qū)?yīng)3 個(gè)預(yù)測(cè)框。下面擬展開研究分述如下。

（1）位置損失：YOLOv5 代碼用IoU值評(píng)價(jià)預(yù)測(cè)框和真實(shí)框的位置損失，本文給出的CIoU［9］指標(biāo)公式具體如下：

這里，IoU為預(yù)測(cè)框和真實(shí)框的比值，叫并比。v是衡量長(zhǎng)寬比一致性的參數(shù)，也可以定義為：

公式參數(shù)代表意義如圖3 所示。

圖3 公式參數(shù)Fig.3 Formula parameters

（2）置信度損失與類損失。YOLOv5 代碼用nn.BCEWithLogitsLoss或FocalLoss［10］來評(píng)價(jià)目標(biāo)框和預(yù)測(cè)框的類損失和置信度損失。其中，nn.BCEWithLogitsLoss是指先對(duì)預(yù)測(cè)輸出做sigmoid變換，接著求變換后的結(jié)果與真實(shí)值的二值交叉熵。

FocalLoss損失考慮的是：目標(biāo)檢測(cè)中正負(fù)樣本嚴(yán)重不均衡的一種策略。該損失函數(shù)的設(shè)計(jì)思想類似于boosting，降低容易分類的樣本對(duì)損失函數(shù)的影響，注重較難分類的樣本的訓(xùn)練。FocalLoss通過提高難分類別的損失函數(shù)來實(shí)現(xiàn)，此處需用到的數(shù)學(xué)公式可寫為：

對(duì)于YOLO 系列針對(duì)大小目標(biāo)檢測(cè)的精度低問題，YOLOv5 算法使用了更加優(yōu)越的特征提取網(wǎng)絡(luò)，這就提升了檢測(cè)的精確度與實(shí)效性。

2 YOLOv5 目標(biāo)檢測(cè)算法實(shí)現(xiàn)

本文通過無人機(jī)拍攝與自制適用于評(píng)估軌道異物檢測(cè)系統(tǒng)性能的數(shù)據(jù)集，對(duì)YOLOv5 算法模型進(jìn)行訓(xùn)練，并通過所得數(shù)據(jù)源以及所用模型架構(gòu)來評(píng)估選用模型的可靠性，進(jìn)而驗(yàn)證本文設(shè)計(jì)算法的可行性。

2.1 數(shù)據(jù)集構(gòu)建

本文所需的軌道數(shù)據(jù)集在網(wǎng)絡(luò)現(xiàn)有圖片庫中篩選得到9 714張，但帶異物的軌道圖片僅有411張，由于數(shù)據(jù)集較少，無法發(fā)揮深度學(xué)習(xí)的優(yōu)勢(shì)，導(dǎo)致后續(xù)機(jī)器識(shí)別準(zhǔn)確率低，出現(xiàn)了識(shí)別失敗或者識(shí)別錯(cuò)誤的情況。

針對(duì)數(shù)據(jù)集的不足，本文通過圖像處理方法，得到的與原圖像略有異處的新圖像如圖4 所示。再通過視頻剪輯分幀方式與Photoshop 軟件對(duì)圖片進(jìn)行拼接整合，得到新的圖片加入到數(shù)據(jù)集當(dāng)中，擴(kuò)充至1 521張軌道異物訓(xùn)練數(shù)據(jù)集.

圖4 軌道圖像樣本處理Fig.4 Track images samples processing

考慮到天氣影響因子，本文整理出帶有天氣影響因子的軌道異物圖片514張，總訓(xùn)練數(shù)據(jù)集擴(kuò)充至2 035張，部分樣本集展示如圖5 所示。

圖5 部分樣本集展示Fig.5 Partial samples set display

測(cè)試數(shù)據(jù)集則是通過無人機(jī)航拍得到，經(jīng)過篩選清洗有872 張適用，包含軌道可能出現(xiàn)的各種異物以及異常天氣下的軌道異物圖。

2.2 數(shù)據(jù)集標(biāo)注與數(shù)據(jù)集路徑規(guī)劃［11］

本文采用的標(biāo)注工具為labelimg，在labelimg 工具中找到需要識(shí)別的異物進(jìn)行框選并輸入需標(biāo)注異物的名稱，標(biāo)注完成后會(huì)自動(dòng)生成訓(xùn)練所需的xml文件，此文件中包含了軌道異物的位置信息與類別名稱。標(biāo)注過程如圖6 所示，標(biāo)注完成后生成的xml 文件如圖7 所示。

圖6 標(biāo)注過程Fig.6 Marking process

圖7 xml 文件Fig.7 xml file

2.3 標(biāo)簽文件配置與轉(zhuǎn)換

由于數(shù)據(jù)集標(biāo)注后生成的xml 文件不適用于YOLOv5 模型訓(xùn)練，所以通過Python 程序?qū)⒛繕?biāo)信息的xml 文件轉(zhuǎn)換為適用于YOLOv5 模型訓(xùn)練的目標(biāo)信息txt 文件，同時(shí)轉(zhuǎn)換過程中也會(huì)生成適用于YOLOv5 的train、test、val 和trainval 的4 個(gè)txt 文件。轉(zhuǎn)換所得txt 文件如圖8 所示。

圖8 txt 文件Fig.8 txt file

2.4 模型訓(xùn)練

模型訓(xùn)練均是在基于Pytorch 框架的YOLOv5模型上進(jìn)行的，采用的是VGG16 網(wǎng)絡(luò)，用于檢測(cè)的預(yù)訓(xùn)練模型為精度最小的YOLOv5s，訓(xùn)練參數(shù)batch -size為16，迭代次數(shù)為300次，訓(xùn)練過程如圖9 所示。

圖9 數(shù)據(jù)集訓(xùn)練Fig.9 Dataset training

模型訓(xùn)練過程中得到的可視化結(jié)果如圖10 所示。由圖10 可知，數(shù)據(jù)集已對(duì)目標(biāo)進(jìn)行標(biāo)記，可在規(guī)定情況下對(duì)復(fù)雜場(chǎng)景的小目標(biāo)進(jìn)行檢測(cè)。

圖10 數(shù)據(jù)集可視化結(jié)果Fig.10 Dataset visualization results

2.5 算法有效性分析

本文通過3 個(gè)指標(biāo)對(duì)算法的有效性進(jìn)行評(píng)估，分別是：準(zhǔn)確率（precision，P）、召回率（recall，R）、平均精度（mean average precision，mAP）。各指標(biāo)的數(shù)學(xué)定義及運(yùn)算公式可進(jìn)行重點(diǎn)闡述如下。

（1）準(zhǔn)確率（Precision）。找對(duì)的正類／所有找到的正類，其值可由如下公式計(jì)算求出：

（2）召回率（Recall）。找對(duì)的正類／所有本應(yīng)該被找對(duì)的正類，其值可由如下公式計(jì)算求出：

（3）平均精度（mAP）。用來衡量識(shí)別精度，由所有類別AP值求均值得到，其值也可用繪制P -R曲線面積表示，并可由如下公式計(jì)算求出：

在訓(xùn)練后，模型的Precision、Recall、mAP_0.5、mAP_0.5：0.95 結(jié)果如圖11 所示，相應(yīng)的最高值分別可達(dá)到0.998、0.994、0.996、0.925。

圖11 模型性能評(píng)價(jià)指標(biāo)Fig.11 Model performance evaluation index

YOLOv5 模型訓(xùn)練準(zhǔn)確率和置信度參數(shù)評(píng)估如圖12 所示。圖12 表明在模型訓(xùn)練過程中，開始時(shí)的epoch訓(xùn)練中精度提升明顯，在驗(yàn)證集已取得不錯(cuò)的表現(xiàn)。

圖12 YOLOv5 參數(shù)評(píng)估圖Fig.12 YOLOv5 parameters evaluation diagram

將訓(xùn)練好的模型與其他檢測(cè)模型、如Faster RCNN 和YOLOv3 模型做對(duì)比，結(jié)果見表1。由表1可以看出，YOLOv5 無論是從精度、還是召回率都好于其他2種，且YOLOv5 作為YOLO 系列最新的模型檢測(cè)網(wǎng)絡(luò)，可以較好地進(jìn)行實(shí)時(shí)檢測(cè)，符合軌道異物檢測(cè)的要求，故選取YOLOv5 為本文檢測(cè)模型。

表1 模型對(duì)比結(jié)果Tab.1 Comparison results of the model %

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

基于機(jī)器訓(xùn)練結(jié)果，使用本文搜集的測(cè)試數(shù)據(jù)集進(jìn)行評(píng)估測(cè)試，部分測(cè)試結(jié)果如圖13 所示，根據(jù)每張圖的類別概率判別，能夠準(zhǔn)確地進(jìn)行分類，且有較高的識(shí)別率。

圖13 YOLOv5 部分測(cè)試結(jié)果Fig.13 Partial test results of YOLOv5

3 結(jié)束語

本文針對(duì)軌道異物入侵，結(jié)合機(jī)器視覺與深度學(xué)習(xí)網(wǎng)絡(luò)模型，對(duì)軌道內(nèi)部出現(xiàn)的異物進(jìn)行識(shí)別，相較于人工檢測(cè)提高了目標(biāo)識(shí)別的精確度，同時(shí)也大大縮減了時(shí)間。高性能的目標(biāo)檢測(cè)算法是防止異物入侵軌道的關(guān)鍵技術(shù)，也是實(shí)現(xiàn)軌道安全通行的強(qiáng)有力保障。本文針對(duì)異物特征在識(shí)別準(zhǔn)確率與模型運(yùn)行時(shí)長(zhǎng)上，選擇了最適用于軌道異物識(shí)別的YOLOv5 算法模型，對(duì)目標(biāo)識(shí)別的精確度有較大提升，對(duì)軌道異物識(shí)別所要求的實(shí)時(shí)性也有較好的效果。研究成果對(duì)于軌道領(lǐng)域的異物檢測(cè)有較好的借鑒價(jià)值。但是在軌道中具有較多異物目標(biāo)時(shí)，卻也面對(duì)著容易對(duì)小目標(biāo)及被遮擋異物出現(xiàn)漏檢的情況，因此后續(xù)研究將在這方面做進(jìn)一步的改善。