基于YOLOv5的鋼材表面缺陷檢測算法

徐明升，?？≥x，干家欣，侯津津，王 圓，周賢勇，陳 琳

(長江大學 計算機科學學院，湖北 荊州 434023)

0 引言

鋼材是國家建設必不可少的重要物資,其應用極為廣泛,尤其在基礎建設當中更是發揮著不可替代的作用?，F階段,大部分鋼材是通過原材料進行壓力加工制成的,具有一定的形狀和性能。在壓力加工過程中,受限于成本、設備和技術等因素,鋼材會出現不可避免的缺陷,如裂紋、內含物、斑塊和點蝕面等,一旦鋼材有了這些缺陷,其抗壓、抗拉伸和耐腐蝕等能力勢必會受到不同程度的影響。因此,缺陷檢測在鋼材出長前尤為重要。

傳統的鋼材缺陷檢測的主要方法有漏磁檢測法、渦流檢測法等。漏磁檢測法在閉合性裂紋上檢測效果不佳;渦流檢測法受環境影響較大,具有一定的局限性。此外,以上傳統方法都需要人工和精密儀器的參與,檢測成本較高。

近年來,隨著人工智能的發展,基于深度學習的檢測方法已逐漸在缺陷檢測行業中使用。目前,主流的深度學習目標檢測算法有以下2類:一階段目標檢測算法和二階段目標檢測算法[1]。二階段目標檢測算法以Faster R-CNN[2]、Mask R-CNN[3]為主;一階段目標檢查算法以Single Shot MultiBox Detector(SSD)[4-5]、You Only Look Once(YOLO)系列[6-8]為主。由于二階段目標檢測算法在檢測過程中會生成一系列候選框,但是一階段目標檢測算法則將其轉換為回歸問題處理,所以二階段目標檢測算法檢測精度相對較高,但是檢測速度相對于一階段目標檢測算法較慢[9]。楊莉等[10]提出了一種特征融合和級聯檢測網絡的Faster R-CNN鋼材表面缺陷檢測算法,能夠較好地完成鋼材表面缺陷的檢測任務,但是在檢測速度上很難滿足工業上的要求。一階段目標檢測算法的精度雖然相對略低,但是擁有較高的檢測速度,十分貼合工業上鋼材表面缺陷檢測的任務。鄒旺等[11]提出了一種改進YOLOv4-tiny的帶鋼表面缺陷實時檢測方法,該方法是基于卷積神經網絡的輕量級目標檢測器,雖然模型尺寸小、檢測速度快,但在NEU-DET數據集上的平均精度均值(mean Average Precision,mAP)為73.29%,精度相對較低。李鑫等[12]提出改進YOLOv5的鋼材表面缺陷檢測算法,其體積減小的同時,進一步提升了一定的檢測速度,在NEU-DET數據集上的mAP值從73.5%提升到了76.8%。

在鋼材表面缺陷檢測應用中,檢測精度在實際應用中尤為重要,因此本文提出一種基于改進YOLOv5的鋼材表面缺陷檢測算法,進一步提升識別準確率。實驗表明,改進后的YOLOv5s算法在NEU-DET數據集上具有良好的檢測性能,并且滿足工業對鋼材表面缺陷檢測實時性的需求。

1 YOLOv5算法簡介

在YOLOv4發布不久,YOLOv5便橫空出世。YOLOv5是一種新穎的檢測算法,該算法在先前版本的基礎上添加了數據增強等新的改進,泛化能力較好、檢測精度高,很受工業的青睞。

目前YOLOv5已經更新到了6.2版本,共有5個模型,按照模型的大小,可分為YOLOv5n、YOLOv5s、YOLOv5m、YOLOv5l、YOLOv5x。本文選擇YOLOv5s作為預訓練模型,YOLOv5s是這5個模型中量級較小的,同時還能保證較為高速度的網絡。將其應用于鋼材表面缺陷檢測,旨在保證擁有一定檢測速度的前提下進一步提升模型的準確度。YOLOv5s網絡結構共分為4個部分,分別是Input(輸入端)、Backbone(主干網絡),Neck(頸部)和Head(頭部)。YOLOv5s的網絡結構如圖1所示。

圖1 YOLOv5s網絡結構Fig.1 Network structure of YOLOv5s

Input采用了圖片尺寸處理、自適應錨框計算和Mosaic數據增強。圖片尺寸處理通過添加少量黑邊的方式將不同長寬的圖片縮放為相應的尺寸;自適應錨框計算是計算預測框與真實框差距后再反向更新,直到獲取最合適的錨框值;Mosaic數據增強將4張圖片進行隨機縮放、排列等操作,組合成1張圖片[13]。

Backbone的任務為特征提取,主要包含了Focus和BottleneckCSP模塊。其中Focus模塊主要對圖片進行切片操作,切片操作可以降低FLOPs,且不影響mAP。BottleneckCSP模塊在切片卷積之后增強整個卷積神經網絡學習性能的同時大幅減少了計算量。

Neck采用了特征金字塔(FPN)[14]與路徑聚合網絡(PAN)[15]聯合的結構,其任務是對來自Backbone提取到的特征進行混合與組合,即結合高層語義信息和底層位置信息,得到更加豐富的特征信息,并把這些特征信息傳遞給Head部分。

Head主要由3個Detect組成,其對應Neck部分輸出的3個特征圖,共輸出大、中、小3個預測分支,之后根據設定的閾值和NMS選出最為合適的檢測結果。

2 改進的YOLOv5算法

改進后的網絡結構如圖2所示,首先在Backbone層引入多個坐標注意力(Coordinate Attention,CA)機制模塊,增強網絡對于重要特征的關注度。使用GhostBottleneck結構與主干網絡中的部分卷積模塊和C3模塊進行替換,構建輕量化網絡。圖中GB-1表示步長為1的GhostBottleneck,用來代替主干網絡中的部分卷積結構,GB-2表示步長為2的GhostBottleneck,用來代替主干網絡中的C3結構。其次在Neck層采用雙向特征金字塔網絡(Bidirectional Feature Pyramid Network,BiFPN)結構替換原Concat結構,提升模型的檢測效果。最后在Detect部分使用K-means算法增加一個尺度為10×10特征層,并接入到相應的Neck部分,提升對缺陷占比較大目標的檢測效果。

圖2 YOLOv5s改進后的網絡結構Fig.2 Improved network structure of YOLOv5s

2.1 引入GhostBottleneck結構

考慮到對模型進行優化會增加額外的參數,導致模型占用內存增大。為減少參數量,獲得更加輕量化的模型,本文將YOLOv5s中Head部分的卷積替換為具有更少參數且對計算精度影響較小的GhostBottleneck[16]。

GhostBottleneck是個可復用模塊,以Ghost模塊為構建基礎,主要由2個Ghost模塊堆疊而成。圖3為Ghost模塊的結構示意,由圖3可以看出,Ghost模塊一開始將Input的特征圖進行卷積,生成的結果與普通卷積中Output的通道數不同,其首先得到特征圖的通道數為輸入特征圖通道數的一半,之后將得到特征圖的每個通道數單獨進行某種線性變換,最后將每個通道的特征圖與線性變換的結果結合起來,得到近似于原始的特征圖。圖4所示為本文所采用的GhostBottleneck結構示意,當Stride=1時,其構建類似于ResNet的瓶頸網絡,由2個堆疊的Ghost模塊組成,第一個Ghost模塊增加通道數,第二個Ghost模塊則將通道數減少為輸入的通道數,以此與shortcut支路相匹配。當Stride=2時,其在2個Ghost模塊中插入了DWConv,主要是為了減少計算量。

圖3 Ghost模塊結構示意Fig.3 Structure of Ghost module

圖4 GhostBottleneck結構示意Fig.4 Structure of GhostBottleneck

在本文的研究中,將主干網絡中的CBS模塊用Stride=2的GhostBottleneck代替,C3模塊用Stride=1的GhostBottleneck代替,采用這種結構對網絡進行優化,提高網絡在實際檢測中的魯棒性。

2.2 增加注意力機制模塊

在鋼材圖像中,有些類型的缺陷像素占比較少,且背景相似度高,網絡容易流失一些缺陷的關鍵特征信息。為了提高網絡對鋼材表面缺陷的關注,實現更精準地檢測,本文在YOLOv5s主干特征提取網絡中引入CA模塊[17]。CA模塊結構如圖5所示。

圖5 CA模塊結構Fig.5 CA module structure

CA是一種新的、效率高的且較為輕量級的注意力機制。該機制在通道注意力中嵌入位置信息,使得模型能夠關注大范圍的位置信息又不會帶來過多的計算量。其足夠的靈活和輕量,能夠簡單地插入到移動網絡的核心結構中。CA的編碼過程有兩部分:Coordinate信息嵌入和CA生成。

首先是Coordinate信息嵌入的部分,假設輸入為x,接下來用(H,1)和(1,W)的池化核分別沿著水平方向和豎直方向做相應處理[18],計算一維平均池化,可以得到高度為h的第c個通道:

(1)

寬度為w的第c個通道的表達式如下:

(2)

其次是CA生成的部分。首先級聯之前模塊生成的2個特征圖,通過卷積模塊后,其維度將會降低為原來的C/r,之后將處理過的特征圖F1送入Sigmoid激活函數得到形如1×(W+H)×C/r的特征圖f,表達式如下:

f=δ(F1([zh,zw]))。

(3)

之后將f切割為2個單獨的張量,分別為fh∈RC/r×H和fw∈RC/r×W,然后fh和fw再分別通過1×1卷積得到和輸入A同樣的通道數,結果如下:

gh=σ(Fh(fh)),

(4)

gw=σ(Fw(fw))。

(5)

最后把空間信息通過加權的方式進行融合,其最終的輸出如下:

(6)

本文對YOLOv5s原始網絡中的Backbone層引入卷積注意力機制模塊CA,旨在強化YOLOv5s網絡對方向和位置等信息的敏感度,增強網絡對于重要特征的關注度,進一步提高目標檢測的準確率。

2.3 引入BiFPN特征融合網絡

在進行特征提取時,伴隨網絡層數的加深,特征信息會有必不可少的丟失,因此,融合不同尺度的特征被廣泛應用于目標檢測任務中,以提高網絡模型的檢測性能。目前已有的特征融合網絡有FPN、路徑聚合網絡(PANet)、BiFPN[19]等。FPN結構是自頂向下傳遞語義特征,如圖6(a)所示。YOLOv5在Neck部分使用了PANet結構,PANet在FPN的基礎上增加了一個自底向上的二次融合路徑,實現上下層語義傳遞及融合,如圖6(b)所示。PANet雙向融合較為簡單,為了Neck能對特征進行更好地融合,本文采用如圖6(c)所示的BiFPN特征融合網絡替換原始結構。

圖6 特征融合網絡Fig.6 Feature fusion network

BiFPN為加權雙向(自頂向下+自底向上)特征金字塔網絡,其主要思想包括高效的雙向跨尺度連接和加權特征圖融合。

① 跨尺度連接。相比較于FPN+PAN,BiFPN在設計上的改變如下:如果1個節點只有1條輸入邊,那么其對特征網絡的貢獻就相對較小,故直接刪除只有一條輸入邊的節點,即P3和P7的中間節點,這樣一來簡化了網絡結構;為融合更多的特征,在相同層的輸入節點和輸出節點之間增加1個跳躍連接;將每個雙向的自頂向下的和自底向上的路徑視作為一個特征網絡層,以獲得更高層次的特征融合。

② 加權特征融合。不同的輸入特征對于特征融合的貢獻是不同的,應當對其進行加權融合,加權共分3種情況:其一是Unbounded Fusion,該方法加入可學習的權重,過程極為簡單,但是會引起訓練的不穩定;其二是Softmax-based Fusion,該方法可以將范圍縮小到[0,1],訓練穩定且效果最好,但是大大增加了訓練時間,不適合用于本實驗的鋼材檢測;其三是Fast Normalized Fusion,該方法和Softmax-based Fusion相似,也能將范圍縮小到[0,1],但是Fast Normalized Fusion中引入了ReLU激活函數保證數值穩定,避免了Softmax操作的資源消耗,效率更高,訓練速度比Softmax-based Fusion快30%。

因此本文將原模型中的FPN+PAN結構替換為BiFPN代替原先的特征,快速進行多尺度特征融合。

2.4 增加檢測層

原YOLOv5s網絡Head層中的主體部分就是3個Detect檢測器,即利用基于網格的Anchor在不同尺度的特征圖上進行目標檢測的過程。當輸入圖像的尺寸為640 pixel×640 pixel時,檢測層3個尺度上的特征圖輸出尺寸分別為20 pixel×20 pixel、40 pixel×40 pixel、80 pixel×80 pixel,用于檢測大、中、小目標。由于數據集中有些類別缺陷占比較大,為加強網絡對這些大目標缺陷的檢測能力,在原YOLOv5s網絡的Detect部分增加一個尺度為10 pixel×10 pixel特征層,并接入到相應的Neck部分。由于增加了檢測層,需對檢測錨框進行重新配置,本文使用K-means聚類方法獲取數據集先驗錨框,錨框配置如表1所示。

表1 錨框配置

3 實驗及結果分析

3.1 數據集

本文通過鋼材表面缺陷數據集(NEU-DET)驗證YOLOv5s-GCBD(GhostBottleneck-CA-BiFPN-Anchor)算法的有效性。

實驗選用的數據集是公開的NEU-DET(鋼材表面缺陷)數據集,包括氧化軋皮(rolled-in_scale)、內含物(inclusion)、裂紋(crazing)、劃痕(scratches)、斑塊(patches)和點蝕面(pitted_surface)六種典型的鋼材表面缺陷。每個類別有300張圖片,共1 800張。圖片分辨率為200 pixel×200 pixel的灰度圖。數據集提供了注釋,標注出了缺陷圖像的位置和類別信息。數據集共有1 800張圖片,隨機抽取訓練集1 440張,驗證集180張,測試集180張,即按照8∶1∶1的比例分配數據集。數據集部分圖像如圖7所示。

圖7 6種類別的缺陷圖像示例Fig.7 Example of defect images in six categories

3.2 實驗環境與參數設定

本實驗在服務器上運行,實驗環境配置如表2所示。模型訓練時,開啟Mosaic數據增強,提高模型的泛化能力;使用余弦退火算法動態調整學習率;采用SGD優化器來進行參數更新,參數設置如表3所示。

表2 實驗環境配置

表3 參數設置

3.3 度量指標

本文從mAP、參數量、P-R曲線(Precision-Recall Curve)、平均推理時間(Speed)對模型進行比較。平均精度綜合考慮了目標檢測的精確率(Precision, P)和召回率(Recall, R),表達式分別如式(7)和式(8)所示。以R為橫軸,P為縱軸,構成P-R曲線,AP即P-R曲線下的面積,如式(9)所示。mAP以IoU閾值為0.5來計算平均精度均值,mAP值越高,模型檢測性能越好,如式(10)所示。

(7)

(8)

(9)

(10)

式中:TP(True Positives)、FP(False Positives)、FN(False Negatives)分別為指正樣本被預測為正例的數量、負樣本被預測為正例的數量、正樣本被預測為負例的數量;N為預測類別數量,PAi為第i類樣本的準確度。

3.4 實驗結果分析

為了驗證改變后的網絡對鋼材表面缺陷檢測的性能有所提升,在NEU-DET進行了消融實驗,結果如表4所示。共設計了5組消融實驗,每組實驗都有相同的參數和網絡環境。

表4 消融實驗結果

分析表4結果可知,增加1個檢測層、引入CA機制、替換特征融合網絡為BiFPN都能提升檢測性能,使用GhostBottleneck結構雖然不能提高檢測性能,但能夠大大減少模型參數數量。第2組實驗在原模型的基礎上增加1個檢測層,可以看出參數量有所增加,平均推理時間增加了1.5 ms,這是由于增加了檢測層導致模型復雜度增加、計算量增加,但是mAP值達到了78.6%,提升了1.3%;第3組實驗引入CA機制后,平均推理時間增加了1.5 ms,mAP值有了明顯提升,為1.5%,證明了引入CA機制雖然會增加模型復雜度、加大計算量,但對鋼材表面缺陷特征具有強化作用,使模型擁有更好的檢測性能;在第4組實驗中替換特征融合網絡為BiFPN,模型mAP值又提升了1.1%,說明BiFPN能夠快速進行多尺度融合,提高鋼材表面缺陷的檢測精度;在第5組實驗中引入GhostBottleneck結構,mAP值略微降低,這是因為GhostNet卷積模塊的引入會造成算法各性能參數的下降,但參數量明顯減少,減少了17%,平均推理時間也減少了1 ms,這得益于GhostBottleneck結構能夠減少計算量,獲得了更輕量化的模型。

YOLOv5s與YOLOv5s-GCBD的各類缺陷mAP值對比如圖8所示,其中all classes表示所有類別mAP值的平均值。改進后與改進前的結果相比,整體的mAP值從76.7%提升到80.2%,提升了3.5%,結果表明改進的模型能夠在多種缺陷類別上提升檢測準確率。6類缺陷中除了pitted_surface類別的AP值略微降低了0.9%外,其他類別均有不同程度的提高,尤其是rolled-in_scale類別的mAP值從63.8%提高到了74.6%,提高了10.8%。scratches類別的mAP值最高,達到92.3%。

圖8 改進前后對比Fig.8 Comparison before and after improvement

實驗結果表明,改進的YOLOv5s-GCBD模型比原YOLOv5s模型的mAP值提升了3.5%,雖然模型大小和參數數量略微增加,但是其檢測速度和檢測精度仍然可以勝任工業上的要求。

3.5 檢測效果對比

為驗證YOLOv5s-GCBD模型的檢測性能,使用原YOLOv5s模型和改進后的YOLOv5s-GCBD模型針對缺陷圖片在測試集上進行了測試,圖9所示為6類缺陷檢測效果對比圖,第一行為ground truth圖,第二行為YOLOv5s算法檢測效果圖,第三行為本文提出的YOLOv5s-GCBD算法檢測效果圖。

圖9 6類缺陷檢測效果對比Fig.9 Comparison of six kinds of defect detection results

從圖示的結果來看,原YOLOv5s模型對于一些缺陷目標存在漏檢的情況,而改進后的YOLOv5s-GCBD模型能夠更好地解決漏檢問題。例如圖中inclusion類中只檢測出了1處目標缺陷,這是由于缺陷目標不明顯,難以檢測出缺陷;改進后的模型檢測出了原模型漏掉的部分缺陷目標,對鋼材表面缺陷的定位更加準確。尤其是rolled-in_scale類別缺陷,改進后的模型檢測效果尤為突出。從整體的檢測效果來看,YOLOv5s-GCBD模型的檢測性能更好,檢測效果有明顯提升。

3.6 對比試驗

為了驗證YOLOv5s-GCBD模型的優勢,將本文算法與SSD、Faster R-CNN、YOLOv4、YOLOv5、文獻[12]在NEU-DET數據集上進行訓練與測試,并應用mAP和模型權重文件大小進行模型性能的對比,對比結果如表5所示。

表5 對比實驗結果

從表5可以看出,本文所提出的YOLOv5s-GCBD模型的mAP值達到了80.2%,比YOLOv5提升了3.5%。相比Faster R-CNN、SSD、YOLOv4、文獻[12]算法,本文算法在檢測準確率和模型大小上均有所優化。綜合上述,本文所提出的模型相比其他主流算法在綜合性能上表現更好。

4 結束語

為了解決傳統鋼材表面缺陷檢測問題,提高檢測精度,本文提出了一種新的基于YOLOv5算法的鋼材表面缺陷檢測模型。該模型增加檢測層,降低漏檢率;引入CA機制,提升特征提取能力;采用BiFPN結構,提高特征融合效率;使用GhostBottleneck結構減少模型的參數量。雖然改進之后的模型比原模型的參數量有略微提升,但在檢測準確率上有更大優勢。后續將繼續優化網絡結構、降低網絡的推理復雜度,并將改進后的模型應用在不同的數據集上進行訓練和檢測效果。