基于改進YOLO v4的熱軋帶鋼表面缺陷檢測

季娟娟，王 佳+，陳亞杰，盧道華，3

(1.江蘇科技大學 機械工程學院，江蘇 鎮江 212100；2.中國船舶重工集團公司 上海船舶設備研究所，上海 200031；3.江蘇科技大學 海洋裝備研究院，江蘇 鎮江 212003)

0 引 言

基于深度學習的檢測方法因對任意復雜函數都能實現較好地擬合且特征提取能力較強，被廣泛地應用于產品缺陷檢測領域[1]。程婧怡等[2]利用改進的YOLO v3[3]算法檢測熱軋鋼條表面缺陷，通過使用DIoU[4]邊框回歸損失并利用K-Means++[5]聚類算法來提取更多小缺陷目標特征。孫永鵬等[6]針對沖壓件缺陷檢測提出了一種基于改進YOLO v4[7]的算法，使用改進的MobileNet v3[8]網絡代替CSPDarknet-53[9]，并結合SE Block[10](sequeze and excitation block)來提高檢測精度。

卷積神經網絡通常提取的是局部特征，不能很好捕捉全局特征信息間的依賴關系，隨著網絡深度的增加，參數量與計算量也會隨之增加，檢測速度變慢。隨著Transformer的快速發展，其憑借著優于CNN(convolutional neural network)的全局信息特征捕獲能力，在CNN的應用領域中也取得了不錯的成果。鑒于CNN與Transformer的不同優勢，本文提出了一種基于改進YOLO v4的熱軋帶鋼表面缺陷檢測方法，主要改進有以下4點：①將YOLO v4的主干網絡CSPDarknet-53替換成MobileViT[11]輕量化網絡，在YOLO模型中嵌入Transformer，在不降低檢測速度的同時，提高檢測精度并大大減少參數量；②利用K-Means++聚類算法來修改先驗框的尺寸，使定位更加準確，減少模型損失；③在PANet[12](path aggregation network)后面添加ASFF[13](adaptively spatial feature fusion)模塊，增強多類型缺陷的特征提取能力；④使用改進的DIoU NMS代替NMS進行后處理操作，減少誤刪量。為了驗證改進的有效性，本文做了多個對比實驗以及對每一個改進方法做了消融實驗。

1 YOLO v4算法

將大小為W×H×S的圖片輸入到YOLO v4中。首先利用CSPDarknet-53進行特征提取；接著通過Neck部分進行多尺度特征融合；然后在Yolo Head部分利用3個不同尺度的特征層進行多尺度檢測并得到初步的預測結果；最后對初步的預測結果做置信度過濾、NMS等后處理操作得到最終的輸出結果，輸出的結果經過式(1)的換算，可得到實際預測框的寬(bw)、高(bh)、中心點坐標 ((bx，by))。 式(1)中的cx、cy為當前錨框所屬的網格左上角坐標值，pw、ph為錨框的寬高，σ為sigmoid激活函數，將函數內的值約束在0～1內YOLO v4的網絡結構如圖1所示。

圖1 YOLO v4網絡結構

(1)

2 改進的YOLO v4算法

為了充分利用輸入圖片全局特征的空間位置與語義之間的關系，本文算法將YOLO v4的主干網絡CSPDarknet-53替換成MobileViT輕量化網絡，并將前兩個MobileViT block(MVIT)和最后的3*CBL模塊提出作為Neck部分的輸入；Neck部分通過SPP結構利用4個不同Pooling Size的最大池化對特征圖進行不同程度的感受野增大，生成4個特征圖后進行特征拼接，接著通過PANet結構進行從下至上，從上至下的多次特征提取拼接，多尺度地融合語義信息和空間信息，最后通過ASFF模塊自適應地融合多尺度特征信息，加強特征的提取；Yolo Head部分通過3個不同大小的特征層進行多尺度檢測，每個特征層由1個CBL模塊和普通卷積組成。本文算法的網絡結構如圖2所示。

圖2 本文算法網絡結構

2.1 替換特征提取網絡

CSPDarknet-53是典型的CNN網絡模型，有較好的空間局部表征信息提取能力，然而計算量大，網絡深度較深，容易丟失重要特征。與CNN不同，Transformer有更好的全局特征學習能力，在識別分類任務中有更高的精度，但是參數量大，在實時檢測中達不到低延遲的要求。而MobileViT則使用Transformer將卷積的局部建模換成全局建模，結合了CNN和Transformer的優勢，在參數和精度方面大幅度超過了MobileNet v3等輕量級網絡架構[11]。

MobileViT主要是由多個MobileViT block和MobileNet v2[14]block組成。MobileViT block的結構如圖3所示。

圖3 MobileVit block結構

Local representations部分主要是用來提取局部空間信息，利用n×n與1×1的卷積對輸入特征進行局部特征提取并投影到高維空間；Global representations部分主要用來學習全局特征，將特征展開為N個non-overlapping flattened patches，再利用Transformer對多個patches之間的關系進行編碼；Fusion部分利用n×n卷積來融合前兩個部分得到的局部、全局特征。

MobileNet v2 block又稱為倒殘差。先通過1×1的卷積升維，進一步提取更豐富的特征；然后利用3×3的深度可分離卷積，在使用較少計算量的情況下來提取更大感受野的特征；再利用1×1的卷積將輸入的特征圖降維，減少參數量；最后再利用殘差連接，使其在較少增加計算量與參數量的同時又可以捕獲到更多的特征信息。結構如圖4所示。

圖4 MobileNet v2 block結構

2.2 采用K-Means++錨框聚類

YOLO v4設定的先驗框尺寸是針對VOC數據集上的圖片聚類得到的，并不適合熱軋帶鋼表面的小尺寸缺陷檢測，因此，要采用聚類算法針對本文所研究的數據集生成更利于模型收斂的先驗框尺寸。聚類算法最常用的是K-Means[15]算法和K-Means++算法，聚類中心初始值設定的好壞對聚類結果有著一定的影響。K-Means算法的聚類中心是提前指定好的，而K-Means++算法是在聚類中心的選擇過程中自動選擇較優的聚類中心，相較于K-Means算法，能更加靈活的生成代表數據集的先驗框尺寸，因此本文采用K-Means++算法對數據集進行聚類。聚類結果為(19，48)，(25，170)，(40，85)，(50，39)，(67，157)，(81，72)，(146，76)，(154，29)，(156，181)。

2.3 添加ASFF模塊

ASFF模塊可通過學習后的權重參數自動地融合多尺度特征信息，加強特征的提取。本文算法使用的是一個三特征層的ASFF結構，如圖5所示。

圖5 三特征層ASFF結構

每個ASFF的輸入為Level 1、Level 2、Level 3。當Level 3特征圖中既有大物體，又有小物體時，小物體的存在會對特征提取有一定的干擾性，此時應該過濾掉小物體的信息，并結合Level 1與Level 2特征圖中含有的高級語義特征信息來豐富Level 3特征圖中的大物體特征信息。以ASFF-3為例，3個層的特征經過變換后乘上相應的權重系數并相加，即可得到新的融合后的特征圖ASFF-3，其計算公式如式(2)所示

ASFF-3=X1→3*α3+X2→3*β3+X3→3*γ3

(2)

式中：X1→3為Level 1特征層通過通道數調整、尺寸縮放，變換至同Level 3特征層相同大小的特征層；同理，X2→3為Level 2變換至同Level 3特征層相同大小的特征層；X3→3為Level 3特征層；α，β，γ為各自的權重系數，在訓練中通過反向傳播獲得。

2.4 改進的DIoU NMS

在通過Yolo Head得到初步的預測結果后，還需進行一系列的后處理操作才能得到真正的預測結果。后處理操作中，對結果影響最大的就是非極大值抑制(NMS)，NMS能過濾掉冗余的預測框，提高檢測的準確度。NMS的具體步驟如下：

步驟1 對所有的候選預測框按照置信度值進行降序排列；

步驟2 選擇置信度值最高的預測框A添加到輸出列表中，并將其從候選列表中刪除；

步驟3 計算框A與其它候選預測框的IoU值，并將大于NMSthreshold的候選預測框視為重疊預測，將其從候選列表中刪除；

步驟4 重復上述3個步驟，直至候選列表為空后，返回輸出列表，即真正的預測框。

(3)

(4)

圖6 DIoU

采用大小不同的λ值進行多次實驗對比，作出不同λ值的mAP(mean average precision)對比，如圖7所示，可知λ取0.83時檢測精度值最高，mAP值為80.16%。

圖7 不同λ值的mAP對比

3 實驗結果與分析

3.1 實驗數據

圖8 6類缺陷

3.2 評價指標

目標檢測的評價指標主要有準確率(Accuracy)、檢測速度(FPS)、召回率(Recall)、精準率(Precision)、平均精度AP(Average Precision)、mAP等等。本實驗主要研究的是工業生產環境下的表面缺陷，既要滿足實時性又要有好的檢測性能，所以采用FPS與mAP作為評價指標來比較各算法。FPS是每秒可檢測完的圖片數量，mAP是各類缺陷平均精度的平均值。mAP、Precision、Recall的計算公式如下

(5)

(6)

(7)

其中，n為缺陷類別數；AP為每類缺陷的平均精度值，可由P-R曲線下的陰影面積獲得；TP為被模型預測為正類的正樣本；FN為被模型預測為負類的正樣本；FP為被模型預測為正類的負樣本。

3.3 實驗結果

本實驗是在ubuntu20.04操作系統下基于pytorh深度學習框架實現的，GPU選用內存12 G的NVIDIA GeForce GTX 3060，CUDA11.1，CUDNN11.1。訓練150個輪數，凍結與解凍時的batchsize設為4，使用Adam優化器進行參數處理，凍結時的學習率為le-3，解凍時的學習率為le-4，訓練集在解凍訓練時采用在線數據增強。

本實驗訓練所得的損失函數曲線如圖9所示。

圖9 損失函數曲線

其中，橫坐標是訓練輪數，縱坐標是平均損失，由圖9可看出，本文算法收斂速度較快。前50個訓練輪數為凍結訓練部分，后100個訓練輪數為解凍訓練部分，驗證集損失在前50個訓練輪數中的趨勢從下降變為開始上升，在第50個訓練輪數處由于模型解凍且學習率改變，驗證集損失與訓練集損失的損失曲線發生輕微的突變，但又迅速下降并開始收斂，尤其是驗證集損失不僅停止了上升的趨勢且變得下降收斂，因此損失函數曲線的整體走向趨勢是下降的。

本實驗算法檢測得到的各類P-R曲線如圖10所示。

圖10 各類P-R曲線

P-R曲線是由橫軸為召回率，縱軸為精準率構成的曲線，由圖10可知，各類的P-R曲線往右上外凸，越接近右上角，P-R曲線下的面積越大，AP值就越高，模型的性能也就越好。

3.4 對比分析與消融實驗

3.4.1 聚類算法對比

將K-Means、K-Means++兩種算法在本實驗數據集上聚類生成的先驗框與原始的YOLO v4先驗框，分別應用在YOLO v4算法上進行對比，對比結果見表1。

表1 3種先驗框的mAP值

由表1可知，K-Means++算法聚類生成的先驗框在YOLOv4上有更高的精度值，分別比原YOLO anchor、K-Means anchor高了2.04%、1.41%，對熱軋鋼表面小缺陷的檢測準確度更高，可提高邊框回歸位置預測的精確度。

3.4.2 YOLO v4算法與本文算法結果對比

將YOLO v4算法與本文算法分別對數據集進行檢測，所得到的AP值的對比如圖11所示。

圖11 兩種算法mAP值對比

由圖11可知，本文算法相較于YOLO v4算法，夾雜類與斑塊類的平均精度均提升了5%，麻點類提升了10%，劃痕類提升了11%，壓入氧化皮類提升了12%，裂紋類更是提升了26%。6類缺陷的檢測精度都得到了不同程度的提升，尤其是裂紋類。

從6類缺陷的檢測結果中，分別抽出一張進行對比，對比的結果如圖12所示。

圖左列的圖(a)是YOLO v4算法的結果，右列的圖(b)是本文算法的檢測結果。由圖12對比可知，YOLO v4算法更容易出現漏檢的情況，且部分缺陷無法識別或者識別出的置信度值較低。而本文算法則能夠盡可能檢測出圖上應有的缺陷。

3.4.3 不同算法對比

為驗證本文算法的綜合性能，用Faster R-CNN、YOLO v3、YOLO v4、YOLO v5、文獻[2]、文獻[6]算法對實驗數據集進行檢測，并與本文算法的結果進行對比，對比的結果見表2。

由表2可知，本文算法的性能均優于Faster R-CNN與YOLO v4，檢測速度雖然比YOLO v3與YOLO v5略低，但精度值最高。文獻[2]采用四特征層輸出，相較于YOLO v3精度值有所提高，但由于使用的是全卷積神經網絡，特征提取不是很全面，因此對檢測精度值的提高低于本文算法。文獻[6]采用MobileNet v3作為主干網絡，由于MobileNet v3能在保證一定檢測精度的前提下擁有較強的提速能力，因此在檢測速度方面超過了本文算法，但由于跟文獻[2]一樣，均采用全卷積神經網絡，所以檢測精度值依然低于本文算法。

本文算法由于結合了Transformer和卷積神經網絡的優勢，具有較強的全局、局部特征提取能力，再加上使用了ASFF模塊，K-Means++算法與改進的DIoU NMS算法來進一步地提高檢測精度，所以本文算法的檢測精度最高。綜合兩個評價指標可知，本文算法的綜合檢測性能最佳，做到了準確度更高的實時檢測。

3.4.4 消融實驗

將本文算法對原YOLO v4算法的每一個改進做消融實驗，結果見表3。

表3 消融比較

由表3可知，僅更換主干網絡，檢測的精度值與速度就能得到提高；利用K-Means++重新聚類先驗框，精度提升的幅度更大；再利用ASFF模塊增強對特征的利用率，精度得以提升，但速度相較于之前有所下降；最后再加上改進的DIoU NMS的替換，進一步提高檢測精度。最終，整個改進的算法使熱軋帶鋼表面缺陷檢測精度值得到了11.57%的提升。

4 結束語

本文針對熱軋帶鋼表面缺陷檢測難度大，識別準確率低等問題提出了一種基于改進YOLO v4的檢測算法，該算法在YOLO v4框架下，結合Transformer技術重構出一個新的深度學習網絡。采用新主干網絡MobileViT并增加ASFF模塊來增強對微小缺陷的特征提取能力；使用K-Means++聚類和改進的DIoU NMS來增強算法對本實驗數據集檢測的魯棒性。總體上是，結合一個工業檢測應用場景，組合兩種主流技術形成的一種解決方案。最后通過多個對比實驗與消融實驗驗證了本文算法改進的有效性。之后，將針對檢測精度與檢測速度進行進一步提升。