基于CycleGAN和Pix2Pix的軌面缺陷圖像智能生成技術

陳嘉欣，孫傳猛，葛耀棟，李欣宇，靳書云，李 勇

(1.中北大學 省部共建動態測試技術國家重點實驗室，山西 太原 030051；2.中北大學 電氣與控制工程學院，山西 太原 030051；3.華北計算機系統工程研究所，北京 100083；4.重慶大學 煤礦災害動力學與控制國家重點實驗室，重慶 400044)

在國家“科技強國”“交通強國”戰略支撐下,高速鐵路等新型載運工具得到了普遍關注和認可,并積極進行布局。其中,高動態、高精度軌道巡檢技術對高速鐵路的智能運維至關重要[1]。然而,基于人工智能技術開展高速條件下鋼軌軌面缺陷自主檢測面臨以下小樣本學習難題:

1)數據分布不平衡[2]。已有軌道在通常情況下都是無缺陷的,對剛修建完或正在建設中的高速鐵路軌道而言,獲得軌面缺陷數據更是難上加難。顯然,目前構建的軌面缺陷數據集存在正負樣本分布不平衡問題,這嚴重降低了相關缺陷檢測模型的性能,極端情況下模型會將缺陷識別為噪聲從而導致缺陷檢測結果無效。

2)數據缺乏多樣性。軌面缺陷數據獲取困難,相關數據集規模偏小,這意味著數據正負樣本的形態偏少,背景與目標樣式多樣性不足,導致模型只能學到部分特征而影響模型的識別性能,甚至引發“過擬合”問題。

3)數據標注難度大[3]。軌面缺陷數據的標注存在鑒定難等問題,需要專業人員完成。此外,通過實地拍攝得到覆蓋面廣的數據集并準確標注,將會耗費大量的人力、物力。

針對上述小樣本學習難題,一種可靠的解決方法是遷移學習,它通過人工生成缺陷形態符合真實分布的、目標與背景近似真實的模擬圖像來使數據集的增強,訓練增強數據集實現缺陷檢測模型的卓越性能;隨著新的、真實的缺陷數據不斷加入訓練集,不斷微調已有模型,實現性能的最優化。

顯然,模擬圖像生成的質量直接關系到遷移學習的成敗。生成對抗網絡(generative adversarial network, GAN)[4-8]是一種基于博弈論的深度學習生成模型,可以學習樣本數據高階語義特征并生成具有該特征的新數據。GAN模型是在2014年由Goodfellow等提出,相對其他數據增強方法,其生成的圖像具有清晰度、分辨率更高且速度更快等優點[9-10]。在圖像生成領域,相關學者對GAN模型做了大量研究工作。2019年,Bailo等[11]采用條件生成對抗網絡(conditional generative adversarial network, CGAN),對顯微鏡下的紅細胞圖像進行生成,得出在生成數據對檢測任務的增益大于分割任務的結論。2020年,Suh等[12]提出分類增強型生成對抗網絡(classification enhancement generative adversarial networks, CEGAN),可以在數據不平衡的情況下生成合成的少數群體數據來增強分類,提高了非平衡數據條件下分類的穩定性和性能。同年,Niu等[13]提出一種新的缺陷圖像生成方法,引入兩個鑒別器和周期一致性損失來生成換向器缺損圖像。2021年,Belhiti等[14]改進了StackGAN,實現對人臉數據的生成。同年,Cui等[15]提出生成式協同訓練網絡GenCo,通過引入多個互補的鑒別器來緩解鑒別器過擬合的問題,這些鑒別器在訓練中可以在多個不同的尺度上進行檢測。

雖然軌面缺陷形態無定式,但類別是固定的。因此,生成的軌面缺陷模擬圖像,既不能樣式單一、與小樣本數據集中缺陷形態完全一致,又不能形態過分夸張、跳脫已有類別限制。因此,本文采用CycleGAN和Pix2Pix相結合的方法,實現小樣本數據集語義特征學習和軌面缺陷數據的智能生成。其中,Pix2Pix輸入的是成對的缺陷圖像和背景圖像,生成特定類別的缺陷圖像;CycleGAN模型則生成不受定式約束的缺陷圖像。二者結合,實現模擬圖像缺陷類別不變而形態各異,從而達到數據集增強的目的。

1 相關理論

1.1 GAN

GAN[4]結構見圖1。

圖1 GAN模型基本結構

通過生成器網絡G和判別器網絡D不斷進行二元極大、極小博弈,達到同時優化兩個模型的目的。其中,生成器G的目標是提高生成樣本與真實樣本的相似度,即實現生成樣本pg與真實樣本pd分布相同;判別器D的目標是分辨真實樣本與生成樣本。GAN損失函數(目標函數)LGAN的表達式為

E(ln{1-fD[fG(z)]})}

(1)

式中:x為圖像數據;z為隨機噪聲;E為期望值;fD為判別器D的輸出;fG為生成器G的輸出。

1.2 Pix2Pix

經典GAN模型生成內容由參數fG和隨機噪聲z決定,無法控制生成內容。為控制生成器生成內容,引入條件信息y。在Pix2Pix模型中,y為與m對應的另一個圖像域數據。在生成器中,隨機噪聲z和數據y同時作為輸入,生成一個跨模態特征;在判別器中,數據m及對應的y作為輸入,生成跨模態向量,同時判斷m的真偽。這樣,隨著條件信息y的引入,生成器逐漸被控,生成特定結果。加入條件信息y,Pix2Pix的損失函數[16-17]Lpix2pix的表達式為

E(ln{1-fD[fG(z|y)]})}

(2)

通常,在損失函數中引入L1正則項可以使模型參數稀疏化,從而提升泛化能力,有助于生成更清晰的數據[18],則L1正則項損失函數LL1為

LL1(fG)=E[‖y-fG(m,z)‖1]

(3)

在正則約束情況下,添加懲罰因子λ為0.01的懲罰項,則模型參數fG*為

(4)

不同于經典GAN使用MLP作為模型結構,Pix2Pix使用了CNN中常用的卷積+BN+ReLU的模型結構。其中,生成器采用了經典的編碼器-解碼器結構[19],即UNet,其典型結構見圖2。UNet可以保留不同分辨率下像素級的細節信息,實現將高分辨率輸入網格映射到高分辨率輸出網格;判別器采用PatchGAN[20],這不同于經典GAN,將整幅圖作為判別器判別的目標,而是將輸入圖像分成若干個圖像塊(Patch),依次作為判別目標。當圖像塊的大小為1×1時,PatchGAN又被叫做PixelGAN。Pix2Pix對一張圖片切割成不同的N×N大小的Patch,判別器對每一個Patch做真假判別,將一張圖片所有Patch的結果取平均作為最終的判別器輸出。PatchGAN判別器結構見圖3。

圖2 UNet典型結構

圖3 PatchGAN判別器結構

最終,Pix2Pix形成如圖4所示的結構。圖4中,黑白Mask標注即為引入的條件信息y,該標注為二值圖像,反映了缺陷所對應的位置和形狀大小信息。

圖4 Pix2Pix模型基本結構

1.3 CycleGAN

Pix2Pix模型訓練需要輸入成對的缺陷圖像和背景圖像。顯然,實際成對的軌面圖像數據是很難獲取的。CycleGAN[21],即循環生成對抗網絡,采用雙向循環生成的結構,突破Pix2Pix對數據集圖片一一對應的限制。

CycleGAN結構見圖5。圖5中包含兩個生成器網絡和兩個判別器網絡。生成器網絡結構見圖6。其中,兩個生成器網絡結構相同,包含下采樣、轉碼和上采樣3部分;兩個判別器結構采用PatchGAN,同Pix2Pix模型。

圖5 CycleGAN模型基本結構

圖6 生成器網絡結構

設無缺陷軌面圖像為X域,有缺陷軌面圖像為Y域,則生成器網絡G1的映射函數fG1為

fG1(x)=y′x∈X

(5)

生成器網絡G2的映射函數fG2為

fG2(y)=x′y∈Y

(6)

兩個生成器可以實現跨域圖像生成,兩個判別器則分別判斷兩個生成器生成數據與真實數據之間的差異,這與經典GAN網絡是一致的。

由于輸入CycleGAN的兩組(X、Y域)圖像沒有一一對應關系,即輸入的兩組圖片完全是隨機搭配在一起,常規GAN以及Pix2Pix學不到任何聯系,而CycleGAN通過添加“循環生成”并優化循環一致性損失(consistencyloss)來代替Pix2Pix中的約束條件,使生成器網絡保留原域圖片特征。循環一致性體現在:由生成器網絡G1生成的Y域數據,再由生成器網絡G2恢復為X域數據x′,這與初始輸入的X域x應該是一致的。因此,可由L1損失來判斷x與x′的差異,生成器的循環一致性損失Lcyc的表達式為

Lcyc(fG1,fG2)=E{‖fG2[fG1(x)]-x‖1}+

E{‖fG1[fG2(y)]-y‖1}

(7)

同經典GAN一樣,CycleGAN還包含對抗性損失。其中,由X域到Y域的對抗損失為

LGAN(fG1,fD2,X,Y)=E[lnfD2(y)]+

E(ln{1-fD2[fG1(x)]})

(8)

式中: fD2為判別器網絡D2對應的映射函數。

由Y域到X域的對抗損失為

LGAN(fG2,fD1,Y,X)=E[lnfD1(x)]+

E(ln{1-fD1[fG2(y)]})

(9)

式中: fD1為判別器網絡D1對應的映射函數。

此外,CycleGAN還引入了IdentityLoss來確保生成器網絡所生成的數據盡量保留源圖像特性(如色調)。IdentityLoss的表達式為

LIdent(fG1,fG2)=E[‖fG1(x)-x‖1]+E[‖fG2(y)-y‖1]

(10)

這樣,CycleGAN總的損失函數表達式為

λ2LIdent(fG1,fG2)

(11)

式中:λ1、λ2為正則化系數。

1.4 智能生成總體流程

綜上所述,采樣Pix2Pix模型可以約束生成缺陷圖像的類型,采用CycleGAN模型可以生成不受定式約束的軌面缺陷圖像。這樣,將二者結合即可實現模擬圖像缺陷類別均衡而形態各異,達到數據集增強的目的,從而解決基于人工智能的高動態高精度軌道巡檢技術的小樣本學習難題,其具體過程見圖7。

圖7 總體流程

基于CycleGAN和Pix2Pix的軌面缺陷圖像智能生成方法的偽代碼如下:

初始化:初始化Mp、Mc。

1.for(t=1,…,T)do

2. 訓練模型Mp中的生成器Gp、判別器Dp;

3. 損失函數反向傳播,更新模型參數;

4.endfor

5.for(t=1,…,T)do

6. 更新模型Mc中的生成器G1、G2,判別器D1、D2;

7. 損失函數反向傳播,更新模型參數;

8.endfor

9.設定生成黑白Mask標注類型,使用Gp得到缺陷圖像。

10.輸入軌面無缺陷圖像和缺陷圖像,使用G1、G2對應生成軌面缺陷圖像和無缺陷圖像。

2 實驗研究

2.1 實驗方案

為驗證本文技術方法的可行性,設計實驗方案包含以下兩個方面:①制作軌面圖像數據集。數據集包含有缺陷數據集和無缺陷數據集兩大類,有缺陷數據的缺陷均有對應的黑白Mask標注。②分別搭建Pix2Pix和CycleGAN生成對抗網絡,分別輸入相應的數據訓練各自模型,并考察其生成圖像質量。

本次實驗平臺為Pycharm2020,所用GPUNVIDIARTX6000,開發語言為Python3.6,配置環境為Pytorch1.7。

2.2 軌面缺陷數據集制作

本文采用RSDDsType-Ⅰ數據集[22]作為實驗對象。相比其他故障檢測數據集,RSDDs檢測難度更大,體現如下: ①識別功能有限。該數據集視覺特征的局限性使得大多數基于復雜紋理和形狀特征的目標識別方法無法成功訓練[23]。②缺乏可識別的先驗信息。軌面缺陷往往以隨機或任意的方式出現,其形狀和大小各不相同。③較低的灰度值先驗[24]往往被光照不均一、反射特性變化和噪聲引起的不規則物所抑制。

RSDDsType-Ⅰ數據集共有67張真實缺陷軌面圖片和其缺陷對應的黑白Mask標注。數據集部分圖片及其黑白Mask標注見圖8。RSDDsType-Ⅰ數據集如何判斷缺陷說明見文獻[22]。

圖8 RSDDs Type-Ⅰ數據集部分樣本

原始的RSDDsType-Ⅰ數據集不能滿足本實驗方案要求,必須再次加工獲得滿足本次實驗的數據集樣本。首先,利用矩陣運算將原始樣本圖像裁剪成256×256大小的多個新樣本;然后,根據對應Mask標簽將裁剪后的圖片分為有缺陷軌面數據集和無缺陷軌面數據集;最終,制作的有缺陷軌面數據集包含樣本141張及對應黑白Mask標注,部分樣本見圖9(a),無缺陷軌面數據集包含樣本465張,部分樣本見圖9(b)。

圖9 有、無缺陷軌面數據集部分樣本

2.3 實驗結果與分析

2.3.1 基于Pix2Pix的軌道圖像生成實驗

將成對的有缺陷軌面數據集樣本與其黑白Mask標注圖像輸入到Pix2Pix模型進行訓練,訓練共計800輪。為防止過擬合,在Pix2Pix架構中還設置了Dropout層。

訓練完成后,隨機抽取若干樣本測試網絡性能,測試結果見圖10。由圖10可知,通過控制輸入到Pix2Pix網絡的黑白Mask形狀,可以實現對生成缺陷形態的精確控制。

圖10 Pix2Pix模型隨機抽樣測試結果

為進一步驗證Pix2Pix的軌面圖像生成效果,需要更多的Mask圖像,達到豐富數據集的目的。Mask圖像采用如下兩種方法生成: ①對RSDDsType-I數據集自帶的Mask進行隨機變換,包括縮放、鏡像、移位等操作; ②利用規則圖形生成Mask圖像。方法②的具體步驟如下:

Step1生成256×256大小的黑色背景。

Step2生成一個隨機數n作為缺陷數量。

Step3隨機生成一個坐標(x,y)和半徑r,以此為參數畫圓并填充為白色。

Step4確定橫向壓縮率dx、縱向壓縮率dy、旋轉角度θ,并對白色目標進行對應的壓縮、旋轉操作。

Step5不斷重復Step3、Step4,直至完成所有的n個缺陷目標的生成,然后轉入下一個Mask的制作。

Step6利用上述兩種方法獲得了理想的軌面缺陷Mask圖像,部分樣本見圖11。其中,圖11(a)由方法①得到;圖11(b)由方法②得到。

圖11 兩種方法生成的部分黑白Mask圖像

將Mask圖像輸入到訓練好的Pix2Pix模型,生成有缺陷軌面模擬圖像,結果見圖12(因篇幅有限,僅展示部分圖像)。由圖12可知,隨機生成的Mask圖像均能得到對應的模擬缺陷圖像,表明Pix2Pix模型能夠生成特定類別的缺陷圖像。

圖12 設定Mask生成的模擬缺陷圖像

2.3.2 基于CycleGAN的軌面圖像生成實驗

將465張無缺陷軌面數據集樣本作為真實X域數據,將141張有缺陷軌面數據集樣本作為真實Y域數據,輸入CycleGAN進行訓練。訓練共計200輪。

訓練完成后,隨機確定軌面圖像進行CycleGAN模型效果驗證,結果見圖13。由圖13可知,CycleGAN模型成功地將有缺陷軌面圖像轉換成無缺陷軌面圖像,將無缺陷軌面圖像轉換成有缺陷軌面圖像,且缺陷樣式不受定式約束。

圖13 CycleGAN生成的圖像

3 性能測試

本文方法生成的軌面缺陷模擬圖像,從圖像內容和缺陷形態而言,十分接近真實的軌面缺陷圖像。但對于小樣本數據集而言,存在數據分布不平衡、數據缺乏多樣性等問題。由于生成器固有特性,無法利用劃分數據集的方式來評估生成對抗網絡模型生成圖像的性能[25]。因此,采用本文方法生成軌面缺陷模擬圖像進行數據集增強后,是否解決了數據分布平衡、樣本多樣性等問題,當前并沒有恰當的評價指標。如果增強后的數據集在計算機視覺任務中能更好地發揮各自網絡模型的性能,則能間接反映增強數據集解決了原數據集小樣本學習難題。因此,本文利用圖像分類、目標檢測、語義分割這3大典型的計算機視覺任務來測試本文方法生成圖像性能,進而驗證該方法解決軌面缺陷檢測技術小樣本學習難題的有效性。在訓練過程中,所有實驗均同時采用了在線數據增強,最終實驗結果僅突出了數據集擴充后的有效性。

3.1 基于VGG19的圖像分類性能測試

VGG19是牛津大學視覺幾何小組提出的一種深層卷積網絡結構,在圖像分類任務中性能卓越、廣受認可[26]。VGG19訓練不需要Mask標注,只需將有、無缺陷軌道圖像分別賦予不同的分類標簽輸入網絡即可。

首先,構造原始的軌面缺陷小樣本數據集。為模擬數據分布失衡,用于訓練的有、無缺陷軌面圖像樣本分別為20、311張;用于測試的有、無缺陷軌面圖像樣本分別為121、134張,見表1。表1中,def、no_def分別表示有、無缺陷軌面樣本。

表1 原始數據集訓練和測試劃分 張

然后,利用Pix2Pix和CycleGAN對有缺陷軌面數據集進行增強,結果見表2。表2中,ori為原始數據集;aug(Pix2Pix2)、aug(CycleGAN)和aug(Mix)分別為Pix2Pix2、CycleGAN和兩者混合增強的數據集。

表2 增強的有缺陷數據集結果 張

將增強后的有缺陷軌面數據均加入到訓練集中,這樣新的增強數據集中有缺陷和無缺陷樣本數量相等。

最后,分別用原始數據訓練集和增強數據訓練集訓練具有同一結構的VGG19模型,并利用測試集(兩數據集的測試集完全相同)測試各自訓練模型的性能,結果見表3。表3中,準確率為正確分類的樣本數與總樣本數之比。由表3可知,采用Pix2Pix和CycleGAN增強數據集均大幅度提高了VGG19分類模型的性能,而聯合CycleGAN和Pix2Pix增強方法提升幅度最大,達81.177%。這表明本文方法在圖像分類任務中,可顯著改善小樣本特征帶來的數據分布不平衡、缺乏多樣性等問題。

表3 基于VGG19的圖像分類性能測試結果 %

3.2 基于YOLOv5的目標檢測性能測試

YOLOv5是近年來性能最卓越的目標檢測算法之一。本文利用YOLOv5來測試增強數據集對目標檢測算法性能的提高能力。

使用與3.1節所述中相同的有缺陷軌面圖像樣本作為訓練集和測試集,且在不同數據集的實驗中均設置相同的超參數以對比檢測結果。以準確率、召回率以及均值平均精度作為檢測結果的評價指標,數值越高表明對缺陷的定位和識別越加準確。對3種增強數據集進行實驗,結果見表4。表4中,召回率為正樣本被預測正確數量與總正樣本個數之比;均值平均精度為計算所有類別的準確率和召回率曲線下面積的平均值。由表4可知,聯合Pix2Pix和CycleGAN增強數據集方法使得均值平均精度達到最高,相比原數據集提升了18.30%。這表明,聯合兩種方法豐富數據集對模型泛化能力的提升效果最為顯著,部分檢測結果見圖14。

圖14 原始與3種增強數據集的部分缺陷檢測結果

3.3 基于UNet的語義分割性能測試

UNet是一種經典的語義分割深度學習模型。本文利用UNet來測試增強數據集對語義分割算法性能的提高能力。該測試試驗所用的訓練集和測試集同樣為3.1節所述數據集。

由于UNet是有監督學習,需要為訓練集的各個樣本制作相應的標簽圖像。對于原始訓練集,各樣本對應的黑白Mask圖像即可作為標簽;對于Pix2Pix生成的有缺陷軌面數據,其輸入的Mask圖像可以直接作為標簽;對于CycleGAN生成的有缺陷軌面數據,由于是在無缺陷數據上生成的,將生成數據和原數據進行減法運算,并采用全局閾值處理方法[27]轉換為二值圖像,然后使用形態學處理濾除噪聲并填補空隙,即可得到相應的Mask標簽圖像,見圖15。

圖15 CycleGAN生成缺陷圖像標簽

分別用原始數據訓練集和增強數據訓練集訓練具有同一結構的UNet模型,并利用測試集測試各自訓練模型的性能。UNet模型的評估指標為Dice Score,該指標主要用于計算兩集合間的相似性,即兩集合間相似性越高,Dice Score越高。UNet模型在各數據集上的性能見表5。表5中,Dice Score為預測結果與真實標簽的重疊區域與兩者并集之比。由表5可知,采用Pix2Pix增強數據集降低了UNet分割性能,這是因為只用Pix2Pix生成數據沒有有效的增強缺陷的多樣性;只采用CycleGAN方法生成的數據,在多樣性上有較大增強,因而提升了UNet分割性能;而聯合CycleGAN和Pix2Pix增強方法在增強數據分布平衡和缺陷多樣性方面更優,因而取得了最好的分割性能,其UNet分割Dice Score達到71.015,相比原始數據集提高了6%。

表5 基于UNet的圖像分割性能測試結果

4 結論

1)采用CycleGAN和Pix2Pix相結合的方法,實現小樣本數據集語義特征學習和軌面缺陷數據的智能生成,生成數據缺陷類別不變、而形態各異,從而達到數據集增強的目的。

2)利用VGG19、YOLOv5和UNet進行的軌道圖像小樣本數據集增強性能測試實驗表明,本文方法在圖像分類任務中的準確率為81.177%,較原數據集提高23.138%;在目標檢測任務中,準確率為91.90%,提升26.60%,召回率為87.20%,提升16.00%,均值平均精度為93.50%,提升18.30%;在語義分割任務中Dice Score為71.015,較原數據集提高6%。