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

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

(1.青島科技大學 經濟與管理學院,山東 青島 266061;2.九州工業大學 計算機科學與工程研究生院,福岡 北九州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 已經發展成YOLOv1 至v5 的系列算法,而其他代表性的算法還包括SSD、Corner Net等。在上述基于回歸的算法中,YOLOv3[8]是使用最廣泛的算法之一,基于YOLOv3,JING等[9]、LI等[10]和DU等[11]分別對織物、PCB 板、路面等進行了表面缺陷檢測。

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

1 算法

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

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

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

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

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

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

算法:對缺陷檢測網絡訓練

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

輸出:訓練好的YOLOv3模型。

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

1.1)讀取bi,得ai的Ki。

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

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

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

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

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

1.2 CZS算法

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

圖3 CZS算法Fig. 3 CZS algorithm

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

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

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

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

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

1.3 裁減主干網絡算法

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

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

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

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

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

1.4 數據增強算法

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

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

如圖5,用矩形遮蓋模擬標記區域表面的噪聲或缺陷。其位置、大小和顏色可以隨機設置。圖5(a)表示正常;圖5(b)用矩形遮蓋了一個標記區1/3區域,相當于增加了一些噪聲,要保證網絡訓練能識別這樣的正常標記區;圖5(c)和(d)用矩形完全遮蓋一個或兩個標記區來模擬缺陷?；谠摬襟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表示水平翻轉操作。

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

2 實驗分析

2.1 實驗設計

案例研究從某汽車橡塑件生產企業獲得實驗數據,該企業是國內多個知名汽車品牌的特級零部件供應商。因此該算法的成功實驗將對同行業具有示范意義。借助于工業相機的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 算法性能

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

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

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

2.3 綜合對比

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

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

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

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

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

注意力機制對算法性能影響如表3示,該表是對圖7的進一步解釋。CZS操作和數據增強操作都是在圖像預處理步驟增加的環節,裁減YOLOv3主干網絡則是對算法結構的改變。發現數據增強操作對算法精度影響最大,原始圖像只有630張,不足以對算法參數進行充分訓練,精度只有75.3%;而經過數據增強后,精度達到95% 及以上。原版YOLOv3因為比本算法多了14個卷積層,需要更多的計算時間,所以檢測效率打了折扣,只有73 FPS的效率;由于其待訓練參數過多,其精度也有一定降低,只有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也可以通過遷移學習技術快速訓練成型。因此考慮兩種不同的訓練方案。方案一是全新訓練網絡。方案二是基于微軟COCO 數據集預訓練的原版網絡通過遷移學習快速訓練。結果顯示這兩種方法的訓練時間非常相似,都是48 h左右,而m AP 都是0.991左右,也就是說遷移學習并沒有縮短訓練時間。由于本研究的數據集與COCO 數據集之間存在較大差異,因此跨數據集遷移學習的效果并不明顯。

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

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

3 結語

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