基于改進YOLOv4算法的輪轂表面缺陷檢測

吳鳳和, 崔健新, 張 寧, 張志良, 張會龍, 郭保蘇

(1.燕山大學機械工程學院,河北秦皇島066004；2.中信戴卡股份有限公司, 河北秦皇島066004)

1 引 言

輪轂表面缺陷會影響產(chǎn)品的外觀和品牌形象,甚至導致嚴重的安全問題。目前,大部分輪轂生產(chǎn)線都采用人工方式進行輪轂缺陷檢測,操作過程繁瑣,勞動強度大,且效率低下、漏檢率較高,難以滿足輪轂缺陷檢測的精度和效率要求。機器視覺技術具有較高的檢測效率和精度,且對特殊工業(yè)環(huán)境的適應性較強,已在環(huán)境目標檢測、表面缺陷檢測、織物紋理檢測等方面得到了廣泛應用[1-3]。

根據(jù)大致原理可以將各種目標檢測算法分成two-stage和one-stage兩類[4]。two-stage算法需要首先生成可能包含缺陷的候選框,再進行目標檢測,Girshick R等[5]于2014年提出了RCNN算法,隨著深度學習技術的發(fā)展,RCNN算法有了眾多改進,檢測精度及效率均有所提升,如文獻[6]提出的Fast RCNN、文獻[7]提出的Faster RCNN算法等。2017年,Cha Y J等[8]最早將Faster RCNN應用于橋梁表面缺陷定位,其主干網(wǎng)絡被替換為ZF-net,在包含2 366張500×375像素大小圖像的5類橋梁建筑數(shù)據(jù)集中,檢測平均精度Pma值達到87.8%。雖然RCNN系列有多種算法,但由于訓練中需要將多個候選框送入CNN提取特征,存在效率低、提取特征操作冗余等問題,無法滿足大多數(shù)工業(yè)場景的檢測效率要求。

One-stage算法的代表有SSD和Yolo系列算法。文獻[9]最早提出的SSD算法相比于Faster RCNN,提供了一個徹底的端到端的神經(jīng)網(wǎng)絡,通過添加輔助性的特征金字塔,并融入先驗框機制,解決了輸入圖像目標尺寸不同的問題,同時提高了檢測精度。Chen J等[10]采用不同層級的特征圖來改進SSD網(wǎng)絡實現(xiàn)了某種緊固件的缺陷檢測。但因SSD算法先驗框的大小和形狀不能直接通過學習獲得,導致調(diào)試過程過于依賴經(jīng)驗,且針對小目標的檢測效果不太理想。Yolo算法也是一種端到端的神經(jīng)網(wǎng)絡,由Redmon J等[11]在2016年首次提出,在此基礎上結合新的主干網(wǎng)絡和損失函數(shù)又相繼提出了YOLOv2[12]和YOLOv3[13]算法。Yolo算法直接將目標邊框定位問題轉化為回歸問題,并基于圖像的全局信息進行預測,但對近距離物體和較小物體的檢測效果不好,泛化能力相對較弱。譚芳等[14]采用深度學習YOLOv3算法檢測行人目標,實現(xiàn)多目標行人跟蹤的算法,在復雜環(huán)境下均獲得了良好的連續(xù)跟蹤效果。

2020年Bochkovskiy A等[15]發(fā)布了YOLOv4。YOLOv4在原理上與前代的YOLOv3相似,但是在主干網(wǎng)絡、網(wǎng)絡頸部、以及解碼預測部分都結合了全新的模塊方法,在保證檢測效率的同時,大大提高了檢測精度。Xin H等[16]使用YOLOv4算法對PCB板的缺陷進行檢測,在數(shù)據(jù)預處理階段根據(jù)缺陷的尺寸分布和標注框的尺寸大小自動細分圖像,并對算法的超參數(shù)進行優(yōu)化,使得檢測精度有所提高。

雖然目標檢測算法種類較多,但本文輪轂缺陷檢測針對的是工業(yè)生產(chǎn)線,對缺陷檢測的精度和效率都有較高要求,且輪轂缺陷的尺度范圍較大,從十幾cm,到幾mm,缺陷的分布位置也存在隨機性,故現(xiàn)有方法難以實現(xiàn)輪轂各類缺陷的高精度有效檢測。因此,本文基于輪轂缺陷數(shù)據(jù)集的特點,以目前檢測效果最為均衡的YOLOv4算法為基礎模型,針對Yolo系列算法對小目標檢測精度低、泛化能力弱的問題,融入細化U型網(wǎng)絡模塊(thinned U-shaped network module, TUM),強化多尺度特征提取與融合能力,使所得每個規(guī)模的特征都包含多個層級的信息；同時引入注意力機制,使網(wǎng)絡對特征信息的處理更具針對性,從而提高網(wǎng)絡檢測精度,并增強網(wǎng)絡泛化能力,在保證效率的前提下,實現(xiàn)缺陷的精確檢測。

2 輪轂數(shù)據(jù)集制作

圖1 輪轂表面缺陷種類Fig.1 Types of wheel hub surface defects

本文針對輪轂生產(chǎn)線中出現(xiàn)頻率較高且對輪轂性能影響較為嚴重的6種缺陷(漏切、纖維、粘鋁、臟污、毛刺、氣孔)進行數(shù)據(jù)采集,自制輪轂缺陷數(shù)據(jù)集。采集到的輪轂缺陷圖片數(shù)量依據(jù)實際工況中各種缺陷出現(xiàn)的比例,圖像大小為4 928l×3 264 pixel,缺陷類型見圖1。其中,粘鋁缺陷是輪轂窗口等邊角處存在粘附鋁屑現(xiàn)象；漏切體現(xiàn)在輪轂的某一個面上存在黑皮或者印痕,主要由毛坯變形,加工量不準確等情況導致；纖維缺陷由于環(huán)境污染或者油漆質(zhì)量差等原因導致輪轂正面有纖維或者毛發(fā)摻雜；毛刺清理未達標,導致毛刺過大形成毛刺缺陷；壓力不穩(wěn)、充型卷氣、排氣不暢等原因容易產(chǎn)生氣孔,氣孔缺陷多發(fā)生在輪轂的輪輞或輪芯處；噴涂環(huán)境或送風設備清理不及時使得輪轂表面在噴涂過程粘有灰塵,從而形成臟污缺陷。

由于在試驗階段采集到的缺陷圖像數(shù)據(jù)量有限,為了進一步增加模型的學習范圍,采用數(shù)據(jù)增廣策略進行數(shù)據(jù)集擴充。數(shù)據(jù)增廣可以提供更具多樣性的訓練數(shù)據(jù),提高CNN性能,降低過擬合概率。本文運用的數(shù)據(jù)增廣手段包括圖像旋轉、圖像鏡像、圖像亮度對比度改變、圖像添加噪聲等。經(jīng)過一系列增廣最終得到2 346幅圖像。為了得到輪轂圖像數(shù)據(jù)標簽(缺陷類型及坐標),運用LabelImg工具,采用手動方式對圖像進行標注。最終,從2 346幅圖像中共標記得到3 554個缺陷標簽。

3 改進的YOLOv4算法

3.1 YOLOv4算法簡介

YOLOv4是2020年最新提出的目標檢測算法,它是在YOLOv3的基礎上結合多種模塊進行創(chuàng)新[15]。如圖2所示,YOLOv4主要分為4個部分,分別是①輸入端、②主干網(wǎng)絡(Backbone)、③網(wǎng)絡頸部(Neck)以及④輸出預測(Prediction)模塊。輸入端會將數(shù)據(jù)集統(tǒng)一處理成相同尺寸的圖片,繼而輸入Backbone。Backbone主要用于特征提取,該部分的網(wǎng)絡被作者命名為DSPdarknet53,它將多個卷積核大小為3×3、步長為2的卷積結合殘差方式(Res unit)共同組成CSP結構,并將其作為基本組件構成了整個主干網(wǎng)絡,加入殘差結構可以有效的抑制梯度消失的問題,保證在加深網(wǎng)絡層數(shù)的同時不會導致網(wǎng)絡退化。Neck部分運用了空間金字塔池化(spatial pyramid pooling, SPP)結構和路徑聚合網(wǎng)絡(PANet)結構,SPP結構是對Backbone輸出的特征層經(jīng)過3次卷積后的結果進行最大池化,池化過程中共使用4種不同尺度的池化層進行處理,池化核大小分別為 1×1、 5×5、 9×9、 13×13,經(jīng)過SPP處理后可有效增加感受野,分離出顯著的上下文特征。PANet結構是由卷積操作、上采樣操作、特征層融合、下采樣操作構成的循環(huán)金字塔結構。輸出預測的Prediction即輸出端,對Neck部分處理后的3個特征層進行結果預測,網(wǎng)絡的輸出大小與輸入端圖片大小相關,網(wǎng)絡的輸出維度與需要識別的目標種類數(shù)量相關。最后依據(jù)CIOU損失函數(shù)(complete intersection over union loss)對每個特征層的3個先驗框進行判別,判斷其內(nèi)部是否包含目標并識別目標種類,同時進行非極大抑制處理和預測框調(diào)整以實現(xiàn)缺陷的分類及定位。

圖2 YOLOv4網(wǎng)絡結構圖Fig.2 YOLOv4 network structure diagram

3.2 先驗框標定

Yolo系列的目標檢測算法需要在訓練時預先設定先驗框。原始YOLOv4算法提供的先驗框是由COCO數(shù)據(jù)集聚類得到的,而COCO數(shù)據(jù)集圖片與輪轂表面缺陷數(shù)據(jù)集圖片中的目標大小和目標類型存在較大差異,因此采用K-means聚類算法重新生成先驗框。在YOLOv4算法中,檢測分為3個特征層,每個特征層中含有3個大小不同的先驗框。先驗框通過算法對數(shù)據(jù)集中標注好的檢測目標聚類得到。在聚類過程中,首先隨機化初始位置并選定9個聚類中心,計算檢測目標中每個標注框與聚類中心點的交并比,并將標注框分配給交并比最大的聚類中心,分配結束后重新計算聚類中心,直到聚類中心不再發(fā)生改變,并獲得最終聚類效果最好的9個先驗框的寬、高。具體的目標函數(shù)D計算公式如下：

(1)

式中：Sboxi為檢測目標中第i個標注框的區(qū)域面積；Scenj為第j個聚類中心的區(qū)域面積；Sboxi∩Scenj為標注框與聚類中心區(qū)域交集的面積；Sboxi∪Scenj為標注框與聚類中心區(qū)域并集的面積；n為檢測目標數(shù)量；k為聚類中心個數(shù)。

3.3 Neck部分改進

輪轂表面缺陷尺度差異較大,如漏切缺陷通常占據(jù)輪轂表面的較大面積,而纖維缺陷往往只是一小片區(qū)域。one-stage算法受限于端到端的組合方式,雖然檢測效率大幅提升,但對于小缺陷的檢測效果往往較差。一般來說,尺度越大的特征層包含的信息越多,特征層經(jīng)由下采樣之后由于其尺度變小,缺陷信息會大大減少,幾乎只會包含大缺陷的信息。針對該問題,對YOLOv4算法進行了改進,由于主干網(wǎng)絡的作用主要在于對特征的提取,而Neck部分的作用在于強化提取與特征融合。為了進一步挖掘由網(wǎng)絡前端提取到的特征,更大程度地融合更多不同尺度的特征,使得小缺陷的特征得以保留,本文剔除原YOLOv4 Backbone中38×38 pixel的特征輸出,將Backbone中76×76 pixel的特征層經(jīng)過CBL直接與SPP模塊上采樣的結果進行級聯(lián)(Concat)操作；Neck部分只保留SPP模塊,在原本的SPP模塊后面添加多個細化U型網(wǎng)絡(TUM)與注意力機制；經(jīng)由改進Neck部分的輸出特征層直接與Prediction部分連接,如圖3所示。

3.3.1 TUM模塊

本文設計的TUM模塊位于SPP模塊之后,分為下采樣和融合上采樣兩部分。在下采樣部分,進行1×1卷積和多個步長為2的3×3卷積；融合上采樣部分對下采樣的輸出層進行上采樣和特征元素融合操作,然后在最終輸出層前添加1×1卷積層,以增強學習能力,保持特征的平滑。每個TUM中的輸出特征層共有3種尺寸(76×76 pixel,38×38 pixel,19×19 pixel),從整體上看,分別代表淺層、中層以及深層的特征。這種多尺度的特征輸出,有利于后續(xù)完成最大程度的特征提取與融合。其結構如圖4所示。

3.3.2 注意力機制模塊

在嵌入TUM模塊后,需要將多個TUM中相同尺寸的輸出特征層分別進行Concat操作,但直接的特征堆疊可能導致信息冗余,因此本文引入SE-block方式的注意力機制(圖5),注意力機制可以自適應優(yōu)化特征權重,增強有效特征,抑制無效特征,得到表征能力更強的特征,使得特征處理過程中小目標信息丟失的問題得以改善。

具體過程如下：①對TUM模塊Concat得到的特征層進行特征壓縮,將每個特征通道變成一個實數(shù),這個實數(shù)在某種程度上具有全局感受野,并且輸出的維度和輸入的特征通道數(shù)相匹配；②壓縮結果經(jīng)由卷積網(wǎng)絡與激活函數(shù)激活,對每個特征通道生成一個權重,輸出權重描述每個特征通道的重要性；③將權重加權到由TUM模塊堆疊的特征上,完成在通道維度上的對原始特征的重標定。經(jīng)由以上處理得到的3個特征層便可直接輸入Prediction模塊。

圖4 TUM模塊結構圖Fig.4 TUM mod ule structure diagram

圖5 注意力機制模塊Fig.5 Attention mechanism module

4 實驗與分析

為驗證算法的有效性,對提出的改進算法進行消融實驗并與其他常用算法進行對比。所有算法均基于Python語言開發(fā),使用模塊包括numpy、OpenCV、tensorflow等,運行在Windows10平臺下。實驗所用的硬件配置為Intel Core i7 11700K CPU,16G RAM以及NVIDIA GeForce RTX3070 GPU。

4.1 評價標準

在缺陷檢測中,通常將預選框與目標真實框的交并比大于0.5設定為成功預測到目標位置,通過準確率(P)和召回率(R)計算平均精確度(Pa)和加權調(diào)和平均值(Fβ),并以Pa的均值Pma作為算法精度的綜合評價指標。P代表預測為正的樣本中有多少是真正的正樣本,與誤檢率相關；而R代表數(shù)據(jù)集的正樣本中有多少被正確檢測,與漏檢率相關。Fβ是基于P和R的加權調(diào)和平均,將誤檢與漏檢加權計算,反映了算法的綜合性能,其值越大表示算法性能越佳。Pa用于評估模型在單個檢測類別上的精度表現(xiàn),由P-R曲線圍成的面積進行表示,而Pma代表的是整體精度,Pma越大表明整體檢測精度越高。

上述各指標的計算公式如下：

(2)

(3)

(4)

(5)

(6)

式中：PT表示檢測正確的目標數(shù)量；PF表示檢測錯誤的目標數(shù)量；NF表示漏檢的目標數(shù)量；β為常數(shù),∑Pa表示所有缺陷類別的Pa值總和；Nc表示缺陷總類別數(shù)。

此外,檢測效率也是判定算法性能的重要指標之一,本文采用單張圖片平均檢測時間T作為檢測效率的評價指標。T越小,表示檢測效率越高,更容易達到實時檢測。其計算公式為

(7)

式中：Ttotal為檢測的總時間；Fn為檢測的圖片數(shù)量。

4.2 實驗及對比結果

本文在訓練時把數(shù)據(jù)集分成訓練集和測試集,其中隨機選取20%的各種缺陷數(shù)據(jù)圖作為測試集,其余圖像作為訓練集。訓練、測試集的比例以及數(shù)量如表1所示。

表1訓練集與測試集圖片數(shù)量對比Tab.1 The proportion of training and validation sets.

4.2.1 TUM模塊數(shù)量對比實驗

網(wǎng)絡訓練的學習率依據(jù)余弦退火的形式進行下降,余弦函數(shù)中隨著自變量的增加余弦值首先緩慢下降,然后加速下降,再次緩慢下降。這種下降模式能與學習率匹配,提升訓練效果。網(wǎng)絡結構設計上采用不同的TUM模塊數(shù)量,為了保證訓練效率以及訓練結果的可信度,選用TUM數(shù)量分別為3、4、5、6、7、8、9這7種進行多次實驗,每次實驗設定的迭代次數(shù)為100,置信度設置為0.5。采用3.3節(jié)中描述的改進YOLOv4算法對自制數(shù)據(jù)集進行訓練,并通過測試集對其檢測精度進行評估。每個案例的訓練時間大約為5 h,每幅圖像的測試時間大約為0.04 s。通過Pma值來衡量不同TUM數(shù)量下的模型性能,如圖6所示。

圖6 不同TUM數(shù)量的精度對比Fig.6 Accuracy comparison based on the number of TUM

從圖6可以看出,對于數(shù)據(jù)集中6種缺陷的檢測精度可以大致分為3個層次。漏切、臟污、毛刺3種缺陷由于其缺陷面積較大,特征明顯,檢測精度普遍較高；氣孔、纖維缺陷是小缺陷,包含的特征信息較少,故檢測精度較低；而粘鋁缺陷雖然同樣屬于小缺陷,但由于其多發(fā)生在輪轂邊角處,缺陷位置較為固定,缺陷特征更易提取,所以檢測精度相對其他小缺陷有所提升。不同TUM數(shù)量下的檢測精度有較大波動,一般來說,隨著神經(jīng)網(wǎng)絡層數(shù)的加深,其提取的特征信息會愈加復雜抽象,會從簡單的形狀向更“高級”的信息變化,這也是加深神經(jīng)網(wǎng)絡層數(shù)的意義所在。但由于數(shù)據(jù)集規(guī)模的限制,神經(jīng)網(wǎng)絡學習到的信息有限,且神經(jīng)網(wǎng)絡每一次輸入到輸出的迭代過程可能會造成一些不可逆的信息丟失,在網(wǎng)絡加深后可能導致過擬合問題。因此,本文以實驗效果最佳作為確定TUM數(shù)量的依據(jù)。通過分析實驗結果可知,TUM模塊數(shù)量為8時檢測精度最高,模型的檢測平均精度可達85.8%,因此將TUM模塊數(shù)量設置為8。

4.2.2 消融實驗

為了驗證本文對原算法改進的有效性,針對K-means聚類、TUM模塊及注意力機制模塊設置消融實驗,實驗將原算法與采用不同配置的改進后算法進行對比,對比結果見表2。

表2改進YOLOv4算法的消融實驗Tab.2 Ablation experiment for improved YOLOv4 algorithm

如表2所示,運用K-means聚類得到的先驗框進行訓練后,相比原始YOLOv4算法平均精度(82.5%)有所提升,表明針對數(shù)據(jù)集特點重新標定先驗框具有一定價值；在此基礎上嵌入TUM模塊,網(wǎng)絡平均檢測精度由83.33%提升到84.7%,這些提升主要反應在小尺度缺陷上；分析了TUM模塊直接堆疊特征層可能存在問題而引入擠壓激勵模塊(squeeze and excitation block, SE block)機制后,算法的檢測精度進一步提升,達到了85.8%,從而證明了本文改進有效性。

4.2.3 算法對比實驗

為了評估本文改進算法的性能,在相同的數(shù)據(jù)集下,訓練并對比本文改進算法與SSD、Faster RCNN、YOLOv3以及原始YOLOv4算法的評價指標。由于車輪輪轂表面缺陷對于整車外觀甚至安全性有較大影響,在進行缺陷檢測時更希望盡可能避免漏檢,因此,除了平均檢測精度和檢測效率外,Fβ也是值得參考的重要指標(Fβ中β值的設定不宜過低,本文設定該值為1.2,以強調(diào)召回率的重要程度)。上述各種網(wǎng)絡在目標檢測領域應用廣泛,認可度較高。最終對比結果如表3～表5所示。

表3不同算法P,R,Fβ對比Tab.3 Comparison of different algorithms P,R,Fβ (%)

表4不同算法精度對比Tab.4 Comparison of accuracy of different algorithms (%)

表5不同算法檢測效率對比Tab.5 Comparison of detection speed of different algorithms ms

表3中記錄了不同算法對應不同缺陷種類的P、R、Fβ值。如表3所示,對于Fβ值,無論針對哪種缺陷本文算法都是最高的,表明本文算法可以最大程度上避免誤檢、漏檢的發(fā)生。

從表4、表5所反映的信息可以看出,Faster RCNN算法精度為75%,但其檢測效率較低,單幅圖像的檢測時間需要127 ms；而SSD算法正好相反,雖然其檢測效率占優(yōu),但是精度較低,僅僅達到55%,尤其是針對粘鋁、纖維等小缺陷的檢測效果很不理想；YOLOv3算法相較前兩種較為均衡,但精度上與YOLOv4還有一定差距。

對于缺陷檢測,精度是相對重要的評價指標,本文算法在檢測效率上與原算法差距不大,但檢測精度提升了3.3%,多種缺陷在改進算法下的檢測精度均有提升,尤其小缺陷檢測精度的提升幅度更為明顯,如粘鋁缺陷的檢測精度由原來的78%提升到了85%,證明了本文算法的優(yōu)越性。

上述對比數(shù)據(jù)表明,本文提出的改進算法在輪轂表面多類缺陷檢測上,具有最優(yōu)的檢測精度和較高的檢測效率。實際檢測效果示例如圖7所示。

圖7 算法檢測結果Fig.7 Algorithm detection result

5 結 論

本文提出一種基于改進YOLOv4算法的輪轂表面缺陷檢測方法。針對YOLOv4算法對小尺寸缺陷檢測精度不足問題,通過融入TUM模塊和注意力機制,提高算法的多尺度特征提取及融合能力,并依據(jù)缺陷種類,自適應分配缺陷特征權重,使得模型對小尺寸缺陷的表達能力得到提升；使用自主采集并制作的輪轂表面缺陷數(shù)據(jù)集進行模型性能驗證,結果表明：改進YOLOv4算法的平均檢測精度和單幅圖像檢測時間分別為85.8%和38 ms,在保證檢測效率的前提下,有效提高了模型對小尺寸缺陷的檢測效果。