基于YOLOv5 改進(jìn)的鋁片表面缺陷檢測(cè)

謝延楠，劉興德

（1.吉林化工學(xué)院 信息與控制工程學(xué)院，吉林吉林，132022；2.吉林化工學(xué)院 機(jī)電工程學(xué)院，吉林吉林，132022）

0 引言

全球各方面的發(fā)展不斷加快，鋁制品的需求量不斷增加，鋁片的質(zhì)量對(duì)于鋁制品的性能具有重要影響。因此，對(duì)鋁片表面缺陷進(jìn)行有效檢測(cè)具有重要的實(shí)際意義。傳統(tǒng)的缺陷檢測(cè)有準(zhǔn)確性差和檢測(cè)速度慢的問題。隨著計(jì)算機(jī)視覺技術(shù)的快速發(fā)展，使用深度學(xué)習(xí)的目標(biāo)檢測(cè)方法成為主流。

隨著其精度和速度不斷提升，服務(wù)于目標(biāo)檢測(cè)需求的神經(jīng)網(wǎng)絡(luò)模型也越來越多，眾多的網(wǎng)絡(luò)模型為工廠車間內(nèi)的工件檢測(cè)提供了新的可能。常見的目標(biāo)檢測(cè)算法分為兩大類，單階段檢測(cè)算法和雙階段檢測(cè)算法[1]，單階段算法以SSD、YOLO 系列為典型，而雙階段則以Faster R-CNN 網(wǎng)絡(luò)為主。其中基于 YOLOv5 的檢測(cè)方法在速度和準(zhǔn)確性上具有一定優(yōu)勢(shì)。

針對(duì)表面缺陷檢測(cè)問題，王子豫等[2]人提出基于改進(jìn)的YOLOv5 算法的零件表面缺陷檢測(cè)技術(shù)，引入CBAM 注意力機(jī)制，并融合BiFPN 網(wǎng)絡(luò)與可變性卷積策略，并采取剪枝方法提升模型性能。杜少聰?shù)萚3]人提出基于改進(jìn)YOLOv5 的鋼軌表面缺陷檢測(cè)技術(shù)，將多頭注意力層嵌入到骨干網(wǎng)絡(luò)末端，并引入全局依賴關(guān)系，構(gòu)建跨層加權(quán)級(jí)聯(lián)結(jié)構(gòu)，將淺層網(wǎng)絡(luò)融入深層網(wǎng)絡(luò)中，提高檢測(cè)性能。吳健生[4]等人提出基于改進(jìn)Faster R-CNN 算法的帶鋼缺陷檢測(cè)，在主干網(wǎng)絡(luò)采用可變形卷積模塊，F(xiàn)PN 多尺度檢測(cè)模塊以及在RPN 網(wǎng)絡(luò)中融合CBAM 模塊，提高檢測(cè)精度。劉艷菊等[5]人提出基于改進(jìn)SSD 的工件表面缺陷檢測(cè)，采用了以ResNet 替換SSD 中原始的VGGNet 的方法，研究了小目標(biāo)檢測(cè)的問題；采用了對(duì)深層特征進(jìn)行反卷積且將深層特征與淺層特征融合的方法，提升網(wǎng)絡(luò)整體性能。

針對(duì)工業(yè)鋁片缺陷數(shù)據(jù)集，本文研究一種基于YOLOv5的網(wǎng)絡(luò)模型，嵌入不同的卷積注意力機(jī)制模塊以增強(qiáng)圖像的特征提取能力，之后修改YOLOv5 本身的CIOU 損失函數(shù)，降低預(yù)測(cè)框的回歸損失。

1 YOLOv5 算法概述

YOLOv5 是一種高效的目標(biāo)檢測(cè)算法，由YOLOv4 算法改進(jìn)而來。YOLOv5 算法具有檢測(cè)速度快、實(shí)時(shí)性高、魯棒性好等優(yōu)點(diǎn)，廣泛應(yīng)用于工業(yè)檢測(cè)領(lǐng)域。YOLOv5 作為YOLO 系列最常用的模型，提供了四種大小的模型。其中YOLOv5s 的模型結(jié)構(gòu)最小，其他幾種模型結(jié)構(gòu)相比更為復(fù)雜，但平均精度會(huì)有所提升。

圖1 為YOLOv5-6.0 版本的模型結(jié)構(gòu)。

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

下面分別對(duì)四個(gè)部分進(jìn)行分解說明：

(1)Input 輸入端：YOLOv5 在輸入階段結(jié)合Mosaic 數(shù)據(jù)增強(qiáng)和自適應(yīng)圖像縮放，增強(qiáng)推理能力，提高網(wǎng)絡(luò)效率，使其能夠很好地適應(yīng)各種輸入。

(2)Backbone 主干網(wǎng)絡(luò)：相比于Yolov5-5.0 版本，Yolov5-6.0 版本使用卷積核代替Focus 模塊，便于模型的導(dǎo)出，且效率更高。該部分三個(gè)部分組成，分別是CBS、CSP1_X、SPPF，主要負(fù)責(zé)對(duì)輸入圖像進(jìn)行特征提取。三個(gè)模塊的詳細(xì)信息如圖1 中所示。

(3)Neck 頸部網(wǎng)絡(luò)：Neck 部分主要是對(duì)之前提取的特征進(jìn)行整理使其能夠被更好地利用從而提取出更有效的特征，采用基于多特征圖融合的FPN+PANet 結(jié)構(gòu)，從自頂向下和自底向上兩個(gè)方向出發(fā)，同時(shí)增強(qiáng)了語義信息和位置信息的提取，同時(shí)設(shè)計(jì)了的CSP2_X 結(jié)構(gòu)，進(jìn)一步加強(qiáng)了網(wǎng)絡(luò)的特征融合能力。

(4)Head：采用CIOU Loss 作為模型的損失函數(shù)，生成邊界框和類別預(yù)測(cè)，同時(shí)采用非極大值抑制NMS 的方法，篩選保留最佳的目標(biāo)框，完成預(yù)測(cè)任務(wù)并記錄缺陷位置及類別[6]。

2 改進(jìn)YOLOv5 網(wǎng)絡(luò)模型

■2.1 注意力機(jī)制

為提高檢測(cè)精度，提升對(duì)定位內(nèi)容的關(guān)注，加入注意力機(jī)制模塊對(duì)模型的性能進(jìn)行優(yōu)化，本文使用的是通道注意力和空間注意力的集成者CBAM[7]，這是一種用于前饋卷積神經(jīng)網(wǎng)絡(luò)的簡(jiǎn)單而有效的注意模塊[8]，CBAM 為輕量級(jí)通用模塊，即插即用。

CBAM 由通道和空間注意力兩個(gè)子模塊組成，圖2 為CBAM 模塊的結(jié)構(gòu)圖。

圖2 CBAM 模型結(jié)構(gòu)圖

在通道注意力模塊（CAM）的結(jié)構(gòu)中，這個(gè)過程的目標(biāo)是通過對(duì)輸入特征圖F 進(jìn)行不同池化方式的處理，然后通過共享的神經(jīng)網(wǎng)絡(luò)來捕捉全局信息，并通過逐元素相加和激活函數(shù)來生成通道注意力特征圖Mc。這個(gè)特征圖可以幫助網(wǎng)絡(luò)更好地聚焦于輸入數(shù)據(jù)中最重要的通道信息，從而改善模型性能。CAM 的模型結(jié)構(gòu)圖如圖3 所示。

圖3 CAM 模型結(jié)構(gòu)圖

空間注意力模塊（SAM），在SAM 中，輸入是通道注意力模塊得到的特征圖Mc 和輸入特征圖 F 進(jìn)行乘法處理后的結(jié)果F'。下面是SAM 的步驟：

首先特征圖 F' 分別進(jìn)行MaxPooling 和AvgPooling 操作，產(chǎn)生兩個(gè)池化特征圖。然后，將處理得到的池化特征圖進(jìn)行通道拼接，將它們?cè)谕ǖ谰S度上連接起來后，將合并后的特征圖輸入一個(gè)卷積操作，將通道數(shù)量減少到1。對(duì)卷積輸出的特征圖進(jìn)行 sigmoid 激活，生成范圍在0 到1 之間的空間注意力特征圖 Ms[9]，圖4 為SAM 的結(jié)構(gòu)圖。

圖4 SAM 模塊結(jié)構(gòu)圖

將得到的空間注意力特征圖Ms 與SAM 的輸入特征圖F'進(jìn)行逐元素乘法操作，最終得到生成的特征。這樣的操作將融合通道和空間注意力的信息。

總之，通道注意力模塊得到的特征圖F'執(zhí)行池化、通道拼接、卷積和激活操作，生成一個(gè)用于強(qiáng)調(diào)輸入特征圖中重要空間位置的SAM 特征圖 Ms。最終，通過將Ms 與該模塊輸入特征圖F'相乘，得到一個(gè)融合了通道和空間的注意力信息的最終特征。這有助于網(wǎng)絡(luò)更好地關(guān)注輸入數(shù)據(jù)中的重要通道和空間信息，從而提升模型性能。

■2.2 損失函數(shù)

損失函數(shù)的作用就是判斷預(yù)測(cè)缺陷和真實(shí)缺陷之間的差距，若預(yù)測(cè)缺陷信息越接近真實(shí)缺陷信息，則損失函數(shù)值越小。YOLOv5 的損失主要包含三個(gè)方面：邊界框損失(bbox_loss)、分類損失(cls_loss)、置信度損失(obj_loss)?？倱p失的表達(dá)式為三個(gè)損失的和：

邊界框損失，關(guān)注的重點(diǎn)是預(yù)測(cè)邊界框和真實(shí)邊界框的交集比例，用交并比即IOU 表示[10]，其中設(shè)預(yù)測(cè)邊界框?yàn)锳，真實(shí)邊界框?yàn)锽，則IOU 表示為：

當(dāng)預(yù)測(cè)邊界框與真實(shí)邊界框完全重疊時(shí)，也就是說它們具有相交部分，這時(shí)候我們可以說它們的IOU（交并比）為1。相反地，如果預(yù)測(cè)邊界框與真實(shí)邊界框沒有相交部分，那么它們的IOU 為0。在YOLOv5 網(wǎng)絡(luò)中，默認(rèn)使用CIOU Loss 來計(jì)算邊界框的損失。相比于沒有改進(jìn)的IOU 損失，CIOU 損失考慮了邊界框的尺寸信息，并綜合考慮了邊界框之間的位置關(guān)系。具體而言，CIOU 損失通過計(jì)算預(yù)測(cè)邊界框與真實(shí)邊界框之間的距離，將這一距離作為損失函數(shù)的一部分，用于優(yōu)化目標(biāo)檢測(cè)算法[11]，以下為CIOU 的損失計(jì)算公式：

其中ρ為兩個(gè)邊界框中心點(diǎn)間的歐氏距離，b和bgt為預(yù)測(cè)邊界框的中心點(diǎn)和真實(shí)框的中心點(diǎn)，c 為預(yù)測(cè)邊界框和真實(shí)框的最小閉包區(qū)域的對(duì)角線距離，α為權(quán)重參數(shù)，v 是用來衡量寬高比的一致性[13]。

分類損失，原網(wǎng)絡(luò)分類損失函數(shù)使用二元交叉熵函數(shù)，該函數(shù)的定義為：

其中y 為輸入樣本對(duì)應(yīng)標(biāo)簽，p 為預(yù)測(cè)樣本為正樣本的概率。

置信度損失，一個(gè)預(yù)測(cè)邊界框?qū)?yīng)一個(gè)置信度值，該值表示接近真實(shí)邊界框的程度，置信度值越高，與真實(shí)框越吻合。計(jì)算置信度損失的函數(shù)是二元交叉熵函數(shù)。

本文使用的是Alpha-IoU Loss，基于現(xiàn)有IoU Loss 的統(tǒng)一冪化，即對(duì)所有的IoU Loss，增加α冪，當(dāng)α等于1時(shí)[13]，則回歸到原始各個(gè)Loss 中：

■2.3 改進(jìn)后的模型

將CBAM 注意力機(jī)制模塊添加至Neck 部分，使網(wǎng)絡(luò)能夠從通道和空間兩個(gè)方面提升提取特征的能力，并且替換Alpha-IoU Loss 的損失函數(shù)，提高性能，更適用于本數(shù)據(jù)集檢測(cè)，圖5 為更改后的網(wǎng)絡(luò)結(jié)構(gòu)模型。

圖5 改進(jìn)后的網(wǎng)絡(luò)模型結(jié)構(gòu)

圖6 缺陷數(shù)據(jù)集

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

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

本文所使用的實(shí)驗(yàn)數(shù)據(jù)來自飛槳訓(xùn)練平臺(tái)上的公開數(shù)據(jù)集，鋁片表面工業(yè)缺陷檢測(cè)數(shù)據(jù)集，共412 張韓有缺陷的鋁片數(shù)據(jù)，共包含4 類缺陷（臟污dirty，褶皺fold，針孔pinhole，擦傷scratch），隨機(jī)對(duì)該數(shù)據(jù)集進(jìn)行劃分，其中80%為訓(xùn)練集和20%為測(cè)試集，使用LabelImg 打標(biāo)簽軟件對(duì)所得數(shù)據(jù)進(jìn)行打標(biāo)簽處理。

本文中的訓(xùn)練使用的環(huán)境配置為：操作系統(tǒng)：Ubuntu 20.04.6 LTS；CPU：Intel(R) CPU @ 2.20GHz；GPU：NVIDIA Tesla T4(16G)；Python：3.10；深度學(xué)習(xí)框架：Pytorch 2.0.1+cu118。

■3.2 性能分析

本文實(shí)驗(yàn)使用平均精度均值mAP 來評(píng)價(jià)模型的準(zhǔn)確性和FPS 用來評(píng)估檢測(cè)的速度，即每秒內(nèi)可以處理的圖片數(shù)量。

3.2.1 對(duì)比試驗(yàn)

對(duì)比試驗(yàn)可以通過和其他網(wǎng)絡(luò)模型進(jìn)行對(duì)比，更好地體現(xiàn)本文中改進(jìn)模型的優(yōu)勢(shì)，本文實(shí)驗(yàn)將YOLOv5-αC（添加CBAM 注意力機(jī)制和Alpha-IOU 損失函數(shù)）與YOLOv3、Faster-RCNN、SSD、YOLOv7 以 及YOLOv5 模塊添加不同種注意力機(jī)制模型進(jìn)行對(duì)比，采用同樣的訓(xùn)練集和測(cè)試集進(jìn)行實(shí)驗(yàn)，不同模型對(duì)該數(shù)據(jù)集的性能對(duì)比如表1 所示。

表1 不同模型對(duì)比實(shí)驗(yàn)

YOLOv5 模型屬于輕量級(jí)模型，并且簡(jiǎn)單易上手，檢測(cè)速度和檢測(cè)精度相對(duì)而言有所提高。雖然Faster-RCNN 網(wǎng)絡(luò)模型的平均檢測(cè)精度很高，但是檢測(cè)速度很慢，無法做到檢測(cè)的實(shí)時(shí)性，因此不考慮該網(wǎng)絡(luò)模型。

3.2.2 消融實(shí)驗(yàn)

該實(shí)驗(yàn)內(nèi)容的主要是明確地對(duì)比出有無注意力機(jī)制模塊和替換損失函數(shù)對(duì)檢測(cè)效果的影響，如表2 所示。

表2 消融實(shí)驗(yàn)結(jié)果

為了明確地看出模型改進(jìn)的效果，采用測(cè)試集的部分圖片進(jìn)行測(cè)試，測(cè)試集圖片在改進(jìn)模型上的檢測(cè)效果如圖7所示。

圖7 檢測(cè)結(jié)果對(duì)比

4 結(jié)論

提高鋁片工業(yè)缺陷的檢測(cè)精度和檢測(cè)速度在生產(chǎn)過程中具有重要意義，針對(duì)鋁片表面工業(yè)缺陷問題，本文提出的一種添加CBAM 注意力機(jī)制和替換損失函數(shù)的改進(jìn)YOLOv5網(wǎng)絡(luò)能夠有效地提高檢測(cè)精度和速度。實(shí)驗(yàn)表明，與原模型比較，平均檢測(cè)精度mAP 上漲3.1%，模型檢測(cè)速度上漲27.784 幀/秒，符合工業(yè)實(shí)時(shí)檢測(cè)要求。