一種基于pix2pix改進的工業缺陷數據增強方法*

羅月童,段 昶,江佩峰,周 波

(合肥工業大學計算機與信息學院，安徽 合肥 230601)

1 引言

基于深度學習的目標檢測算法能夠學習到目標在各種復雜多變情形下的特征，具有強大的健壯性和泛化能力[1]，特別是基于有監督的深度學習目標檢測算法，在目標檢測任務中表現出了優異的性能。但是，有監督的深度學習目標檢測算法運行的前提是擁有大量有標記的缺陷數據，在工業檢測中缺陷數據的獲取成本也往往較高，數據增強則為降低工業成本、提高網絡的泛化能力提供了一種行之有效的手段。常規數據增強方法包含幾何變換類方法和顏色變換類方法。幾何變換類方法通過對圖像進行翻轉、裁剪和縮放等操作[2]來實現數據擴充;顏色變換類方法則是通過噪聲注入[3]、光度變換[4]和隨機擦除[5]等方法來實現數據擴充。以上數據增強方法操作簡單，能夠迅速擴大數據集，但是傳統數據增強方法也存在嚴重的弊端，通過幾何變換生成的圖像不能有效擴增樣本特征的多樣性，通過顏色變換生成的樣本的多樣性依賴人工設定[6]，因此傳統數據增強方法只能在一定程度上緩解神經網絡過擬合的問題，難以從根本上解決數據不足的問題。

生成對抗網絡GAN(Generative Adversarial Network)[7]經過訓練能夠學習到已有數據的分布，并生成與當前數據分布類似的數據，相較于傳統數據增強方法，其通過神經網絡自動學習圖像特征，無需人工干預，且基于GAN生成的數據幾乎與真實數據的分布相同，從根本上解決了數據不足的問題。因此，GAN逐漸成為一種新的數據增強方法，通過GAN生成的數據能夠有效提升有監督深度學習網絡的泛化性[8]。近年來研究人員主要從GAN的生成器結構、判別器結構和損失函數3個方面進行了改進。林志鵬等人[9]基于深度卷積生成對抗網絡DCGAN(Deep Convolutional GAN)[10]進行改進，重新設計了DCGAN的生成器和判別器的網絡結構，生成了分辨率更高的宮頸細胞圖像用于數據增強，顯著提升了宮頸細胞圖像分類的正確率。Ouyang等人[11]基于pix2pix[12]的判別器結構和損失函數進行改進，提出了PSGAN(Pedestrian-Synthesis-GAN)，PSGAN能夠在街道圖像中添加行人用于數據擴充，從而提升行人檢測的準確率。王金娜等人[6]通過改進Cycle-GAN來提升絕緣子圖像的生成效果，增加了絕緣子圖像分類的正確率。孫曉等人[13]通過改進StarGAN來生成人臉表情圖像，有效提升了人臉表情的識別精度。劉坤等人[14]提出了NSGGAN(Negative Sample Guided Generative Adversarial Network)來進行太陽能電池的數據增強。陳文兵等人[15]提出了數據增強方法GMM-CGAN(Gaussian Mixture Model and Conditional Generative Adversarial Network)。陳佛計等人[16]提出了一種基于GAN的紅外圖像數據增強方法。

盡管當前一些基于GAN的數據增強方法取得了較好的效果，但是這些方法受應用場景的限制，無法直接應用于工業場景。工業場景中通常存在正常數據易收集、缺陷數據難獲取的問題，導致收集到的良品非常多、缺陷數據卻非常有限。同時，在一個確定的工業場景中，良品的特征相對一致，缺陷數據則是通過改變良品上局部位置的特征得到，兩者可以看作是2種不同風格的配對數據。而pix2pix是一個通用的圖像風格轉換框架，通過pix2pix能夠利用大量的正常數據生成缺陷數據，從而解決缺陷數據不足的問題。因此，本文基于pix2pix進行研究。由于pix2pix采用了patchGAN作為判別器，其更重視圖像整體的轉換質量，而在工業場景中，缺陷往往只占據整幅圖像的小部分，即使生成的缺陷質量不高，從整體上看，生成的圖像也可能被判別為真，在小缺陷居多的工業場景中，該問題表現得更加突出。為了解決該問題，本文從增加生成器和判別器對缺陷區域的注意力，進而提升生成缺陷質量的角度出發，在pix2pix的基礎上，改進了其判別器和損失函數，提出了DS-GAN(Defects-Synthesis GAN)的工業數據增強方法。

2 本文方法

本文方法是基于pix2pix改進而來的，以下將從pix2pix基本原理、DS-GAN設計思路、DS-GAN算法流程、DS-GAN生成器結構、DS-GAN判別器結構和DS-GAN損失函數6個方面進行介紹。

2.1 pix2pix

pix2pix是Isola等人[12]在2017年提出的一個基于條件生成對抗網絡cGAN(conditional Generative Adversarial Nets)[17]的圖像風格轉換框架，由1個生成器和1個判別器組成。pix2pix采用的生成器結構為U-NET，U-NET通過在特征層之間進行跳躍連接的方式保留了輸入圖像的低維信息；pix2pix的判別器采用了patchGAN判別器，其將輸入圖像切分成70*70大小的不同圖像塊(patch)，并對每一個圖像塊進行判別，將所有圖像塊的判別結果進行求和平均，作為輸入圖像的判別結果，通過該方式避免了判別器因為某個區域生成效果非常逼真而將整幅圖像判別為真，其更注重輸出圖像整體的轉換效果。

pix2pix的訓練數據為成對的圖像(x，y)，其中,x為源圖像，y為目的圖像,將x輸入到生成器可得到圖像z，將z和x基于通道維度合并在一起，作為判別器的負樣本輸入，得到輸入為真實圖像的概率值，并且y和x也會進行通道維度合并，作為判別器的正樣本輸入,得到輸入圖像為真實圖像的概率值。

pix2pix的損失函數包含2部分：cGAN損失和L1損失,如式(1)～式(3)所示：

(1)

LcGAN(G,D)=Ex,y[logD(x,y)]+

Ex,rn[log(1-D(x,G(x,rn)))]

(2)

LL1(G)=Ex,y,rn[‖y-G(x,rn)‖1]

(3)

其中，λ表示L1損失函數在整個損失函數中的比例系數，整體損失函數是在cGAN損失的基礎上添加L1損失，使得生成的圖像除了能夠欺騙過判別器以外，還需要盡量接近目標域的圖像；rn(random noise)表示隨機噪聲；Ex,y[logD(x,y)]表示logD(x,y)的數學期望；Ex,rn[log(1-D(x,rn))]表示log(1-D(x,rn))的數學期望；Ex,y,rn[‖y-G(x,rn)‖1]表示‖y-G(x,rn)‖1的數學期望。

2.2 DS-GAN的設計

pix2pix在訓練判別器時，將生成的圖像等分為多個70*70大小的圖像塊(patch)，并分別判斷每個patch區域的真假，最后取平均值作為其輸出結果，使用該方法生成的圖像從整體上更加趨近于真實圖像。然而，在工業缺陷圖像中，缺陷往往只占據整幅圖像的小部分，對于一幅生成的工業缺陷圖像，大部分patch區域為背景特征，生成圖像背景區域的逼真程度直接決定了pix2pix的判別結果。在訓練過程中，為了生成能夠欺騙過判別器的圖像，生成器將更加注重學習背景區域的特征而忽略了缺陷區域的學習，雖然生成的工業缺陷圖像在背景上比較接近真實圖像，但仍存在生成缺陷質量不高的問題。因此，pix2pix生成的圖像難以滿足工業檢測的數據增強要求。

基于此，本文對pix2pix的判別器進行了重新設計，并對損失函數進行了改進。改進思路如下：(1)針對判別器的改進：本文提出了僅將圖像局部缺陷作為判別器輸入的策略，在改變判別器輸入后，針對輸入缺陷的大小，重新設計了判別器的網絡結構，引入了更恰當的卷積核用于提取缺陷特征。(2)針對生成對抗損失的改進：本文僅將缺陷區域的判別結果作為生成對抗損失的輸入，分別計算一幅圖像中所有缺陷的損失，最后將所有缺陷的損失進行加權平均，作為整幅圖像的生成對抗損失。

2.3 DS-GAN流程

DS-GAN的流程主要包含3個部分：制作配對數據、訓練生成器及判別器和生成缺陷圖像。

(1)制作配對數據。

本文將缺陷圖像縮放至256*256的固定尺寸，使用矩形框標記出圖像中的缺陷區域，并根據缺陷的形狀，將缺陷區域的3個通道的像素值設置為255，得到與缺陷圖像一一配對的帶有白斑的圖像，記作掩膜(mask)圖像。

(2)訓練生成器及判別器。

本文將帶有矩形標記的缺陷圖像和掩膜圖像配對后輸入生成器，得到生成圖像。根據矩形標記分別提取缺陷圖像、生成圖像和掩膜圖像中標記區域的像素，得到裁剪圖像。將裁剪的掩膜圖像和缺陷區域圖像配對，作為正樣本輸入判別器，將裁剪的掩膜圖像和生成的缺陷區域圖像配對，作為負樣本輸入判別器，判別器輸出裁剪圖像為真實缺陷的概率。在訓練階段，固定生成器參數，通過優化判別器損失，使其能夠準確判斷輸入圖像是否為真；再固定判別器參數，通過不斷優化生成器參數，使其能夠將掩膜圖像轉換成逼真的缺陷圖像來欺騙判別器。通過生成器和判別器的反復對抗訓練，掩膜圖像能夠通過生成器轉換為缺陷圖像。

(3)生成缺陷圖像。

經過生成對抗訓練得到模型，并在收集到的無缺陷圖像中，人工將一些可能出現缺陷的區域的像素值設置為255，得到大量掩膜圖像。將掩膜圖像輸入生成器，得到合成的缺陷圖像。

2.4 DS-GAN生成器

本文生成器網絡結構為U-Net。U-Net是一種encoder-decoder模型，encoder通過一系列卷積操作提取輸入圖像的特征，每經過一次卷積操作，特征圖大小減小至原來的一半，通道數變為原來的2倍；decoder以encoder輸出的高維特征圖作為輸入，通過反卷積操作對輸入特征圖進行上采樣，每經過一次反卷積操作，輸出特征圖的大小將擴大為原來的2倍，通道數縮小為原來的一半，由于encoder和decode結構對稱，因此decoder最終能將高維特征還原至輸入圖像大小。U-Net采用了跳躍連接，其將第i層的特征圖和第n-i層的特征圖進行拼接(其中n為總的特征層數)，通過這種方式將底層信息直接復制到高層特征圖中，在進行圖像轉換時保留一些如顏色、輪廓等低級信息。網絡結構中通過添加LeakyReLU、ReLU和Tanh等激活函數來增加神經網絡的表達能力，添加BN(Batch Normalization)層來進行歸一化處理。

2.5 DS-GAN判別器

PSGAN用于生成街道圖像中的行人，其為了解決pix2pix生成的行人圖像質量差的問題，在pix2pix的基礎上增加了行人判別器，對圖像中的行人進行裁剪后作為行人判別器的輸入，以此來強化“行人”區域的學習。受PSGAN中行人判別器設計思路的啟發，本文改變了pix2pix判別器的輸入，對真實圖像和生成圖像的缺陷區域進行裁剪，分別作為判別器的正樣本輸入和負樣本輸入，以此強化“缺陷區域”的學習。此時，生成圖像的背景部分將不再參與判別器結果的計算，判別器的結果用于反映輸入的缺陷是否為真。在訓練過程中，生成缺陷與真實缺陷的分布存在較大差異時，判別器則將其判別為假，在這種情況下，生成器需要更加注重缺陷區域的學習，才能夠生成逼真的缺陷并欺騙過判別器。本文改變了傳統方式中將整幅圖像作為判別器輸入的方式，僅將裁剪后的缺陷區域、掩膜區域和生成圖像的缺陷區域進行配對，作為判別器輸入，使得判別器完全關注缺陷區域的信息，而忽略背景區域的信息。通過這種方式提升判別器對缺陷區域的注意力，進而使生成器生成高質量缺陷。本文所采用的掩膜圖像和真實圖像的背景區域是完全一致的，這意味著背景區域的映射關系相對簡單，背景部分在不進行生成對抗訓練的情況下，生成器僅通過優化L1損失就能夠學習到其特征。

pix2pix判別器的卷積核是用于提取256*256分辨率的圖像特征。本文實驗數據中缺陷的分辨率大部分小于35*35，對于小于35*35像素的缺陷，經過第3個卷積操作后，其特征圖的寬高為3*3，而第4層的卷積核大小為4*4，存在中間特征層的寬高小于卷積核的情況，顯然，pix2pix判別器不能完成輸入圖像的特征提取，更無法基于提取到的特征完成圖像真偽的判別。針對以上問題，本文重新設計了判別器卷積核的大小，使判別器能夠完成裁剪缺陷的特征提取，并能夠準確判斷裁剪缺陷的真假。DS-GAN判別器網絡結構如圖1所示，其由6個卷積核和一些激活函數組成。輸入的配對缺陷圖像先經過5個卷積核大小為2*2、步長為2、填充寬度為1的卷積操作，然后再經過1個卷積核大小為2*2、步長為1、填充寬度為0的卷積操作，最后通過Sigmoid激活函數將輸出結果轉換成0～1的概率值。

Figure 1 Discriminator network structure

2.6 DS-GAN損失函數

pix2pix中的生成對抗損失為所有patch區域的平均損失，而一個patch區域既可能包含缺陷特征，又可能包含背景特征，在缺陷區域遠小于背景區域的工業芯片中，缺陷區域的損失占總損失的比值較小，導致缺陷的分布不能夠被生成器很好地學習。為了解決該問題，本文將圖像中缺陷區域的判別結果作為生成對抗損失的輸入，圖像的生成對抗損失直接由缺陷損失組成，在不斷優化損失的過程中，生成的缺陷將越來越接近真實缺陷的分布。由于一幅缺陷圖像中往往存在多處缺陷，且不同圖像中缺陷的數量也可能不同，因此圖像間缺陷數量的差異可能導致不同圖像的損失相差很大，進而造成網絡難擬合。為了避免該問題的發生，本文采用加權平均的策略來計算圖像中的缺陷損失。本文損失函數主要包含生成對抗損失和L1損失，如式(4)所示，其中，生成對抗損失如式(5)所示,L1損失函數如式(6)所示：

(4)

LGAN(G,D)=

Epi～mask(pi),ri～generator(ri)[log(1-D(pi,ri))])

(5)

LL1(G)=Exmask,ydefect[‖ydefect-G(xmask)‖1]

(6)

其中，s表示一幅真實圖像中有s個缺陷；一幅掩膜圖像中有s個掩膜區域；pi表示掩膜圖像中的第i個掩膜區域；qi表示真實圖像中對應的缺陷區域；ri表示在生成圖像中裁剪的第i個缺陷區域；Epi～mask(pi),qi～gt(qi)[logD(pi,qi)]表示logD(pi,qi)的數學期望；Epi～mask(pi),ri～generator(ri)[log(1-D(pi,ri))]表示log(1-D(pi,ri))的數學期望；xmask表示帶有掩膜區域的圖像；ydefect表示真實的缺陷圖像；Exmask,ydefect[‖ydefect-G(xmask)‖1]表示‖ydefect-G(xmask)‖1的數學期望。

3 實驗與分析

3.1 實驗內容

實驗共分為2個部分：(1)缺陷生成實驗，分別使用pix2pix和DS-GAN生成缺陷圖像，比較生成圖像的FID(Frechet Inception Distance)[18]，衡量2種方法生成的圖像質量。(2)缺陷檢測實驗，分別使用pix2pix和DS-GAN來進行數據增強，比較2種數據增強方法對RetinaNet[19]算法的平均精度AP(Average Precision)的影響，評估pix2pix和DS-GAN的數據增強能力。

本文實驗采用的數據是工業LED(Light-Emitting Diode)芯片表面圖像，數據集由公司提供，LED芯片及缺陷特征如圖2所示，圖像分辨率為256*256，圖中用白色框標記的為“氣泡”缺陷。實驗所采用的操作系統為Ubuntu18.04，GPU為NVIDIA GeForce GTX Titan X，顯存為12 GB，CPU為Intel(R)Core(TM)i7-7700K，內存為32 GB，實驗算法基于PyTorch框架，采用的Python版本為3.6.12。

Figure 2 LED chip and defect characteristics

3.2 評價指標

本文采用FID和AP2個指標對生成數據的質量進行評估。

FID是度量真實圖像和生成圖像的特征向量之間的距離，其將真實數據和生成數據輸入Inception V3網絡用于提取圖像特征，并將特征向量歸納為多變量高斯分布，使用Wasserstein-2距離來計算2個分布之間的距離。FID越小，表明生成的圖像質量越高。FID的計算公式如式(7)所示:

FID(a,G(rn))=‖μa-μG(rn)‖2+

(7)

其中，a、rn和G(rn)分別表示真實圖像、隨機噪聲和基于隨機噪聲生成的圖像；μa和μG(rn)表示真實圖像和生成圖像對應特征空間的均值；Ca和CG(rn)表示對應特征空間的協方差矩陣，tr()表示矩陣的跡。

AP是衡量目標檢測算法精確度的指標，其值為P-R曲線圍成的面積，其中，P表示準確率,R表示召回率,P和R的計算公式分別如式(8)和式(9)所示:

P=TP/(TP+FP)

(8)

R=TP/(TP+FN)

(9)

其中，TP為被正確檢測為正例的個數，FP表示被錯誤檢測為正例的個數，FN表示被錯誤檢測為負例的個數。

3.3 數據生成實驗結果及分析

在工業場景中，部分類別的缺陷數據十分稀少，為了使實驗條件更加接近現實場景，本文所使用的訓練數據共50幅，其中氣泡缺陷為230個。訓練數據制作方法為：將缺陷數據氣泡區域的3個通道的像素值設置為255，得到掩膜圖像，將掩膜圖像和缺陷圖像進行配對，得到訓練數據。本文實驗使用配對數據集分別訓練pix2pix和DS-GAN，并使用2個模型在148幅掩膜圖像上生成缺陷，其中生成缺陷共823個。實驗具體參數如下：pix2pix和DS-GAN均訓練300輪次(epoch)；Adam初始學習率為0.000 2，Adam動向量為0.5，學習率在前150個輪次保持不變，在后150個輪次線性減少為0，生成器L1損失系數為10.0。

實驗生成缺陷圖像如圖3所示，對比DS-GAN和pix2pix生成圖像的背景部分，兩者均非常接近真實圖像的背景部分。而對比2種方法生成圖像的缺陷部分，可以發現：pix2pix生成的缺陷比較模糊，且缺陷部分和背景部分有一條明顯的界限，存在缺陷和背景融合不自然的問題；DS-GAN生成的缺陷則相對清晰，且缺陷與背景融合得更加自然，從圖像上看不出二者明顯的界限。FID評估結果如表1所示，其中，情況1為真實數據的FID，情況2為pix2pix生成數據的FID，情況3為本文方法生成數據的FID。經過對比可以看出，DS-GAN生成缺陷圖像的FID要低于pix2pix生成圖像的FID，從評估指標上看，DS-GAN生成的缺陷圖也更接近真實圖像的分布。

Figure 3 Images generated by different methods

Table 1 FID comparison of images generated by different data augmentation methods

以上結果表明：(1)從視覺上看，DS-GAN生成的背景區域與真實圖像的背景區域保持了高度一致，而DS-GAN中背景區域并未進行生成對抗訓練，表明DS-GAN中背景部分的映射關系是通過優化L1損失來學習的，其訓練不依賴生成對抗損失；(2)從圖像效果上看，DS-GAN生成的缺陷區域更加接近真實缺陷，表明將背景區域進行對抗訓練不會提升生成圖像背景區域的質量，背景區域的引入反而使得判別結果無法準確表示生成缺陷的質量，從而影響生成器的訓練，導致生成的缺陷區域質量不高；(3)對比圖3中生成圖像的缺陷細節，可以看出本文方法對判別器和損失函數的改進是行之有效的。僅將缺陷區域輸入判別器的策略使得判別器能夠充分學習到真實缺陷區域和生成缺陷區域的特征，且優化每個缺陷區域的生成對抗損失的操作有效地增強了缺陷區域的訓練，使得判別器對缺陷區域具有更準確的判斷能力，進而優化生成器生成更加高質量的缺陷，顯著降低了生成缺陷的FID。

3.4 缺陷檢測實驗

數據預處理流程如下：使用LabelImg標記工具分別對50幅真實缺陷圖像、200幅pix2pix生成圖像、200幅DS-GAN生成圖像進行標記，并將標記好的數據轉換成COCO數據集格式。具體實驗中，首先使用真實圖像對RetinaNet進行訓練；再分別將pix2pix和DS-GAN生成的200幅缺陷數據加入訓練集，訓練2個RetinaNet檢測模型用于實驗對比；最后使用148幅真實數據用于缺陷檢測模型的測試。

從表2中pix2pix的實驗結果可以看出，pix2pix生成圖像加入訓練之后，對缺陷檢測模型的精度提升并不明顯，表明通過pix2pix生成缺陷的分布偏離了真實缺陷的分布，生成圖像只能稍微提高模型精度。從表2中DS-GAN的實驗結果可以看出，DS-GAN生成圖像加入訓練之后，能夠有效提升缺陷檢測模型的精度，表明DS-GAN生成的圖像和真實圖像的分布是接近的，DS-GAN可以用于工業數據增強。

綜合數據生成實驗和工業缺陷檢測實驗，DS-GAN生成的數據更具數據增強能力，對目標檢測性能的提升更加明顯。

Table 2 Influence of different data augmentation methods on object detection accuracy

4 結束語

針對工業缺陷數據收集成本高，使用pix2pix生成的缺陷數據難以滿足深度學習訓練要求這一問題，本文提出了一種基于pix2pix改進的工業缺陷數據增強方法。一方面，對工業芯片圖像中的缺陷區域進行裁剪，作為判別器的輸入，并針對缺陷區域大小，重新設計了判別器的網絡結構；另一方面，重新設計了生成對抗損失的計算方式，單幅圖像的生成對抗損失由圖像中的所有缺陷區域的生成對抗損失組成，通過優化生成對抗損失，引導生成器學習掩膜區域和局部缺陷的映射關系，使得生成器能夠生成高質量缺陷。

通過數據生成實驗和缺陷檢測實驗的實驗結果可知，相較于pix2pix，本文方法提升了生成缺陷數據的質量，豐富了工業芯片缺陷數據集，改善了工業缺陷數據難收集的問題，并且使用該方法生成的數據顯著增加了基于深度學習進行工業缺陷檢測的精度。然而，該方法并不適用所有工業缺陷數據的生成，當訓練樣本中的缺陷數量過少同時特征多變時，該方法生成的缺陷數據存在不夠真實的問題。未來將繼續對GAN網絡進行研究，使得GAN網絡能夠在缺陷特征更加復雜的情況下，生成的數據依然具有很高的質量。