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

曹樂樂，羅 恒，張 鵬

（大連工業大學 機械工程與自動化學院，遼寧 大連 116000）

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

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

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

1 YOLOv5 算法概述

YOLO 系列算法將對象檢測重新定義為一個回歸問題，通過卷積操作，直接在輸出層回歸位置信息和類別信息。YOLOv5 算法有5 種模型結構，其之間的主要區別是網絡的深度和寬度不同，網絡深度和寬度越大，識別精度越高，但計算量也更大。考慮到識別的準確率和識別效率，選取YOLOv5s 網絡進行優化設計，并在NEU-DET 數據集上進行訓練。YOLOv5 網絡結構由4部分組成，分別是Input 輸入層、Backbone 主干網絡層、Neck 特征融合層和Head 預測層[9]。YOLOv5s 的網絡結構圖如圖1 所示。

圖1 YOLOv5s 網絡結構圖

YOLOv5 輸入層包含Mosaic 數據增強、自適應圖片縮放和自適應錨框計算等功能。主干網絡Backbone層的主要作用是特征提取，由Conv、C3、SPPF 等模塊組成。Neck 層采用FPN（特征金字塔網絡）+PAN（金字塔注意力網絡）結構，主要作用是對Backbone 層提取的特征進行融合。Head 層是該算法的預測部分，主要通過對預測框類別、坐標、置信度等信息進行預測。

2 YOLOv5 算法改進

2.1 錨框優化

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

2.2 Backbone 層優化改進

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

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

圖2 Bockbone 層改進前后的網絡結構圖

2.3 Neck 層改進及損失函數優化

YOLOv5 的Neck 層采用FPN+PAN 的網絡結構，該網絡雖然解決了自頂而下的單向信息限制問題，但對多目標檢測任務所表現出的性能不佳，為了提高目標檢測的準確性，本文引用BiFPN 特征金字塔網絡，以實現更佳的特征融合。

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

式中：c 表示最小邊框回歸（bbox）對角線長度，α 取值為3，ρ2α（b，bgt）分別表示為預測框和真實框的中心點的歐式距離。ωgt、hgt表示真實框的寬高，ω、h 表示預測框的寬高。

3 實驗驗證及結果分析

3.1 實驗數據集

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

3.2 配置環境及評測標準

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

本文主要從均平均精度（mAP）這點對模型進行評估，同時考慮精確率（Precision）、召回率（Recall），其表達式如式（4）、式（5）所示。

式中：TP 為正樣本預測為正，FP 為負樣本預測為正，FN 為正樣本預測為負，P 為查準率，R 為查全率，AP 值為PR 曲線圍成的面積，本文中n 等于6。

3.3 對比實驗及結果分析

為了更直觀對比YOLOv5 改進前后的模型，對鋼材缺陷檢測的精準程度，分別將改進前后的YOLOv5模型進行訓練，epoch 設置為300 輪，1 800 張鋼材缺陷數據選用其中的80%作為訓練集，20%作為驗證集，初始學習率為0.01，衰減系數為0.000 5，改進前后訓練結果如圖3、圖4 所示。

圖3 YOLOv5 改進前后PR 曲線圖

圖4 YOLOv5 改進前后檢測結果圖

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

圖5 不同檢測算法在NEU-DET 數據集上的mAP 值

通過對比實驗發現，改進后的YOLOv5 算法相對其他幾種主流算法，檢測性能有明顯優勢。尤其在裂紋和軋制氧化皮2 種較難檢測的缺陷，均有20%左右的mAP 值提高。考慮到精度、速度以及硬件條件，本文算法具有很強的實用性。

4 結論

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