自注意力SRGAN在巖石CT圖像超分辨中的應(yīng)用研究

朱聯(lián)祥，鄭 逸

(西安石油大學(xué) 計(jì)算機(jī)學(xué)院，陜西 西安 710065)

引 言

巖石顯微圖像可以在不改變巖石形態(tài)及內(nèi)部結(jié)構(gòu)的條件下，用于巖石的物性特征分析與計(jì)算，故而被廣泛應(yīng)用于石油地質(zhì)領(lǐng)域[1]。然而在實(shí)際中，受限于采集設(shè)備成本及自然環(huán)境因素，一般難以獲得足夠高分辨率的巖石圖像，從而對(duì)巖石分析的質(zhì)量帶來(lái)影響。單圖像超分辨率技術(shù)可通過(guò)深度學(xué)習(xí)等方式，從單幅低分辨率圖像構(gòu)建出高分辨率的圖像，具有節(jié)約成本且處理時(shí)間短的優(yōu)點(diǎn)。

單張圖像超分辨率重建算法主要經(jīng)歷了基于插值的算法、基于重建的算法和基于深度學(xué)習(xí)的算法3個(gè)階段[2]?；诓逯档乃惴ㄊ歉鶕?jù)已知像素通過(guò)某種算法估計(jì)當(dāng)前像素值，基于重建的算法則是在空域或頻率域進(jìn)行操作。針對(duì)這兩種方法均存在的受圖像內(nèi)容限制及重建圖像質(zhì)量差的問(wèn)題，Dong等人在2016年通過(guò)將深度學(xué)習(xí)引入到圖像超分辨率研究，提出了SRCNN[3]。雖然所提算法只是采用了3層的卷積神經(jīng)網(wǎng)絡(luò)結(jié)構(gòu)，但與雙三次插值(Bicubic)及基于重建的算法相比，無(wú)論是在視覺(jué)效果，還是圖像的峰值信噪比(PSNR)及結(jié)構(gòu)相似度(SSIM)等方面，都有了明顯進(jìn)步。在此之后，亞像素卷積[4]、殘差網(wǎng)絡(luò)[5]、拉普拉斯金字塔[6]等網(wǎng)絡(luò)結(jié)構(gòu)陸續(xù)被應(yīng)用于圖像的超分辨率處理，并取得了較為顯著的效果。

2017年，Ledig等人將生成對(duì)抗網(wǎng)絡(luò)(Generative Adversarial Network，GAN[7])應(yīng)用超分辨率領(lǐng)域，提出了SRGAN[8]。該方法將低分辨率的圖像作為GAN的生成網(wǎng)絡(luò)的輸入，通過(guò)訓(xùn)練得到超分辨率的生成圖像，并將其和原始高分辨率圖像一起作為判別網(wǎng)絡(luò)的輸入。判別網(wǎng)絡(luò)試圖區(qū)分生成的圖像和原始的高分辨率圖像，而生成網(wǎng)絡(luò)則試圖通過(guò)訓(xùn)練網(wǎng)絡(luò)參數(shù)來(lái)生成可以“欺騙”判別網(wǎng)絡(luò)的超分辨率圖像。通過(guò)判別網(wǎng)絡(luò)和生成網(wǎng)絡(luò)反復(fù)地迭代學(xué)習(xí)，最終達(dá)到納什均衡，即生成網(wǎng)絡(luò)可以生成足以亂真的超分辨圖像，判別網(wǎng)絡(luò)只能以50%的概率區(qū)分生成圖像和原始高分辨率圖像，從而實(shí)現(xiàn)低分辨率圖像的超分辨處理。

最早提出并應(yīng)用注意力機(jī)制的是自然語(yǔ)言處理任務(wù)[9]。隨著時(shí)間的推移，注意力機(jī)制逐漸得到廣泛應(yīng)用。相較于傳統(tǒng)的卷積方法，注意力機(jī)制可以通過(guò)得到的稀疏矩陣從全局范圍內(nèi)篩選出重要的部分區(qū)域，從而能更好地處理圖像中的長(zhǎng)范圍依賴關(guān)系，協(xié)調(diào)生成圖像的紋理細(xì)節(jié)，在一定程度上彌補(bǔ)卷積的不足。

本文以SRGAN為基礎(chǔ)，通過(guò)引入自注意力機(jī)制[10]，對(duì)生成器和判別器的網(wǎng)絡(luò)結(jié)構(gòu)進(jìn)行調(diào)整，以更好地表達(dá)高頻特征。同時(shí)考慮到巖石顯微圖像的特殊性，通過(guò)在損失函數(shù)中引入巖石圖像所特有的孔隙度特征來(lái)構(gòu)建新的損失函數(shù)，實(shí)現(xiàn)了適合于巖石顯微圖像分辨率提升的超分辨率處理算法。最后用所提算法實(shí)驗(yàn)實(shí)現(xiàn)了巖石顯微圖像的超分辨率處理，并采用客觀評(píng)價(jià)指標(biāo)對(duì)生成的超分辨率巖石圖像的性能進(jìn)行了評(píng)價(jià)。

1 模型介紹

1.1 生成對(duì)抗網(wǎng)絡(luò)

生成對(duì)抗網(wǎng)絡(luò)由生成網(wǎng)絡(luò)和判別網(wǎng)絡(luò)組成，低分辨率圖像通過(guò)生成網(wǎng)絡(luò)得到超分辨率圖像，判別網(wǎng)絡(luò)對(duì)超分辨率圖像和高分辨率圖像進(jìn)行判斷。在訓(xùn)練過(guò)程中，生成網(wǎng)絡(luò)和判別網(wǎng)絡(luò)的對(duì)抗過(guò)程如下式所示：

(1)

其中：x表示真實(shí)數(shù)據(jù)，Pdata表示其概率分布；z表示生成器的噪聲輸入，Pz為其相應(yīng)的概率分布。在生成對(duì)抗網(wǎng)絡(luò)中，要計(jì)算的納什均衡點(diǎn)正是使生成網(wǎng)絡(luò)和判別網(wǎng)絡(luò)各自損失函數(shù)最小的情況。在上述公式的基礎(chǔ)上，可以分別求得理想的生成網(wǎng)絡(luò)G*和判別網(wǎng)絡(luò)D*：

(2)

(3)

對(duì)于上述的D*，假設(shè)生成網(wǎng)絡(luò)G是固定的，推導(dǎo)如下：

V=Ex～PdatalogD(x)+Ex～Pzlog(1-D(G(z)))

(4)

希望通過(guò)尋找一個(gè)D使得V最大，則應(yīng)對(duì)于積分中的項(xiàng)f(x)=Pdata(x)logD(x)+Pg(x)log(1-D(x))，無(wú)論x取何值都能最大。其中Pdata是固定的，之前假設(shè)了G是固定的，所以Pg也是固定的。通過(guò)求導(dǎo)可得：

(5)

(6)

由于JS散度是非負(fù)的，當(dāng)且僅當(dāng)Pdata=Pg時(shí)，可以取得全局最小值。故所求的最優(yōu)生成網(wǎng)絡(luò)是使得G*的分布滿足Pg=Pdata。

1.2 自注意力機(jī)制

自注意力機(jī)制的使用促使自然語(yǔ)言處理取得了很大的進(jìn)步，同時(shí)上下文信息(context information)對(duì)于諸如語(yǔ)義分割、目標(biāo)檢測(cè)等視覺(jué)任務(wù)至關(guān)重要。自注意力機(jī)制通過(guò)三元組提供了一種有效獲取更大感受野、捕捉全局上下文信息的建模方式(如圖1所示)，推動(dòng)了視覺(jué)任務(wù)的發(fā)展。

圖1 自注意力機(jī)制模塊圖

(7)

2 改進(jìn)方法

2.1 算法原理

在SRGAN中，通過(guò)殘差網(wǎng)絡(luò)提取特征，該網(wǎng)絡(luò)在提取特征時(shí)對(duì)于部分邊緣信息不敏感，由于巖石顯微圖像邊緣復(fù)雜，在提取特征過(guò)程中更容易丟失特征信息。本文在SRGAN的基礎(chǔ)上，引入了注意力機(jī)制模塊，希望通過(guò)注意力機(jī)制模塊加強(qiáng)對(duì)于邊緣信息的描述，并針對(duì)巖石顯微圖像調(diào)整了損失函數(shù)。

2.2 自注意力機(jī)制

本文采用自注意力機(jī)制模塊加強(qiáng)紋理特征和邊緣特征，自注意力機(jī)制模塊與生成網(wǎng)絡(luò)的結(jié)合如圖2所示。自注意力機(jī)制模塊接收來(lái)自上層結(jié)構(gòu)的輸出特征圖，在模塊內(nèi)部經(jīng)過(guò)相似度計(jì)算得到權(quán)重并歸一化，最終將權(quán)重與特征圖加權(quán)求和，得到特征圖中任意兩個(gè)位置的全局依賴關(guān)系。

圖2 自注意力機(jī)制模塊與Generator的結(jié)合

2.3 網(wǎng)絡(luò)結(jié)構(gòu)

本文生成網(wǎng)絡(luò)和判別網(wǎng)絡(luò)如圖3所示，生成網(wǎng)絡(luò)在提取特征時(shí)采用了殘差網(wǎng)絡(luò)和自注意力機(jī)制模塊，放大時(shí)采用了亞像素卷積模塊；判別網(wǎng)絡(luò)由卷積、注意力機(jī)制和激活函數(shù)組成。

圖3 生成網(wǎng)絡(luò)和判別網(wǎng)絡(luò)結(jié)構(gòu)圖

生成網(wǎng)絡(luò)輸入一個(gè)3×56×56的巖石圖像，先后經(jīng)過(guò)一個(gè)卷積PReLU層，5個(gè)殘差塊，1個(gè)自注意力機(jī)制模塊，1個(gè)卷積BN層，和log2upscale個(gè)亞像素卷積塊得到大小為3×224×224的超分辨率圖像，其中卷積核大小均為3×3，步長(zhǎng)和填充均為1。

判別網(wǎng)絡(luò)輸入圖像尺寸為3×224×224，輸出為經(jīng)過(guò)卷積BN層和Sigmoid激活函數(shù)處理之后的概率值。在卷積過(guò)程中仍保持3×3大小的卷積核，同時(shí)為了更好地對(duì)圖像進(jìn)行判斷，在其中也加入了自注意力機(jī)制模塊。

2.4 損失函數(shù)

巖石CT圖像通過(guò)灰度值對(duì)巖石內(nèi)部成分進(jìn)行了表征，可以對(duì)巖體和孔隙進(jìn)行區(qū)分。針對(duì)工業(yè)CT掃描得到的巖石CT圖像，采用圖像邊緣檢測(cè)的方法，可以對(duì)CT圖像進(jìn)行分割，得到巖石中固體成分和孔隙部分的分類，統(tǒng)計(jì)每一類的像素?cái)?shù)量即可得到孔隙度。采用灰度直方圖來(lái)統(tǒng)計(jì)圖像的灰度分布，當(dāng)直方圖呈現(xiàn)雙峰型[11]時(shí)，一般選取雙峰之間的峰谷所對(duì)應(yīng)的灰度值作為閾值，從而實(shí)現(xiàn)巖石固體成分和空隙部分的區(qū)分。圖4是碳酸鹽的灰度直方圖，通過(guò)分析該圖，可以發(fā)現(xiàn)對(duì)該樣本而言，應(yīng)該選擇灰度值85作為巖石固體成分和空隙部分的區(qū)分指標(biāo)，且所得到的孔隙度符合國(guó)家給定的碳酸鹽孔隙度范圍，側(cè)面驗(yàn)證了該處理方法的可行性。

圖4 碳酸鹽灰度直方圖

用psolid和ppore表示巖石固體成分像素?cái)?shù)和孔隙像素?cái)?shù)，對(duì)第n幅CT切片圖像I進(jìn)行分析，得到樣品的孔隙度φ，即

φ=ppore(I)/(ppore(I)+psolid(I))。

(8)

(9)

(10)

φi,j(GθG(ILR))x,y)2；

(11)

(12)

lporosity=(φ(ISR)-φ(IHR))2；

(13)

((ISR)x+1,y-(ISR)x,y)2)。

(14)

其中：λ和η為權(quán)重參數(shù)；r是放大倍數(shù)；W是圖像寬；H是圖像高；IHR表示高分辨率圖像；ISR表示超分辨率圖像；ILR表示低分辨率圖像；θG表示生成網(wǎng)絡(luò)參數(shù)；θD表示判別網(wǎng)絡(luò)參數(shù)；Wi,j,Hi,j分別表示VGG網(wǎng)絡(luò)中某一層特征圖的寬和高；φi,j是VGG網(wǎng)絡(luò)中某一層的計(jì)算結(jié)果。

3 實(shí)驗(yàn)結(jié)構(gòu)及分析

3.1 實(shí)驗(yàn)環(huán)境、數(shù)據(jù)集和評(píng)價(jià)指標(biāo)

本文算法在GPU計(jì)算型GN10X，8核40GB顯存的環(huán)境下，通過(guò)pytorch框架實(shí)現(xiàn)。數(shù)據(jù)集采用新南威爾士大學(xué)的數(shù)字巖心超分辨率數(shù)據(jù)集DRSRD1_2D中的碳酸鹽巖石圖像。該數(shù)據(jù)集提供了經(jīng)過(guò)處理之后的灰度數(shù)字巖石圖像，專門用于圖像超分辨率處理。

圖像超分辨率處理問(wèn)題主要研究放大倍數(shù)為2和4兩種情況，且放大倍數(shù)為2的情況在傳統(tǒng)神經(jīng)網(wǎng)絡(luò)上已經(jīng)有了很好的效果，因此本文主要研究放大倍數(shù)為4的情況。

隨機(jī)裁剪可以增加數(shù)據(jù)，還可以弱化數(shù)據(jù)噪聲和增加模型穩(wěn)定性，于是本文在數(shù)據(jù)集中對(duì)訓(xùn)練集中每張圖像進(jìn)行隨機(jī)裁剪，獲得224×224的圖片作為高分辨率圖像，通過(guò)對(duì)數(shù)據(jù)集中低分辨率圖片進(jìn)行相同操作得到與之對(duì)應(yīng)的56×56分辨率圖像。

在超分辨率圖像重建的研究中，恢復(fù)圖像質(zhì)量的評(píng)價(jià)指標(biāo)通常分為主觀評(píng)價(jià)和客觀評(píng)價(jià)。由于主觀評(píng)價(jià)通常容易受到參與人員個(gè)人因素的影響，本文采用常用客觀評(píng)價(jià)標(biāo)準(zhǔn)峰值信噪比(Peak Signal to Noise Ratio，PSNR)和結(jié)構(gòu)相似度(Structural Similarity Index，SSIM)進(jìn)行量化。

3.2 實(shí)驗(yàn)結(jié)果與分析

3.2.1 巖石孔隙度對(duì)結(jié)果的影響

本文把巖石孔隙度作為一種約束條件。為衡量其效果，實(shí)驗(yàn)通過(guò)對(duì)放大倍數(shù)為4的圖像進(jìn)行了客觀指標(biāo)的評(píng)價(jià)，加入lporosity的SRGAN與SRGAN在相同實(shí)驗(yàn)環(huán)境下PSNR和SSIM的對(duì)比如圖5所示。

圖5 SRGAN和加入lporosity的SRGAN對(duì)于PSNR和SSIM的對(duì)比

通過(guò)圖5可以發(fā)現(xiàn)，增加了損失函數(shù)lporosity作為約束條件可以在保證SSIM穩(wěn)定的前提下，使PSNR的波動(dòng)更加緩慢，減少了PSNR的突變情況，同時(shí)可以發(fā)現(xiàn)PSNR可以更快地趨于穩(wěn)定，加速了網(wǎng)絡(luò)的收斂過(guò)程。

通過(guò)對(duì)測(cè)試集上全體數(shù)據(jù)不同情況下孔隙度進(jìn)行分析可得：在加入lporosity作為約束條件之后，孔隙度較SRGAN平均上漲了0.5%左右，更加接近原圖的值。在超分辨處理之后能使其孔隙度與原始圖片相接近，減少了超分辨處理過(guò)程中隨機(jī)生成造成的影響和處理前后巖石物性特征的變化，對(duì)于成礦流體的研究有著極其重要的意義。

隨機(jī)采樣得到5張測(cè)試集圖片，對(duì)其不同情況下的孔隙度進(jìn)行具體展示，見(jiàn)表1。

表1 原始圖像、SRGAN和加入lporosity的SRGAN的孔隙度對(duì)比

3.2.2 自注意力機(jī)制對(duì)結(jié)果的影響

為衡量本文所加入的自注意力機(jī)制的效果，對(duì)放大倍數(shù)為4的圖像進(jìn)行了客觀指標(biāo)的評(píng)價(jià)，并與具有代表性的幾種SR方法在同一數(shù)據(jù)集上進(jìn)行了對(duì)比，結(jié)果見(jiàn)表2。

由表2可知，通過(guò)在SRGAN中增加注意力機(jī)制模塊，可以使其獲得更高的PSNR值，得到的圖像與原圖相比失真更小。

表2 測(cè)試集DRSRD1_2D中的對(duì)比實(shí)驗(yàn)結(jié)果

通過(guò)圖6中的SRGAN和本文模型重建后視覺(jué)效果的對(duì)比，可以發(fā)現(xiàn)，在增加了注意力機(jī)制之后，對(duì)于一些細(xì)小的突變位置，SRGAN往往會(huì)恢復(fù)得相對(duì)圓滑，而本文模型對(duì)于這些細(xì)小的突變則更加敏感，同時(shí)在圖片的整體感知方面，本文模型和SRGAN效果相近。因此本文在高頻特征的恢復(fù)上有一定的優(yōu)勢(shì)，恢復(fù)出來(lái)的圖像更加接近原始圖片。

圖6 各算法重建效果比較

3.2.3 3種算法的定量對(duì)比

為定量分析加入lporosity和自注意力機(jī)制對(duì)時(shí)間、空間的影響，取得了訓(xùn)練過(guò)程中一個(gè)epoch所用時(shí)間和空間大小，見(jiàn)表3。

表3 一輪訓(xùn)練中SRGAN、porosity-SRGAN和引入自注意力機(jī)制SRGAN在時(shí)間、空間和模型大小上的對(duì)比

通過(guò)分析可以發(fā)現(xiàn)，加入lporosity會(huì)增加其運(yùn)行時(shí)間用于孔隙度的計(jì)算，引入自注意力機(jī)制則會(huì)增加空間用于構(gòu)建自注意力映射表。自注意力映射表的大小由于受到圖像本身大小的影響，在訓(xùn)練和測(cè)試時(shí)會(huì)占用較大的顯存空間，但由于其在運(yùn)行中產(chǎn)生，并不影響模型大小。

4 結(jié) 論

本文針對(duì)SRGAN超分辨率網(wǎng)絡(luò)在巖石圖像一些細(xì)節(jié)的恢復(fù)上過(guò)于圓滑的問(wèn)題，引入了注意力通道機(jī)制，并通過(guò)在損失函數(shù)中引入巖石圖像所特有的孔隙度約束條件，實(shí)現(xiàn)了在加速GAN網(wǎng)絡(luò)訓(xùn)練過(guò)程的同時(shí)，進(jìn)一步提升超分辨率圖像的視覺(jué)效果。

在所采用的數(shù)據(jù)集上，本文所采用的方法得到的PSNR相比于SRGAN和EDSR分別提升了0.89 dB和0.63 dB。然而由于超分辨率重構(gòu)本質(zhì)上是一個(gè)病態(tài)問(wèn)題，找到合適的高低分辨率圖像映射關(guān)系極其困難，因此與原圖相比仍有提升的空間。由于注意力機(jī)制的引入導(dǎo)致模型訓(xùn)練和測(cè)試時(shí)會(huì)占用大量的顯存，同時(shí)也會(huì)消耗較多的時(shí)間，故本文下一步將著手研究在保障重構(gòu)質(zhì)量的前提下提升算法的訓(xùn)練、測(cè)試效率。