基于YOLOv5 算法改進(jìn)的鋼材表面缺陷檢測(cè)

曹樂(lè)樂(lè)，羅 恒，張 鵬

（大連工業(yè)大學(xué) 機(jī)械工程與自動(dòng)化學(xué)院，遼寧 大連 116000）

在機(jī)械設(shè)計(jì)及制造領(lǐng)域，鋼材是廣泛應(yīng)用的金屬材料，也是重要的工業(yè)原料，廣泛應(yīng)用于刀具制造、航空航天、國(guó)防科技等領(lǐng)域。隨著我國(guó)制造業(yè)水平的不斷提高，鋼材加工技術(shù)也取得了巨大進(jìn)步[1]。但是在鋼材加工中也難免會(huì)產(chǎn)生刮痕、裂紋、夾雜等缺陷，這些缺陷對(duì)鋼材產(chǎn)品的耐磨性、抗腐性，強(qiáng)度和硬度等物理性能產(chǎn)生影響。因此，為了保證鋼材產(chǎn)品的安全性能，需要對(duì)鋼材表面進(jìn)行缺陷檢測(cè)[2]。

傳統(tǒng)鋼材缺陷檢測(cè)主要依靠人力觀察，這種檢測(cè)方式極易使檢測(cè)人員產(chǎn)生視覺(jué)疲勞，從而出現(xiàn)誤檢漏檢的情況，且整體檢測(cè)效率不高。隨著計(jì)算機(jī)視覺(jué)技術(shù)的不斷發(fā)展，利用圖像處理技術(shù)進(jìn)行鋼材缺陷檢測(cè)成為了主流?；谏疃葘W(xué)習(xí)的目標(biāo)檢測(cè)算法可以分為兩大類，一類是以R-CNN、Faster R-CNN 等為代表的兩階段目標(biāo)檢測(cè)算法，另一類是以SSD、YOLO[3]等為代表的單階段目標(biāo)檢測(cè)算法。He 等[4]通過(guò)Faster-RCNN 算法提出了一種鋼材缺陷檢測(cè)方法，并在東北大學(xué)開(kāi)源數(shù)據(jù)集NEU-DET 上進(jìn)行訓(xùn)練，檢測(cè)任務(wù)的mAP 值達(dá)到82.3%，但此方案的檢測(cè)速度難以滿足工業(yè)領(lǐng)域的檢測(cè)需求。Cheng 等[5]結(jié)合通道注意力機(jī)制，利用自適應(yīng)空間特征融合的方法，在NEU-DET 數(shù)據(jù)集上進(jìn)行訓(xùn)練，其mAP 值達(dá)到78.25%，但是該模型相對(duì)于其他模型而言，檢測(cè)精度較低。李維剛等[6]通過(guò)改進(jìn)YOLOv3 算法模型，并在NEU-DET 數(shù)據(jù)集上進(jìn)行訓(xùn)練，得到的平均精度（mAP）值為80%，同時(shí)檢測(cè)速度保持在50 fps 左右，但是該網(wǎng)絡(luò)模型參數(shù)較大，難以部署。

本文針對(duì)鋼材表面常見(jiàn)的缺陷問(wèn)題，提出一種以YOLOv5 為基礎(chǔ)框架的改進(jìn)算法，并在NEU-DET 數(shù)據(jù)集上進(jìn)行算法驗(yàn)證。改進(jìn)算法主要在Backbone 層引入RepVGG[7]和BoTNet[8]網(wǎng)絡(luò)以提高特征提取能力。并在Neck 層引入BiFPN 網(wǎng)絡(luò)以替代原有的特征金字塔網(wǎng)絡(luò)。為了獲得更好的檢測(cè)效果，對(duì)K-means 聚類算法、損失函數(shù)等進(jìn)行了優(yōu)化，提高了回歸精度和模型的魯棒性。

1 YOLOv5 算法概述

YOLO 系列算法將對(duì)象檢測(cè)重新定義為一個(gè)回歸問(wèn)題，通過(guò)卷積操作，直接在輸出層回歸位置信息和類別信息。YOLOv5 算法有5 種模型結(jié)構(gòu)，其之間的主要區(qū)別是網(wǎng)絡(luò)的深度和寬度不同，網(wǎng)絡(luò)深度和寬度越大，識(shí)別精度越高，但計(jì)算量也更大?？紤]到識(shí)別的準(zhǔn)確率和識(shí)別效率，選取YOLOv5s 網(wǎng)絡(luò)進(jìn)行優(yōu)化設(shè)計(jì)，并在NEU-DET 數(shù)據(jù)集上進(jìn)行訓(xùn)練。YOLOv5 網(wǎng)絡(luò)結(jié)構(gòu)由4部分組成，分別是Input 輸入層、Backbone 主干網(wǎng)絡(luò)層、Neck 特征融合層和Head 預(yù)測(cè)層[9]。YOLOv5s 的網(wǎng)絡(luò)結(jié)構(gòu)圖如圖1 所示。

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

YOLOv5 輸入層包含Mosaic 數(shù)據(jù)增強(qiáng)、自適應(yīng)圖片縮放和自適應(yīng)錨框計(jì)算等功能。主干網(wǎng)絡(luò)Backbone層的主要作用是特征提取，由Conv、C3、SPPF 等模塊組成。Neck 層采用FPN（特征金字塔網(wǎng)絡(luò)）+PAN（金字塔注意力網(wǎng)絡(luò)）結(jié)構(gòu)，主要作用是對(duì)Backbone 層提取的特征進(jìn)行融合。Head 層是該算法的預(yù)測(cè)部分，主要通過(guò)對(duì)預(yù)測(cè)框類別、坐標(biāo)、置信度等信息進(jìn)行預(yù)測(cè)。

2 YOLOv5 算法改進(jìn)

2.1 錨框優(yōu)化

YOLOv5 錨框是依據(jù)COCO 數(shù)據(jù)集，進(jìn)行Kmeans 聚類得到，為了更好地對(duì)鋼材進(jìn)行缺陷檢測(cè)，本文針對(duì)NEU-DET 數(shù)據(jù)集進(jìn)行K-means++算法聚類，重新設(shè)計(jì)錨框的大小。K-means++算法相對(duì)于Kmeans 算法可以更好地獲得聚類中心，防止聚類中心局部化。通過(guò)K-means++聚類算法最終得到的9 組錨框尺寸分別為（19，45）、（21，170）、（37，90）、（52，46）、（58，173）、（78，88）、（128，50）、（147，93）、（158，188）。

2.2 Backbone 層優(yōu)化改進(jìn)

2021 年，丁霄漢等通過(guò)在VGG 網(wǎng)絡(luò)模型中引入identity 殘差分支，提出RepVGG 網(wǎng)絡(luò)結(jié)構(gòu)，該結(jié)構(gòu)在訓(xùn)練時(shí)使用殘差結(jié)構(gòu)，在推理時(shí)使用VGG 結(jié)構(gòu)，RepVGG 網(wǎng)絡(luò)結(jié)構(gòu)既可以提高模型推理速度，又不會(huì)有任何的性能損失。RepVGG 訓(xùn)練模型轉(zhuǎn)化為推理模型需要結(jié)構(gòu)重參數(shù)化，重參數(shù)化主要分為兩步，第一步將卷積層和BN 層融合為一個(gè)3×3 卷積層，第二步將3個(gè)分支上的3×3 卷積層融合成一個(gè)3×3 卷積層。本文在主干網(wǎng)絡(luò)中引入RepVGG 模型替代原有的卷積層，同時(shí)兼顧了特征提取的能力和特征提取的速度，更有利于整個(gè)模型的部署。

BoTNet（Bottleneck Transformer Network）是2021年UC Berkeley 和谷歌團(tuán)隊(duì)合作發(fā)布，其中核心模塊是ResNet Bottleneck 中被多頭注意力機(jī)制（MHSA）替代3×3 的卷積，形成新的網(wǎng)絡(luò)結(jié)構(gòu)。MHSA 的輸入尺寸為H×W×d，分別表示輸入特征矩陣的高寬以及單個(gè)維度，維度數(shù)量為H×W。q、k、v、r 均為輸入特征，分別代表查詢、鍵、值的矩陣和位置編碼。Rh和Rw分別表示高度和寬度不同位置的位置編碼。BoTNet 是一個(gè)概念簡(jiǎn)單但功能強(qiáng)大的網(wǎng)絡(luò)架構(gòu)，將自注意結(jié)合到多個(gè)計(jì)算機(jī)視覺(jué)任務(wù)中，對(duì)小物體檢測(cè)有顯著提升[8]，所以該網(wǎng)絡(luò)可以更好識(shí)別鋼材缺陷，提高識(shí)別準(zhǔn)確率。Backbone 層通過(guò)更換卷積層，增添注意力機(jī)制，得到如圖2所示的主干網(wǎng)絡(luò)。

圖2 Bockbone 層改進(jìn)前后的網(wǎng)絡(luò)結(jié)構(gòu)圖

2.3 Neck 層改進(jìn)及損失函數(shù)優(yōu)化

YOLOv5 的Neck 層采用FPN+PAN 的網(wǎng)絡(luò)結(jié)構(gòu)，該網(wǎng)絡(luò)雖然解決了自頂而下的單向信息限制問(wèn)題，但對(duì)多目標(biāo)檢測(cè)任務(wù)所表現(xiàn)出的性能不佳，為了提高目標(biāo)檢測(cè)的準(zhǔn)確性，本文引用BiFPN 特征金字塔網(wǎng)絡(luò)，以實(shí)現(xiàn)更佳的特征融合。

YOLOv5 采用的是LoosGIOU損失函數(shù)，該損失函數(shù)雖然把IOU（預(yù)測(cè)框與真實(shí)框的交集比上并集）的問(wèn)題解決了，但還是基于面積層面度量，并沒(méi)有把2 個(gè)矩形框的距離，寬高比等問(wèn)題考慮進(jìn)去。為了更加精準(zhǔn)定位預(yù)測(cè)框，本文選用Lossα-CIOU損失函數(shù)代替LoosGIOU損失函數(shù)，定義公式如式（1）所示。

式中：c 表示最小邊框回歸（bbox）對(duì)角線長(zhǎng)度，α 取值為3，ρ2α（b，bgt）分別表示為預(yù)測(cè)框和真實(shí)框的中心點(diǎn)的歐式距離。ωgt、hgt表示真實(shí)框的寬高，ω、h 表示預(yù)測(cè)框的寬高。

3 實(shí)驗(yàn)驗(yàn)證及結(jié)果分析

3.1 實(shí)驗(yàn)數(shù)據(jù)集

本文是對(duì)東北大學(xué)公開(kāi)數(shù)據(jù)集NEU-DET 進(jìn)行訓(xùn)練驗(yàn)證，該數(shù)據(jù)集含有1 800 張灰度圖像且分為6類的表面缺陷，每類缺陷含有300 張灰度圖，這6 類缺陷分別是軋制氧化皮（Rolled-in Scale，RS）、斑塊（Patches，Pa）、劃痕（Scratches，Sc）、裂紋（Crazing，Cr）、夾雜（Inclusion，In）以及表面點(diǎn)蝕（Pitted Surface，PS）。

3.2 配置環(huán)境及評(píng)測(cè)標(biāo)準(zhǔn)

本文是基于Windows10 操作系統(tǒng)，顯卡為NVDIA GeForce RTX1650Ti，編譯環(huán)境為PyCharm，軟件環(huán)境基于PyTorch1.10.1 搭建，python 版本3.8，CUDA 版本11.6。

本文主要從均平均精度（mAP）這點(diǎn)對(duì)模型進(jìn)行評(píng)估，同時(shí)考慮精確率（Precision）、召回率（Recall），其表達(dá)式如式（4）、式（5）所示。

式中：TP 為正樣本預(yù)測(cè)為正，F(xiàn)P 為負(fù)樣本預(yù)測(cè)為正，F(xiàn)N 為正樣本預(yù)測(cè)為負(fù)，P 為查準(zhǔn)率，R 為查全率，AP 值為PR 曲線圍成的面積，本文中n 等于6。

3.3 對(duì)比實(shí)驗(yàn)及結(jié)果分析

為了更直觀對(duì)比YOLOv5 改進(jìn)前后的模型，對(duì)鋼材缺陷檢測(cè)的精準(zhǔn)程度，分別將改進(jìn)前后的YOLOv5模型進(jìn)行訓(xùn)練，epoch 設(shè)置為300 輪，1 800 張鋼材缺陷數(shù)據(jù)選用其中的80%作為訓(xùn)練集，20%作為驗(yàn)證集，初始學(xué)習(xí)率為0.01，衰減系數(shù)為0.000 5，改進(jìn)前后訓(xùn)練結(jié)果如圖3、圖4 所示。

圖3 YOLOv5 改進(jìn)前后PR 曲線圖

圖4 YOLOv5 改進(jìn)前后檢測(cè)結(jié)果圖

本文基于YOLOv5 改進(jìn)的檢測(cè)算法在NEU-DET數(shù)據(jù)集上的mAP@0.5 的值達(dá)到了88.5%，相對(duì)于改進(jìn)前的76.8%，提高了11.7%，尤其對(duì)較難檢測(cè)的裂紋缺陷，識(shí)別精度提高了18.7%，檢測(cè)速度為70.1 fps。改進(jìn)后的模型不僅測(cè)精度更高，而且滿足實(shí)時(shí)檢測(cè)的需求。為了驗(yàn)證改進(jìn)后YOLOv5_LH 模型對(duì)鋼材缺陷檢測(cè)的優(yōu)勢(shì)，分別用其他幾種主流的深度學(xué)習(xí)模型進(jìn)行檢測(cè)，并作柱狀圖直觀比較，如圖5 所示。

圖5 不同檢測(cè)算法在NEU-DET 數(shù)據(jù)集上的mAP 值

通過(guò)對(duì)比實(shí)驗(yàn)發(fā)現(xiàn)，改進(jìn)后的YOLOv5 算法相對(duì)其他幾種主流算法，檢測(cè)性能有明顯優(yōu)勢(shì)。尤其在裂紋和軋制氧化皮2 種較難檢測(cè)的缺陷，均有20%左右的mAP 值提高?？紤]到精度、速度以及硬件條件，本文算法具有很強(qiáng)的實(shí)用性。

4 結(jié)論

本文基于YOLOv5 算法提出一種改進(jìn)的鋼材缺陷檢測(cè)模型，該模型在主干網(wǎng)絡(luò)上利用RepVGG 更換卷積層；并在SPPF 層前增添BoTNet 注意力機(jī)制層；在Neck 層使用BiFPN 特征金字塔網(wǎng)絡(luò)，且優(yōu)化了損失函數(shù)和聚類算法。主干網(wǎng)絡(luò)的更改，可以充分利用Backbone 層的特征提取能力，Neck 層的優(yōu)化，可以更好地對(duì)主干網(wǎng)絡(luò)提取的特征進(jìn)行融合，損失函數(shù)優(yōu)化，提高了回歸精度，聚類算法優(yōu)化可以增強(qiáng)模型的魯棒性。YOLOv5_LH 算法訓(xùn)練集mAP 值達(dá)到88.5%，檢測(cè)速度能達(dá)到70.1 fps，既保證了檢測(cè)精度，又保證了檢測(cè)效率，對(duì)鋼材缺陷檢測(cè)有重要的現(xiàn)實(shí)意義。該算法也具有較強(qiáng)的可移植性，對(duì)刀具等磨損缺陷檢測(cè)也有很高的使用價(jià)值。