基于改進YOLOv5的鋁型材瑕疵檢測算法

劉 柱,董 琴,楊國宇,陳朝峰

(鹽城工學(xué)院 信息工程學(xué)院,江蘇 鹽城 224001)

0 引 言

伴隨“中國制造2025”的提出,鋁型材憑借其獨特的質(zhì)輕,柔韌性好等特點,在機械工業(yè)領(lǐng)域有著不可動搖的地位,其全行業(yè)的產(chǎn)量和消費量在世界范圍內(nèi)逐年遞增[1]。在加工速度上漲的同時,由于環(huán)境等不可抗力因素,使得產(chǎn)品表面不可避免地產(chǎn)生瑕疵問題,這不僅僅是鋁型材外觀問題,它嚴重影響到了鋁型材在應(yīng)用時的可靠性和安全性。

通過調(diào)研發(fā)現(xiàn),傳統(tǒng)的檢測方法成本高、操作復(fù)雜度高且對環(huán)境要求苛刻,無法實現(xiàn)可視化過程檢測。伴隨工業(yè)化的發(fā)展,機器視覺檢測憑借精確、無接觸測量等優(yōu)點被廣泛應(yīng)用于瑕疵表面檢測[2]。早期利用機器視覺進行表面瑕疵檢測,主要是通過對缺陷照片進行圖像處理提取瑕疵特征,這類方法能對特定瑕疵較準確地進行檢測[3]。但對于鋁型材瑕疵種類較多、尺度不一的情況,存在泛性較差、特征提取復(fù)雜、檢測精確度不高等問題,因此在實際生產(chǎn)中局限性太大,難以滿足生產(chǎn)需求。

近些年,基于深度學(xué)習(xí)的瑕疵檢測算法因具有穩(wěn)定高效、魯棒性強的特點,在很多場景中得到了廣泛的應(yīng)用,取得了矚目的成績,并逐漸取代傳統(tǒng)表面瑕疵檢測方法。YOLO(You Only Look Once)算法的檢測性能更為突出,在YOLO的基礎(chǔ)上涌現(xiàn)出了YOLOv3[4]、YOLOv4[5]等眾多改進算法。例如,劉浩等人在機床刀具的檢測過程中采用了改進的YOLOv3算法,取代了傳統(tǒng)的人工目測[6]。孫永鵬等人采用改進YOLOv4的檢測方法提升沖壓件的檢測效率和檢測精度并取得了良好的效果[7]。喬歡歡等人使用YOLOv5算法進行交通標志檢測,取得了良好的檢測效果[8]。

上述研究結(jié)果表明,YOLO算法在瑕疵檢測方面具有良好的魯棒性。因此,針對準確快速的檢測需求,提出一種基于改進YOLOv5(You Only Look Once Version-5)的鋁型材瑕疵檢測算法。首先,在heard部分通過K-means++算法設(shè)計新的檢測端,提升小尺度瑕疵的檢測效果。其次,改進原Backbone網(wǎng)絡(luò)結(jié)構(gòu),加入新的卷積層,并插入改進注意力機制E-CBAM,增強整個模型網(wǎng)絡(luò)特征提取能力,減少原始信息丟失。最后,使用EIoU Loss(Efficient IoU Loss)作為整個網(wǎng)絡(luò)結(jié)構(gòu)的損失函數(shù),加快收斂效率,從而解決難易樣本不平衡的問題。通過評估不同改進策略對檢測模型的影響,及真實產(chǎn)線下數(shù)據(jù)集實驗數(shù)據(jù)對比不同檢測算法,確保改進YOLOv5算法在滿足工廠實際需求的同時,兼顧檢測精度和檢測效率。

1 改進YOLOv5算法模型

1.1 改進YOLOv5網(wǎng)絡(luò)模型

YOLOv5是一階段的目標檢測算法,主要模型結(jié)構(gòu)可分為Input(輸入端)、Backbone(骨干網(wǎng)絡(luò))、Neck(頸部網(wǎng)絡(luò))及Head(輸出網(wǎng)絡(luò))。

Input主要用于對輸入數(shù)據(jù)進行規(guī)范化調(diào)整、數(shù)據(jù)增強、自適應(yīng)錨框計算。當(dāng)數(shù)據(jù)進入網(wǎng)絡(luò)進行訓(xùn)練時,輸入端會起到保持圖片的大小一致作用,當(dāng)輸入數(shù)據(jù)量不夠時,會自動對圖像進行裁剪、反轉(zhuǎn)擴充數(shù)據(jù)量,在訓(xùn)練階段在初始錨點框的基礎(chǔ)上輸出對應(yīng)的預(yù)測框。Backbone主要用于輸入數(shù)據(jù)的特征提取,YOLOv5使用CSPDarknet作為骨干網(wǎng)絡(luò),在不同圖像細粒度上聚合并形成圖像特征的卷積神經(jīng)網(wǎng)絡(luò)。Neck網(wǎng)絡(luò)借鑒CSPnet[9]研發(fā)出CSP2結(jié)構(gòu)實現(xiàn)不同深度的特征圖信息融合,增強神經(jīng)網(wǎng)絡(luò)對特征信息的表達能力,有效提高網(wǎng)絡(luò)的特征融合性能。Head以CIOU Loss作為bounding box損失函數(shù),使用三個不同尺度的Detect檢測端在特征圖上進行瑕疵位置確定及種類分別的預(yù)測。

為使得網(wǎng)絡(luò)模型達到檢測更精確、檢測尺度更廣、速度更快的瑕疵檢測效果,針對鋁型材瑕疵檢測模型結(jié)構(gòu)的改進將通過改進原有Backbone網(wǎng)絡(luò)結(jié)構(gòu)、加入優(yōu)化注意力機制(E-CBAM)增強特征信息提取、添加新的檢測端擴大識別尺度且錨框使用K-means++算法聚類生成、采用EIOU邊界損失函數(shù)加快網(wǎng)絡(luò)收斂的方法實現(xiàn)[10],改進后的網(wǎng)絡(luò)結(jié)構(gòu)模型如圖1所示。

圖1 改進后的YOLOv5網(wǎng)絡(luò)框圖

1.2 改進檢測尺度及錨框優(yōu)化

1.2.1 檢測尺度優(yōu)化

原有的YOLOv5模型結(jié)構(gòu)采用三個檢測端,以本實驗為例,若輸入的圖片數(shù)據(jù)為640×640,通過卷積池化等操作步驟最終得到的檢測尺度為80×80、40×40、20×20,這對于一般的檢測任務(wù)是較為合理的,但對于鋁型材表面瑕疵尺度多樣化,且存在尺度差較大的瑕疵類別,為提升檢測精度,在Neck網(wǎng)絡(luò)結(jié)構(gòu)中再一次進行上采樣及張量拼接,將淺層特征與深層特征融合增大檢測特征尺度10×10,同時在Backbone網(wǎng)絡(luò)結(jié)構(gòu)中增添一組CSP及CBS模塊。

相比原有檢測結(jié)構(gòu),新的檢測端進一步加強了對特征信息的提取,并將淺層特征進一步融合得到更加豐富、高效的語義信息,拓寬感受野得到更多有用的特征信息,從而有效提升模型檢測精度。

1.2.2 檢測錨框優(yōu)化

每個檢測端的檢測錨框采用K-means++算法聚類生成[11],原始算法的自適應(yīng)錨框只適用于常規(guī)尺寸的數(shù)據(jù)集,而本次實驗所用的鋁型材瑕疵數(shù)據(jù)集瑕疵復(fù)雜程度較大,尺度大小差距較大,采用K-means++算法減小真實框與預(yù)設(shè)框的差距,避免發(fā)生錯檢、漏檢現(xiàn)象,提升檢測精度。

相比于傳統(tǒng)的K-means[12]算法,K-means++算法使用IOU(Intersection Over Union)也就是預(yù)測框與真實框的交并比來計算樣本與聚類中心的距離,從而控制不同大小邊框的誤差,公式如下:

DiStance(box,cluster)=1-IOU(box,cluster)

(1)

式中,box表示真實框集合,cluster表示聚類中心集合,IOU(box,cluster)表示真實框與聚類中心交集與并集的比例。

IOU值越高表示真實框與聚類中心越接近,相關(guān)度也就越高。初始化簇中心的方法為逐個選取k個簇中心,且離其它簇中心越遠的樣本點越有可能被選為下一個簇中心。計算公式如下:

(2)

其中,α表示數(shù)據(jù)集,D(X)表示每個樣本與當(dāng)前已有聚類中心之間的最短距離,P(X)表示下一個聚類中心的概率大小。

最終利用聚類分析得到最優(yōu)的12個不同比例大小的錨框參數(shù)避免收斂速度慢的現(xiàn)象,其聚類結(jié)果和分配策略如表1所示。

表1 聚類結(jié)果及分配策略

1.3 改進注意力機制E-CBAM

注意力機制(Attention Mechanism)能夠根據(jù)需求選擇最合適的輸入,能夠?qū)R兩個序列之間的token關(guān)系,從而實現(xiàn)更好的效果。為了提升YOLOv5的網(wǎng)絡(luò)表達能力,在Backbone網(wǎng)絡(luò)中對CBAM[13]注意力機制進行改進,并命名為E-CBAM(Efficient Convolutional Block Attention Module)。

如圖2所示,為選出針對當(dāng)前數(shù)據(jù)集最合適的注意力機制,在Backbone網(wǎng)絡(luò)結(jié)構(gòu)中分別加CA[14]、ECA[15]、SE[16]及E-CBAM注意力機制后,通過pytorch與GradCAM[17]實現(xiàn)熱力圖可視化。通過比較特征圖層的熱力圖可以明顯看出,E-CBAM注意力機制針對當(dāng)前數(shù)據(jù)集類別檢測在精度和準確度要優(yōu)于其他注意力機制。

圖2 不同注意力機制熱力圖效果

通道注意力模塊CAM(Channel Attention Module)[18]和經(jīng)過優(yōu)化后的空間注意力模塊E-SAM(Efficient Spatial Attention Module)兩個子模塊共同組成E-CBAM(Efficient Convolutional Block Attention Module),結(jié)構(gòu)如圖3所示。

圖3 E-CBAM結(jié)構(gòu)

從圖中可知,輸入信息先通過一個通道注意力模塊得到加權(quán)結(jié)果之后,再經(jīng)過一個空間注意力模塊,最終進行加權(quán)得到結(jié)果。

改進后的空間注意力模塊結(jié)構(gòu)如圖4所示?？臻g注意力模塊的主要工作包括:將Channel attention模塊輸出的特征圖信息作為輸入數(shù)據(jù);將特征信息進行基于channel的global max pooling和global average pooling的處理,然后將這2個結(jié)果基于channel做concat操作;經(jīng)過3個卷積核大小為3*3的卷積,降維為1個channel;再經(jīng)過sigmoid生成spatial attention feature;最后將該feature和該模塊的輸入feature做乘法,得到最終生成的特征。

圖4 改進空間注意力模塊結(jié)構(gòu)

整體表達公式如下:

Fs=Ms(Fc)?Fc

(3)

MS(F)=σ((f3×3([AugPool(F);

MaxPool(F)]))3)=

(4)

式中,Ms分別表示通道及空間注意力模塊,Fs表示通道注意力特征圖及空間注意力特征圖。通過改進空間注意力的卷積結(jié)構(gòu),使用連續(xù)3*3大小的卷積核在原算法的基礎(chǔ)上降低了模型的計算量。

1.4 損失函數(shù)EIoU(Efficient IoU Loss)

鋁型材瑕疵種類較多,像素數(shù)量不均衡,因而容易忽略個別瑕疵的重要性,原始的損失函數(shù)縱橫比差異不是寬高分別與其置信度的真實差異,因此在很大程度上會降低模型優(yōu)化相似性。EIoU Loss[19]在CIOU[10]的基礎(chǔ)上將縱橫比拆開計算,使目標與錨盒的寬度和高度之差最小,從而加快收斂效率。針對BBox回歸過程中樣本數(shù)據(jù)問題,利用梯度下降方法,EIoU Loss基于Focal Loss把高質(zhì)量的錨框和低質(zhì)量的錨框分開,通過分類加權(quán)的方法提高回歸精度,因此采用EIoU Loss作為整個網(wǎng)絡(luò)結(jié)構(gòu)的損失函數(shù)。該函數(shù)的懲罰公式如下:

LIOU=1-IOU

(5)

(6)

(7)

LEIOU=LIOU+Ldis+Lasp

(8)

LF-EIOU=IOUγLEIOU

(9)

式中,CW和Ch分別表示最小的寬度和高度,LIOU表示IOU損失函數(shù),Ldis表示距離損失函數(shù),Lasp表示寬高損失函數(shù),Υ為控制異常值抑制程度的參數(shù)。b和bgt表示預(yù)測框和真實框的中心點,ρ表示計算兩個中心點間的歐氏距離,c表示能夠同時包含預(yù)測框和真實框的最小閉包區(qū)域的對角線距離。

2 實驗結(jié)果與分析

2.1 實驗環(huán)境及參數(shù)設(shè)定

在實驗過程中選用Windows 10 x64的操作系統(tǒng),處理器為Intel(R) Xeon(R) Gold 6258R CPU @2.70 GHz,顯卡型號是NVIDIA Tesla V100S-PCIE-32 GB,采用Py-Torch1.12框架及CUDA11.3圖形加速,運行內(nèi)存為32 G,顯存是256 G,采用SGD優(yōu)化器,編譯語言為Python3.7。

2.2 數(shù)據(jù)集處理

數(shù)據(jù)集采用廣東某鋁型材廠家提供的競賽專用數(shù)據(jù)集,和實際生產(chǎn)中的情況直接相關(guān),具有公開性、實用性。采用該數(shù)據(jù)集可以使得訓(xùn)練后的模型更具實用性,實驗所用數(shù)據(jù)集中包含不導(dǎo)電、角位漏底、桔皮、漏底、雜色、起坑、漆泡共計7種常見表面瑕疵,如圖5所示。

圖5 數(shù)據(jù)集樣本類別

通過對數(shù)據(jù)集進行離線增強,如旋轉(zhuǎn)、亮度變化、裁剪等處理方式,將每種類型的瑕疵數(shù)據(jù)控制在700張,共計4 900張,并按照8∶2的策略劃分訓(xùn)練集和檢測集,由于原始數(shù)據(jù)集的標簽格式為json,因此需通過腳本文件對數(shù)據(jù)集標簽進行格式轉(zhuǎn)換,調(diào)整為YOLO使用的txt格式數(shù)據(jù)標簽。

2.3 評價指標

實驗采用計算機視覺領(lǐng)域通用檢測評價標準,AP(average precision)、mAP(mean average precision)以及FPS(Frames Per Second)作為模型性能評價指標,AP表示檢測精度,mAP表示平均檢測精度,FPS表示檢測速度,其計算公式如下:

(10)

(11)

(12)

(13)

式中,TP(True Positive)表示檢測正確數(shù)量,FN(False Negative)表示漏檢數(shù)量,FP(False Positive)表示錯檢數(shù)量,P(Precision)為查準率,R(Recall)為查全率,N為目標檢測任務(wù)的類別數(shù)量。AP的值越大,代表這個種類在該網(wǎng)絡(luò)模型上檢測效果越好,mAP值則反映該網(wǎng)絡(luò)模型對目標檢測的效果優(yōu)劣。

2.4 實驗對比

為驗證改進算法檢測性能的優(yōu)越性,將改進后的YOLOv5模型與YOLOv3、YOLO v4、SSD[20]等主流算法進行對比,并以Precision、Recall、mAP、FPS作為評估指標,實驗結(jié)果如表2所示。

表2 不同網(wǎng)絡(luò)結(jié)構(gòu)性能對比

通過表2可知,改進后的YOLOv5在mAP達到了85.8%,相較于YOLOv3精度提升10.64百分點,召回率提升8.8百分點,mAP提升7.45百分點,各項指標均有提升,相較于原始算法,精度提升2.2百分點,召回率提升3.9百分點,mAP提升2.9百分點,FPS提升4.7百分點,同其他算法相比,檢測性能及檢測速度均具有優(yōu)勢。

2.5 消融實驗

為進一步驗證改進策略的有效性,在原有YOLOv5的網(wǎng)絡(luò)結(jié)構(gòu)接觸上設(shè)計消融實驗,主要改進策略包括:anchor、檢測層、注意力機制E-CBAM、EIOU,在實驗過程中為確保實驗數(shù)據(jù)的可靠性,保證超參數(shù)、數(shù)據(jù)集等條件一致。消融實驗結(jié)果如表3所示。

表3 YOLOv5消融實驗

通過表3可知,在原始算法的基礎(chǔ)上,使用K-means++錨框生成方式后,算法的召回率提升了6.4百分點,增加檢測層后,將淺層特征進一步融合得到更加豐富,高效的語義信息,從而精確度提升了5百分點且召回率略有提升,在原有的YOLOv5基礎(chǔ)上加入改進后的注意力模塊E-CBAM發(fā)現(xiàn),其精確度、召回率皆有提升,其中mAP@0.5:0.95更是提升了9.8百分點,提升了邊框回歸效果,對于IOU效果要求較高的任務(wù)是一項有效的改進措施。

通過表2、表3及圖6數(shù)據(jù)展示可知,最終改進算法相比于原始YOLOv5算法,在各項精確度指標皆有提升,證明了算法改進的有效性,同時數(shù)據(jù)集數(shù)量滿足了工業(yè)生產(chǎn)要求,模型的性能滿足了企業(yè)需求。

圖6 不同檢測模型的mAP比較

3 結(jié)束語

針對鋁型材的不同瑕疵類別,提出了一種基于改進YOLOv5的鋁型材瑕疵檢測算法。通過融入改進后的E-CBAM注意力模塊,減少冗余計算,弱化背景因素影響,加強鋁型材瑕疵的特征提取能力;豐富檢測尺度的同時,加入新的卷積核,提升算法對不同尺度目標的檢測能力,加強對鋁型材表面瑕疵的區(qū)域特征學(xué)習(xí);采用K-means++錨框生成方式與EIOU損失函數(shù)相結(jié)合,使得算法回歸更加專注于高質(zhì)量錨框,提升了收斂效率,以及模型的精度和魯棒性。

目前,鋁型材相關(guān)廠家現(xiàn)場配備的攝像頭像素較低,無法對相關(guān)瑕疵進行精確捕捉,針對該問題,可通過后期人工標注,及對采集的數(shù)據(jù)進行高分辨率重建,提升數(shù)據(jù)精度。因此后續(xù)研究方向可從瑕疵的尺度、分布及算法的邊框回歸函數(shù)進一步考慮改進。