一種基于注意力機制的YOLO 缺陷檢測算法

于龍振 ,李逸飛 ,朱建華* ,趙謙 ,王志憲

(1.青島科技大學 經(jīng)濟與管理學院,山東 青島 266061;2.九州工業(yè)大學 計算機科學與工程研究生院,福岡 北九州8 040000)

缺陷檢測的深度學習算法主要可分為兩種類型:基于分類的算法和基于回歸的算法[1-3]?；诜诸惖乃惴ㄒ訰-CNN 系列為代表,包括R-CNN、SPP-Net、快速R-CNN、R-FCN、Mask R-CNN 等。FAN等[4]、JI等[5]、ZHANG等[6]分別對木材、齒輪、金屬等的表面缺陷進行了檢測。由于采用兩階段處理,即特征提取和對象分類,所以R-CNN 算法一般需要較高的算力,故該算法檢測速度相對較慢。而基于回歸的算法的特點是只做一輪處理,即同時完成特征提取和對象識別,因此速度比較快。REDMON等[7]提出了YOLO(you only look once)算法,這是一個基于回歸和端到端的算法。到目前為止,YOLO 已經(jīng)發(fā)展成YOLOv1 至v5 的系列算法,而其他代表性的算法還包括SSD、Corner Net等。在上述基于回歸的算法中,YOLOv3[8]是使用最廣泛的算法之一,基于YOLOv3,JING等[9]、LI等[10]和DU等[11]分別對織物、PCB 板、路面等進行了表面缺陷檢測。

本研究提出一種高效的缺陷檢測算法,該算法對原版YOLOv3提升了注意力。其主要思路在于工業(yè)相機的拍攝角度和缺陷特征相對固定,而注意力機制的提出是基于在神經(jīng)網(wǎng)絡的輸入和輸出之間存在全局依賴的關系[12],這啟示本研究將注意力放在相對固定區(qū)域缺陷檢測上。本研究對原版YOLOv3從三方面提升注意力:首先是采用CZS算法,重新組合圖像的缺陷區(qū)域并刪除無用的識別區(qū)域;其次是采用裁減主干網(wǎng)絡算法,裁減掉原版YOLOv3主干網(wǎng)絡中不必要的識別尺度網(wǎng)絡層;最后是采用數(shù)據(jù)增強算法,添加圖像噪音和使圖像旋轉,擴增訓練數(shù)據(jù)集。通過以上改進,使算法做圖像識別和訓練的計算量減少,使缺陷圖像的檢測效率得到提升,本研究并進行了實驗驗證。

1 算法

1.1 改進的YOLO 缺陷檢測算法原理

與經(jīng)典的深度學習物體檢測算法相比,本研究解決的一類缺陷檢測問題,其突出特點在于圖像上的識別區(qū)域是可預測的。如圖1示,圖1(a)為正常焊接效果,圖1(b)為缺陷焊接效果,兩副圖像拍攝角度是固定的。事實上,本研究只關心2張照片的紅框缺陷檢測相關區(qū)域,只需識別這個區(qū)域,就能確定該零件是否有缺陷,原始照片的其他區(qū)域可以被相應地刪除。

圖1 正常的焊接和帶缺陷的焊接Fig. 1 Normal welding and defective welding

此外本問題來源于工業(yè)現(xiàn)場,而解決方案也要再應用于工業(yè)場景,因此有精度和效率的要求,要盡量的優(yōu)化算法、提高精度和縮短檢測時間。據(jù)此本研究提出基于注意力機制的YOLOv3缺陷檢測算法整體框架如圖2示。

圖2 缺陷檢測算法整體框架Fig. 2 Overrall framework of the defect inspection algorithm

該缺陷檢測算法總體包括訓練算法和推理算法兩部分,此處只給出訓練算法,推理算法主要步驟相當于訓練算法的子集。本算法對YOLO 注意力的改進主要體現(xiàn)在CZS算法、裁減主干網(wǎng)絡算法和數(shù)據(jù)增強算法這3個子算法。在數(shù)據(jù)預處理階段調用CZS算法,能夠將原始大圖像轉變成只留待檢測區(qū)域的小圖像,把圖像上的無效檢測區(qū)域提前處理掉了,能夠提升注意力、節(jié)省算力、提升檢測效率;在圖像檢測階段,基于裁減主干網(wǎng)絡算法,能夠把通用型的原版YOLOv3網(wǎng)絡轉變成針對應用場景的精簡網(wǎng)絡,裁減掉不必要的檢測尺度,能夠有效減少網(wǎng)絡復雜度,使訓練的網(wǎng)絡更專注于針對性的檢測尺度,減少訓練時間,也提高檢測速度;在訓練算法中還使用了數(shù)據(jù)增強算法,這對樣本有限的網(wǎng)絡訓練格外重要,能夠大大增加樣本量和樣本多樣性,使訓練的網(wǎng)絡具有更強的甄別缺陷的能力。樣本量的大幅增加,也會使網(wǎng)絡訓練時間大幅增加,而結合前述兩個優(yōu)化子算法,使YOLOv3網(wǎng)絡訓練壓力在升降之間取得一定的折中。3個子算法的具體細節(jié)放在1.2到1.4節(jié)中依次介紹,而在后文第2節(jié)實驗分析中,對添加這3個子算法是否必要,分別添加前后的總算法性能做了詳細的數(shù)據(jù)對比分析,表明了改進注意力的有效性。

算法:對缺陷檢測網(wǎng)絡訓練

輸入:原始樣本集{(ai∈A,bi∈B)},1≤i≤I。其中ai表示第i張圖像,該圖像對應的缺陷描述文檔是bi,在bi上記錄著ai上的待檢測區(qū)域集Ki。A為圖像集,B為缺陷描述文檔集。

輸出:訓練好的YOLOv3模型。

1) 對任意ai∈A(1≤i≤I),反復執(zhí)行:

1.1)讀取bi,得ai的Ki。

1.2)對Ki執(zhí)行CZS算法,得新圖像a°i和對應缺陷描述文檔b°i。

1.3)分別對ai、bi和a°i、b°i執(zhí)行數(shù)據(jù)增強算法,合并得增強樣本集a°°i、b°°i。

2) 根據(jù)研究情境,使用裁減主干網(wǎng)絡算法,把YOLOv3的主干網(wǎng)絡Darknet-53 裁減為Darknet-39。

3) 使用增強樣本集{(a°°i∈A°°,b°°i∈B°°)},1≤i≤I,訓練2)后的YOLOv3。

4) 當達到精度要求或訓練迭代次數(shù)上限,則訓練結束。得訓練好的YOLOv3。步驟結束。

1.2 CZS算法

注意力機制主要是對系統(tǒng)的輸入和輸出之間的關系進行建模,不受距離的限制[12]。對于缺陷檢測,工業(yè)相機的拍攝角度是相對固定的。因此可以預先定義可能的缺陷區(qū)域,只從特定區(qū)域提取特征。注意力機制的代表形式之一是自我關注機制,原義是指通過關聯(lián)序列中的不同位置來計算序列語義表示。在YOLOv3的預處理中加入了相當于自我關注機制的圖像預處理操作,稱之為CZS 算法(cut,zoom and splice),如圖3。藍色方框代表切割區(qū)域,綠色方框和紅色方框代表兩種缺陷標記區(qū)域。左邊的原始圖像的1～8號區(qū)域對應右邊的圖像拼接區(qū)域。整個過程涉及3個步驟。從原圖像中剪切出若干需檢測區(qū)域,將各待檢測區(qū)域縮放到相同大小,然后再將這些區(qū)域拼接成新的圖像。

圖3 CZS算法Fig. 3 CZS algorithm

α代表一個擴展系數(shù),取值在1到2之間,意味著剪切區(qū)域的面積是缺陷標記區(qū)域的1到2倍。為了保證剪切區(qū)域能夠包含缺陷標記區(qū)域,剪切區(qū)域應該比缺陷標記區(qū)域略大。三元操作符 (邏輯表達式)“? 真操作:假操作”意味著當缺陷標記區(qū)域接近圖片的邊界時,剪切區(qū)域左上角相應延長。

步驟2:縮放算法(zoom)。是將圖像上的所有剪切區(qū)域縮放到相同的尺寸,以便能夠被寬度和高度均為416像素的新圖像完全容納,這是YOLOv3輸入圖像的標準尺寸。如果已知圖像上的缺陷標記區(qū)域的數(shù)量為Nl,那么新圖像的一行或一列中可容納的剪切區(qū)域的數(shù)量均為Nrow=Ncol=ceil(sq rt(Nl));剪切區(qū)域應被縮放的目標尺寸寬和高均為wz=hz=;而縮放系數(shù)為

步驟3:拼接算法(splice)。是將多個剪切區(qū)域經(jīng)過縮放操作后組合成一張新圖像。它主要涉及兩種算法,第一種算法是將剪切區(qū)域映射到拼接區(qū)域;第二種算法是將標記區(qū)域映射到拼接區(qū)域。

對于第一種算法,首先根據(jù)xc值將剪切區(qū)域由小到大排列。如果兩個剪切區(qū)域的xc值相同,那么按照yc值從小到大排列。假設一個剪切區(qū)域被排序為No,則計算中間變量nrow=(No%Ncol≠0)(No//Ncol+1):(No//Ncol)和ncol=(No%Ncol≠0)(No//Ncol):Ncol,那么在新圖像中該區(qū)域左上角的寬度和高度均為wd=hd=wz、x坐標為xd=(ncol-1)×wz,以及y坐標為yd=(nrow-1)×hz。以上公式中//代表整除操作,%代表整除取余操作,“條件? 為真執(zhí)行:為假執(zhí)行”是條件三元操作符。

以上完成了CZS算法,如果生成的新圖像還有空白區(qū)域,可把這些空白區(qū)域像素點填充為0處理。

1.3 裁減主干網(wǎng)絡算法

對應CZS操作后的新圖像,YOLOv3主干網(wǎng)絡可進行裁減,以更有針對性地檢測缺陷區(qū)域。圖4為原版及裁減部分主干網(wǎng)絡的YOLOv3算法,原版YOLOv3的主干網(wǎng)絡包括53層,稱為Darknet-53。YOLOv3的特征圖有3 種檢測尺度,尺度13×13像素對圖像上的大尺寸物體的識別有更好支持,尺度52×52像素對圖像上的小尺寸物體的識別有更好支持,而尺度26×26像素對圖像上的中等尺寸的物體識別有更好支持。

圖4 原版YOLOv3 Darknet-53和裁減部分主干網(wǎng)絡的YOLOv3 Darknet-39(實驗情景)Fig. 4 Original YOLOv3 Darknet-53 and YOLOv3 Darknet-39 by tailoring parts of the backbone network (experimental scenario)

根據(jù)缺陷檢測的案例背景,主干網(wǎng)絡可以相應的進行裁減定制。以汽車橡塑件的缺陷檢測為例,由于檢測照片中橡塑件的缺陷尺寸適中,因此可以簡化小尺度檢測網(wǎng)絡。再以PCB 焊點缺失檢測為例,由于焊點細密,且布局合理,所以可以對大尺度檢測網(wǎng)絡裁減。綜上,具體判斷算法如式(1)和(2)。

tailor函數(shù)表示裁減網(wǎng)絡的函數(shù)。52×52表示YOLOv3主干網(wǎng)絡中的小尺度檢測網(wǎng)絡層,13×13表示YOLOv3 主干網(wǎng)絡中的大尺度檢測網(wǎng)絡層,each函數(shù)表示取得每一個待檢測標記區(qū)域,N表示圖像上待檢測標記區(qū)域的總數(shù)目。wi和hi分別代表第i個待檢測標記區(qū)域的寬度和高度,而W和H分別代表圖像的寬度和高度。

在第3節(jié)實驗案例中,52×52小尺度網(wǎng)絡層被裁減了。具體如圖4所指,原版YOLOv3的第三步下采樣層要執(zhí)行8輪重復的卷積殘差操作,而在本模型中,刪除了其中7輪,即裁減掉了14個卷積層。因此主干網(wǎng)絡從53層縮減為39層,變成Darknet-39。

1.4 數(shù)據(jù)增強算法

為了增強注意力,提高深度學習網(wǎng)絡的識別精度,在原始訓練數(shù)據(jù)集規(guī)模不大的情況下還需要實施數(shù)據(jù)增強算法,以擴大訓練數(shù)據(jù)集,本研究主要基于以下2個步驟的算法。

步驟1:在原始圖像的缺陷標記區(qū)加入隨機噪聲,將正常區(qū)域變?yōu)閹г肼晠^(qū)域或故障缺陷區(qū)域。算法原理如公式(3),Ki,j代表第i張圖像的待檢測區(qū)域集中的第j個待添加遮擋區(qū)域(噪聲)代表添加的遮擋,rd是取的0～1之間的隨機數(shù)。則有遮擋1/3區(qū)域、遮擋當前整個區(qū)域、以及遮擋當前和下一個編號的2個區(qū)域,這3種添加遮擋的概率各占1/3。

如圖5,用矩形遮蓋模擬標記區(qū)域表面的噪聲或缺陷。其位置、大小和顏色可以隨機設置。圖5(a)表示正常;圖5(b)用矩形遮蓋了一個標記區(qū)1/3區(qū)域,相當于增加了一些噪聲,要保證網(wǎng)絡訓練能識別這樣的正常標記區(qū);圖5(c)和(d)用矩形完全遮蓋一個或兩個標記區(qū)來模擬缺陷?；谠摬襟E,可以從1張圖像擴增為4張圖像。

圖5 添加噪聲的演示效果Fig. 5 Demonstration of noisy or missing faults

步驟2:是將圖像旋轉90°、180°和270°,并水平翻轉。這相當于從1張圖像擴增為8張圖像。算法原理如公式(4)表示對第i張圖像進行旋轉和翻轉操作之后生成的圖像集,rotate90°,180°,270°表示依 次旋轉的角度,flip_horizontal表示水平翻轉操作。

采用自動化數(shù)據(jù)集增強操作,可以提高訓練質量,并有效減少人工作業(yè)。本研究案例初始只有630張樣本照片。而通過腳本自動化的CZS算法和數(shù)據(jù)增強算法,訓練和測試的最終數(shù)據(jù)集達到40 320張圖像。

2 實驗分析

2.1 實驗設計

案例研究從某汽車橡塑件生產企業(yè)獲得實驗數(shù)據(jù),該企業(yè)是國內多個知名汽車品牌的特級零部件供應商。因此該算法的成功實驗將對同行業(yè)具有示范意義。借助于工業(yè)相機的10個固定拍攝角度,對若干種產品收集了630張原始樣本照片,每個角度63張照片。每張照片都是1 600×1 200像素分辨率和8位深,共有16種類型的缺陷標記,正常標記和缺陷標記的比例約為9∶1。圖6為對原始圖像旋轉和水平翻轉的演示效果圖。開展研究的深度學習工作站硬件條件包括NVIDIA RTX3090 24 GB GPU;AMD R9 3900X CPU,64 GB DDR4 內存和4 TB 機械硬盤。軟件條件包括Ubuntu 18.04、CUDA 11、Python 3.8、Py Torch、Tensorflow 等。

圖6 對原始圖像旋轉和水平翻轉的演示效果Fig. 6 Demonstration of rotate and flip horizontal the original image

2.2 算法性能

在深度學習工作站上對算法訓練,設定訓練最大次數(shù)為300輪。經(jīng)過CZS算法和數(shù)據(jù)增強算法后,隨機取樣得訓練集32 000張照片,剩下8 320張照片為測試集。訓練過程比較漫長,一輪用時近17 min。但另一方面,精度收斂的非常快。當完成了300輪訓練后,準確率超過了99%。

有了預訓練的算法模型,隨即在工作站上測試時間效率。所花費的時間主要分為兩部分:一是Python庫的加載時間約為1 s,二是檢查時間約為0.01 s,如表1。Python庫的加載過程非常耗時,可以將其做成一個守護程序,使其始終處于加載庫狀態(tài),通過掃描檢測圖像的變化,實時待命檢測。本算法較原版YOLOv3的檢測時間縮短0.004 s,準確率提高4.2%?？梢姳舅惴▽崿F(xiàn)了對原版YOLOv3的性能提升,達到了更好的性能。

表1 本算法與原版YOLOv3的對比Table 1 Comparison between ouralgorithm and YOLOv3

2.3 綜合對比

在本研究中,m AP(平均精度的平均值)和Io U(交并比)被用作主要的精度評價指標。計算16個缺陷類別的平均檢測精度,然后將這16個平均精度的平均值作為m AP,m AP 越接近1越好。在訓練之前,由手動標記的待檢測區(qū)域即實際邊界框,之后由算法推理得到的標記區(qū)域即預測邊界框。Io U 是這兩個區(qū)域的交集除以并集。Io U 越接近1越好,表明算法預測區(qū)域與實際標記區(qū)域一致。

本算法的優(yōu)點是對YOLOv3提升注意力,對增加CZS算法和數(shù)據(jù)增強算法前后的算法進行了比較,效果見圖7。沒有CZS 算法的全局交并差(GIo U)出現(xiàn)過早收斂問題,無法找到最優(yōu)解,如第二行左一圖示。沒有進行數(shù)據(jù)增強算法的精度收斂較慢并震蕩,在訓練期內沒有達到理想的收斂效果,如第一行左四圖示。綜合對比,實施完整注意力改進的本算法,精度收斂更快、性能更穩(wěn)定。

圖7 對YOLOv3改進注意力的效果Fig. 7 Effect of attention on YOLOv3

本算法與以往一些精度較高的缺陷檢測研究的對比如表2示。XIA[13]和GE[14]的研究代表得采用傳統(tǒng)技術進行缺陷檢測,PRAPPACHER[15]、HE[16]則采用了深度學習技術,JING[9]也采用了YOLOv3。從精度來看,本算法唯一超過99%;從運行效率看(1/FPS),本算法位居前列;而由于準確性(m APbbox)的數(shù)據(jù)有限,只列出了本算法的運行結果為0.991,該值取1為最好。應該看到,由于進行缺陷檢測的具體問題域存在差別,所以并不能武斷的就斷定本算法更好。但總體肯定,本算法在固定攝像角度的工業(yè)檢測情境下,具有一定先進性。

表2 本算法與其他精度較高的缺陷檢測算法的性能對比Table 2 Performance comparison of this algorithm with other defect detection algorithms with higher accuracy

注意力機制對算法性能影響如表3示,該表是對圖7的進一步解釋。CZS操作和數(shù)據(jù)增強操作都是在圖像預處理步驟增加的環(huán)節(jié),裁減YOLOv3主干網(wǎng)絡則是對算法結構的改變。發(fā)現(xiàn)數(shù)據(jù)增強操作對算法精度影響最大,原始圖像只有630張,不足以對算法參數(shù)進行充分訓練,精度只有75.3%;而經(jīng)過數(shù)據(jù)增強后,精度達到95% 及以上。原版YOLOv3因為比本算法多了14個卷積層,需要更多的計算時間,所以檢測效率打了折扣,只有73 FPS的效率;由于其待訓練參數(shù)過多,其精度也有一定降低,只有95%。CZS操作使精度從98.3%提升到99.2%,看起來對算法性能影響并不是太顯著;但要看到,在算法精度從0到1的逐漸提升過程中,改進后面的1%和改進前面的1%并不是對等概念,甚至截止目前,沒有任何一種方法敢宣稱其缺陷檢測精度達到了100%,也就是說使精度從99%到100%是一個看來不可能完成的任務,反觀從98%提升到99%,其實是克服了極大困難才能達到的目標,所以CZS算法非常有意義。

表3 注意力機制的性能影響Table 3 Performance impact of attention mechanism

2.4 進一步探討

為了進一步驗證算法的有效性,本研究還進行了其他一些相關實驗并獲得一些結論。

遷移學習。在訓練本算法時,權重是全新訓練的。已有研究認為,原版YOLOv3也可以通過遷移學習技術快速訓練成型。因此考慮兩種不同的訓練方案。方案一是全新訓練網(wǎng)絡。方案二是基于微軟COCO 數(shù)據(jù)集預訓練的原版網(wǎng)絡通過遷移學習快速訓練。結果顯示這兩種方法的訓練時間非常相似,都是48 h左右,而m AP 都是0.991左右,也就是說遷移學習并沒有縮短訓練時間。由于本研究的數(shù)據(jù)集與COCO 數(shù)據(jù)集之間存在較大差異,因此跨數(shù)據(jù)集遷移學習的效果并不明顯。

精度分析。與原版YOLOv3相比,本算法減少了14個神經(jīng)網(wǎng)絡層。然而經(jīng)過300輪的訓練,精確率反而從95%提高到99.2%。這是因為優(yōu)化版的待訓練參數(shù)變少了,比原版更容易收斂。

誤差分析。雖然該方法的識別準確率達到99.2%,但前提是現(xiàn)有只涉及10 個拍攝角度的圖像。而當增加新角度圖像時,即便還是同樣批次零件,雖然要識別的缺陷類型沒增加,但精確率達不到99.2%,說明該算法的自適應可擴展性能有限,還有必要通過其他技術手段再優(yōu)化。

3 結語

本研究的問題來自于常見的一種工廠情境,即工業(yè)相機拍攝角度相對固定,這為設計一種基于注意力機制的高效缺陷檢測算法提供了靈感。首先采用CZS算法,能去掉圖像上無效檢測區(qū)域,最大限度保留圖像上缺陷檢測相關區(qū)域。其次采用裁減主干網(wǎng)絡算法,根據(jù)待檢目標的相對尺寸,對YOLOv3的3種尺度的識別網(wǎng)絡適當裁減。最后采用數(shù)據(jù)增強算法,通過添加圖像噪音和修改圖像角度兩個策略而擴增圖像集。實驗結果表明,通過提升YOLOv3的注意力,算法精度達到99.2%,檢測時間為0.01 s,可以更好的滿足工業(yè)領域的檢測需求。另一方面,本工作未解決問題:一是關于算法選擇,目前已經(jīng)發(fā)布了YOLOv5的第6個版本,亟需對新老YOLO 技術進行驗證比較。二是關于使用情境,本研究訓練和檢測環(huán)境需要深度學習工作站,而從發(fā)展趨勢看,分布式、自適應的邊緣計算和霧計算等將是未來主流,需要將智能算法下沉部署到邊緣芯片層。