基于全局注意力機(jī)制的單像素成像圖像增強(qiáng)方法

劉 輝, 楊照華, 吳 云, 趙梓棟, 余遠(yuǎn)金

1. 北京航空航天大學(xué),北京 100191 2. 北京控制工程研究所,北京 100094 3. 北京理工大學(xué),北京 100081

0 引 言

單像素成像已經(jīng)成為了量子光學(xué)領(lǐng)域的前沿技術(shù),在光譜成像、光學(xué)加密、3D成像、目標(biāo)跟蹤和空間探測(cè)領(lǐng)域上極具潛力[1-5].單像素成像的原理是利用空間光調(diào)制器調(diào)制散斑照射到物體上,同時(shí)使用桶探測(cè)器記錄物體反射的總光強(qiáng)值,通過(guò)關(guān)聯(lián)算法完成圖像重構(gòu).單像素成像與計(jì)算鬼成像的區(qū)別在于成像目標(biāo)和空間光調(diào)制器在光路中的先后順序,在很多論文中,由于成像原理與方法基本一致,所以不再區(qū)分單像素成像與鬼成像,本論文也是如此.為了提高重構(gòu)的圖像質(zhì)量,在低采樣率時(shí),俞文凱等[6]提出一種切蛋糕(cake-cutting,CC)序方法對(duì)空間光調(diào)制器生成的哈達(dá)瑪散斑進(jìn)行排序,確保重構(gòu)出的圖像無(wú)重疊的陰影部分.為了克服光學(xué)設(shè)備不穩(wěn)定帶來(lái)的光強(qiáng)漲落的影響,FERRI等[7]提出使用差分鬼成像的方法利用散斑和桶探測(cè)器值對(duì)圖像進(jìn)行重構(gòu),提高了圖像的重構(gòu)質(zhì)量.盡管使用上述方法對(duì)重構(gòu)圖像進(jìn)行了初步增強(qiáng),恢復(fù)出來(lái)的圖像質(zhì)量仍然不佳.

為了進(jìn)一步提高圖像的質(zhì)量,很多研究人員提出使用深度學(xué)習(xí)的方法進(jìn)行圖像增強(qiáng).LECUN等[8]提出LeNet神經(jīng)網(wǎng)絡(luò),首次將卷積神經(jīng)網(wǎng)絡(luò)使用在圖像分類上.KRIZHEVSKY等[9]提出AlexNet神經(jīng)網(wǎng)絡(luò),由于其優(yōu)異的性能,使卷積神經(jīng)網(wǎng)絡(luò)成為了圖像處理領(lǐng)域的主流算法.為了將卷積神經(jīng)網(wǎng)絡(luò)應(yīng)用到單像素成像處理中,WANG等[10]將單像素成像過(guò)程中生成的桶探測(cè)器序列值輸入到卷積神經(jīng)網(wǎng)絡(luò)中,實(shí)現(xiàn)了更高質(zhì)量的圖像重構(gòu),該方法在訓(xùn)練集上取得了很好的效果,但在測(cè)試集上表現(xiàn)欠佳,問(wèn)題在于網(wǎng)絡(luò)模型過(guò)擬合.RIZVI等[11]提出一種基于自編碼器網(wǎng)絡(luò)的重構(gòu)方案,該方法先使用CGI(computation ghost imaging)算法恢復(fù)出低質(zhì)量的圖像,然后使用編碼器將CGI算法恢復(fù)的圖像先編碼成隱層信息,再根據(jù)隱層信息解碼出原始圖像,編碼器與解碼器的網(wǎng)絡(luò)架構(gòu)解決了模型過(guò)擬合的問(wèn)題,但缺點(diǎn)是圖像重構(gòu)的效果取決于隱層信息的表達(dá),限制了模型的圖像增強(qiáng)性能.KARIM等[12]提出使用生成對(duì)抗網(wǎng)絡(luò)對(duì)單像素成像的圖像做后處理,此模型的骨干網(wǎng)絡(luò)仍然是卷積神經(jīng)網(wǎng)絡(luò),模型的性能受限于感受野的大小.

DEL等[13]提出Transformer神經(jīng)網(wǎng)絡(luò)模型,在自然語(yǔ)言處理(natural language processing, NLP)領(lǐng)域取得了優(yōu)異的成績(jī).受此啟發(fā),計(jì)算機(jī)視覺(jué)(computer vision,CV)領(lǐng)域的科研工作者們提出了ViT(vision transformer)模型[14]和ST(swin transformer)模型[15],解決了傳統(tǒng)卷積神經(jīng)網(wǎng)絡(luò)的平移不變性和無(wú)法獲得全局感受野的問(wèn)題.LIU等[16]提出了SUNet(swin transformer UNet),它是一個(gè)神經(jīng)網(wǎng)絡(luò)降噪模型,基于UNet架構(gòu),使用ST提取圖像的特征;使用塊合并層(patch merging)去除冗余的圖像特征,使用塊擴(kuò)展層(patch expanding)方法恢復(fù)圖像的細(xì)節(jié)信息.SUNet模型采用滑動(dòng)窗口注意力機(jī)制獲得全局的感受野,可進(jìn)一步提升模型的圖像去噪能力.本文將SUNet網(wǎng)絡(luò)引入到單像素成像領(lǐng)域中,采用STL-10圖片數(shù)據(jù)集,根據(jù)CC序?qū)_(dá)瑪散斑排序,使用差分鬼成像算法得到重構(gòu)圖像,確保重構(gòu)圖像具有較高的質(zhì)量,最后將該圖像送入神經(jīng)網(wǎng)絡(luò)模型中訓(xùn)練,并在測(cè)試集上取得了較好的結(jié)果.

1 單像素成像原理

1.1 單像素成像圖像增強(qiáng)算法

圖1 單像素成像圖像增強(qiáng)算法原理圖Fig.1 Schematic diagram of single-pixel imaging image enhancement algorithm

1.2 單像素成像流程

單像素成像包括光場(chǎng)調(diào)制、探測(cè)與重構(gòu)2部分.單像素成像的成像原理如圖2所示,光線從光源發(fā)出經(jīng)物體O反射到數(shù)字微鏡陣列(digital micromirror device,DMD)上,DMD預(yù)先加載的掩膜矩陣序列h會(huì)對(duì)入射光場(chǎng)進(jìn)行空間光場(chǎng)調(diào)制,最終由桶探測(cè)器接收與掩膜矩陣序列一一對(duì)應(yīng)的總光強(qiáng)序列值.第m次桶探測(cè)器探測(cè)的總光強(qiáng)值Sm可表示為

圖2 單像素成像原理圖Fig.2 Schematic of single-pixel imaging

(1)

n=Nβ

(2)

N=pq

(3)

式(1)中,h(x,y)表示掩膜矩陣(h∈Rp×q),O(x,y)表示真值圖像(O∈Rp×q),m=1,2,3,…,n,m為掩膜矩陣的序號(hào),(x,y)為圖像坐標(biāo);式(2)中,n為采樣次數(shù),β為采樣率;式(3)中,N是掩膜矩陣p行與q列的乘積,p、q為掩膜矩陣的維度,也為重構(gòu)圖像的像素?cái)?shù),本文中p=q=64.

1.3 基于哈達(dá)瑪基的光場(chǎng)優(yōu)化調(diào)制

在常見(jiàn)的單像素成像中,通常使用哈達(dá)瑪矩陣來(lái)構(gòu)建掩膜矩陣序列,但使用自然序列的哈達(dá)瑪散斑在低采樣率下重建圖像會(huì)導(dǎo)致重構(gòu)后的圖像出現(xiàn)重影,圖像的質(zhì)量偏低.為了初步提高重構(gòu)圖像的質(zhì)量,可以結(jié)合排序方法如CC序、折紙序等序列進(jìn)行圖像重構(gòu),本文根據(jù)空間探測(cè)實(shí)時(shí)性要求高和硬件資源有限的特點(diǎn),選擇構(gòu)造簡(jiǎn)單且易于實(shí)現(xiàn)哈達(dá)瑪基的CC序進(jìn)行光場(chǎng)優(yōu)化調(diào)制.

根據(jù)式(4)生成自然序列的哈達(dá)瑪矩陣

(4)

式中,?表示克羅內(nèi)克積.將原始哈達(dá)瑪矩陣H∈RN×N的每一行重塑成p行q列的二維掩膜矩陣,共計(jì)N個(gè)掩膜矩陣(使用h表示,每個(gè)h由H中的某一行變換得到),然后計(jì)算每個(gè)掩膜矩陣中連通域(圖像中聯(lián)通像素的集合)的個(gè)數(shù),根據(jù)連通域的個(gè)數(shù)給這N個(gè)掩膜矩陣從小到大排序,再根據(jù)采樣率從前向后選取掩膜矩陣.哈達(dá)瑪散斑排序原理如圖3所示,圖3(a)為16×16的原始哈達(dá)瑪散斑H,圖3(b)上方為原始序列哈達(dá)瑪散斑生成的掩膜矩陣,下方為經(jīng)過(guò)CC序后生成的掩膜矩陣.

圖3 16×16哈達(dá)瑪散斑CC序圖Fig.3 16×16 Hadamard speckle CC sorting map

1.4 基于差分鬼成像的圖像重構(gòu)

在單像素成像的圖像重構(gòu)過(guò)程中,通常使用二階關(guān)聯(lián)算法重構(gòu)圖像.在空間探測(cè)應(yīng)用中,由于光源功率的波動(dòng)或光學(xué)探測(cè)器的不穩(wěn)定性會(huì)對(duì)探測(cè)的總光場(chǎng)強(qiáng)度產(chǎn)生波動(dòng),為此,本文選用差分鬼成像算法如下:

(5)

式中,IDGI是重構(gòu)出來(lái)的圖像,〈·〉表示對(duì)任意函數(shù)求系綜平均,r是參考光路的總光強(qiáng)值,S是物體光路的總光強(qiáng)值.

差分鬼成像重構(gòu)結(jié)果如圖4所示,圖像大小為64×64,(a)為原始圖像,(b)為自然序列哈達(dá)瑪散斑的差分鬼成像重構(gòu)圖,(c)為CC序后哈達(dá)瑪散斑的重構(gòu)圖.對(duì)比圖4可知,差分鬼成像使用CC序后可以消除自然序列哈達(dá)瑪散斑的重影問(wèn)題,進(jìn)而提升圖像的質(zhì)量.

圖4 基于CC序的差分鬼成像實(shí)驗(yàn)結(jié)果圖Fig.4 Differential ghost imaging experiment results based on CC sorting

2 基于SUNet的圖像增強(qiáng)

SUNet的增強(qiáng)過(guò)程如圖5所示.包括淺層特征提取、編碼器、解碼器和重建模塊4部分,該模型包含10張?zhí)卣鲌D,使用標(biāo)號(hào)①、②、③、…、⑩表示,每張?zhí)卣鲌D的分辨率大小如圖5(b)所示.SUNet的主體為編碼器和解碼器,針對(duì)圖4(c)進(jìn)行圖像增強(qiáng).在編碼器部分,先對(duì)該圖經(jīng)過(guò)一個(gè)3×3的卷積操作得到特征圖,再經(jīng)過(guò)ST模塊提取特征和塊合并層下采樣后得到特征向量.在解碼器部分,特征向量經(jīng)過(guò)ST模塊恢復(fù)特征和塊擴(kuò)展層上采樣恢復(fù)到原始分辨率的圖像,再使用一個(gè)3×3的卷積操作得到增強(qiáng)后的圖像.在編碼器和解碼器網(wǎng)絡(luò)中,如果特征圖的分辨率相同,使用跨層連接可以獲得圖像的深層信息和淺層信息,增強(qiáng)模型的特征表達(dá)能力.最后計(jì)算增強(qiáng)后的圖像與原始圖像的均方誤差,并反向傳播梯度更新網(wǎng)絡(luò)參數(shù),數(shù)據(jù)集迭代200次后結(jié)束訓(xùn)練過(guò)程.

圖5 SUNet單像素圖像增強(qiáng)方案圖Fig.5 SUNet single-pixel image enhancement scheme

2.1 SUNet建模

SUNet的架構(gòu)基于圖像分割模型,在此架構(gòu)中應(yīng)用了雙上采樣模塊來(lái)避免棋盤(pán)偽影,并且SUNet中自注意力操作無(wú)法共享內(nèi)核權(quán)重,不同層中的特征使用了不同的內(nèi)核值,所以相較于傳統(tǒng)的CNN架構(gòu)在圖像增強(qiáng)中更為合理.

SUNet由3個(gè)模塊組成:

(1)淺層特征提取模塊

此模塊位于圖5(a)中淺層特征提取部分,由特征圖①、②組成.淺層特征提取是SUNet中的第一個(gè)模塊,用于獲取輸入圖像的低頻信息,如顏色或紋理.在本文中,輸入是重構(gòu)圖像IDGI(IDGI∈R1×1×64×64),64×64是重構(gòu)圖像的分辨率,使用3×3的卷積層MSFE(·)提取圖像的低頻信息,如下:

Fshallow=MSFE(IDGI)

(6)

式中,Fshallow(Fshallow∈R1×4×64×64)表示圖像的淺層特征.

(2)UNet特征提取模塊

此模塊由圖5(a)中編碼器和解碼器共同組成,包含特征圖②、③、…、⑨.UNet特征提取淺層特征Fshallow被輸入到UNet特征提取模塊MUFE(·)中提取高級(jí)別和多尺寸的深層特征,如下:

Fdeep=MUFE(Fshallow)

(7)

式中,Fdeep(Fdeep∈R1×4×64×64)是提取出的高級(jí)別和多尺寸的深層特征.

(3)重建模塊

此模塊位于圖5中重建模塊的紅色矩形框,由特征圖⑨、⑩組成.重建模塊負(fù)責(zé)從深層特征中恢復(fù)圖像,利用一個(gè)3×3的卷積層MR(·)來(lái)實(shí)現(xiàn),輸入是UNet特征提取部分獲得的深層特征,如下:

(8)

2.2 編碼器

編碼器由ST模塊和塊合并層組成.ST模塊將輸入圖像分成大小相同且互不重疊的塊,以塊為最小單位,再將塊聚合在一起形成窗口,分別在窗口內(nèi)和窗口間進(jìn)行注意力計(jì)算,使每個(gè)像素點(diǎn)都能獲得全局的感受野.塊合并層負(fù)責(zé)進(jìn)行下采樣操作,原理圖如圖6所示.塊合并層將圖像全部分成分辨率為2×2的塊,并在第三個(gè)維度中拼接,以此實(shí)現(xiàn)下采樣的功能,最后經(jīng)過(guò)3×3的卷積層進(jìn)行特征降維.

圖6 編碼器的塊合并層操作原理Fig.6 Operating principle of patch merge layer

2.3 編碼器

解碼器由ST模塊和塊擴(kuò)展層組成,負(fù)責(zé)特征恢復(fù)和上采樣.ST模塊的原理與解碼器相同.解碼器的塊擴(kuò)展層原理如圖7所示,它的目的是將圖像分辨率恢復(fù)到原始分辨率大小,每個(gè)塊擴(kuò)展層可以擴(kuò)充圖像至×2分辨率.上采樣由2部分組成:1)直接通過(guò)轉(zhuǎn)置卷積得到×2的特征圖;2)使用3×3的卷積層將圖像特征擴(kuò)充為初始的4倍,然后將這些特征在二維平面展開(kāi),實(shí)現(xiàn)×2的特征圖.將這2部分得到的×2的特征圖在特征維度拼接后送入3×3的卷積層后得到特征降維后的圖像.

圖7 解碼器的塊擴(kuò)展層操作原理Fig.7 Operating principle of patch expanding layer

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

3.1 數(shù)據(jù)集

本文使用STL-10數(shù)據(jù)集,訓(xùn)練集有105 000張圖片,驗(yàn)證集有8 000張圖片,從驗(yàn)證集中選出少量的圖片做測(cè)試集.圖像的原始分辨率是96×96,通過(guò)雙線性插值將分辨率變?yōu)?4×64.STL-10數(shù)據(jù)集一共有10個(gè)大類:飛機(jī)、鳥(niǎo)、汽車(chē)、貓、鹿、狗、馬、猴子、船和卡車(chē),類別較多,有助于提高神經(jīng)網(wǎng)絡(luò)模型的泛化性能.

3.2 模型評(píng)估指標(biāo)

本文中采樣的評(píng)估標(biāo)準(zhǔn)為峰值信噪比(peak signal-to-noise ratio, PSNR)和結(jié)構(gòu)相似度(structural similarity index measure, SSIM).2種指標(biāo)的值越大,表明模型的恢復(fù)效果越好.MSE指標(biāo)表示真值圖像與恢復(fù)圖像的均方誤差,PSNR計(jì)算公式如下:

(9)

(10)

(11)

3.3 圖像增強(qiáng)實(shí)驗(yàn)

為驗(yàn)證設(shè)計(jì)的算法有效性,在 Linux操作系統(tǒng)上,基于Python 環(huán)境下Pytorch深度學(xué)習(xí)工具包實(shí)現(xiàn)算法的仿真實(shí)驗(yàn),仿真計(jì)算機(jī)配置為Intel Core i7@4.00 GHz 處理器、16 G內(nèi)存和2塊Nvidia GeForce Titan Xp顯卡.

本實(shí)驗(yàn)采用基于SUNet模型的圖像增強(qiáng)網(wǎng)絡(luò)對(duì)單像素成像的圖像進(jìn)行增強(qiáng),此方法不需要物體的空間信息,更適合提高在外太空非合作目標(biāo)超遠(yuǎn)距離成像的圖像質(zhì)量.實(shí)驗(yàn)使用STL-10數(shù)據(jù)集,模型的輸入圖像大小為64×64.采樣率從0.05到0.5,步長(zhǎng)為0.05,模型訓(xùn)練次數(shù)為200次.初始學(xué)習(xí)率為0.000 2,使用Adam優(yōu)化器自適應(yīng)調(diào)整學(xué)習(xí)率[17],損失函數(shù)選擇均方誤差.訓(xùn)練方式使用Pytorch的分布式訓(xùn)練,可極大地提高模型的訓(xùn)練速度,每個(gè)采樣率模型訓(xùn)練時(shí)間為25 h.

限于篇幅,圖8給出了采樣率為0.3時(shí),模型訓(xùn)練的次數(shù)與PSNR和SSIM的曲線.由于神經(jīng)網(wǎng)絡(luò)參數(shù)量巨大,在模型訓(xùn)練初期,大部分參數(shù)的梯度下降方向不準(zhǔn)確,所以模型的效果會(huì)出現(xiàn)突然下降又迅速回升的情況.到了模型訓(xùn)練后期,大部分參數(shù)都已經(jīng)收斂到了次優(yōu)值,訓(xùn)練的參數(shù)調(diào)節(jié)幅度小,所以曲線會(huì)變平滑.

圖8 采樣率為0.3時(shí),模型PSNR與SSIM訓(xùn)練結(jié)果圖Fig.8 At the sampling rate of 0.3, PSNR and SSIM vary with model training times

同時(shí),ST使用滑動(dòng)窗口注意力機(jī)制,相對(duì)于ViT大大減少了模型的參數(shù)量,便于模型訓(xùn)練和誤差收斂,完成密集型的像素預(yù)測(cè).表1給出了不同采樣率下不同模型在測(cè)試集上得到的PSNR.表2展示了不同采樣率下不同模型在測(cè)試集上得到的SSIM.對(duì)比模型分別為2010年提出的DGI(differential ghost imaging)模型[18]、2014年提出的稀疏(Sparse)模型[19]、2018年提出的AP(alternating projection)模型[20]、2022年提出的GIDC(ghost imaging using deep neural network constraint)模型[21]和本文提出的SUNet模型.圖9為不同采樣率下的PSNR曲線與SSIM曲線.本文提出的SUNet與2022年提出的GIDC方法相比,在0.1的采樣率下,峰值信噪比從20.05 dB提生到23.34 dB,提升了3.29 dB;結(jié)構(gòu)相似度從0.63提升到了0.71,提升了8%.

表1 不同采樣率下各模型的PSNR值對(duì)比Tab.1 Comparison of PSNR of different models under different sampling rates

表2 不同采樣率下,各模型的SSIM值對(duì)比Tab.2 Comparison of SSIM of different models under different sampling rates

圖9 低采樣率下性能指標(biāo)折線圖Fig.9 Performance metrics line chart under low sampling rates

本文選擇部分測(cè)試集的圖像通過(guò)不同的深度學(xué)習(xí)方法對(duì)圖像進(jìn)行增強(qiáng),如圖10所示.最左邊的一列圖像表示原始圖像.其余各列分別為DGI、AP、稀疏、GIDC模型和本文提出的SUNet模型在25%采樣率下增強(qiáng)后的圖像,其性能指標(biāo)PSNR和SSIM顯示在圖像下方.本文提出的SUNet方法恢復(fù)的圖像質(zhì)量最佳,通過(guò)CC序后的哈達(dá)瑪散斑與桶探測(cè)器值使用差分鬼成像算法恢復(fù)出低質(zhì)量的圖像,然后使用深度學(xué)習(xí)方法對(duì)圖像進(jìn)行增強(qiáng),并使用均方誤差反向傳播梯度更新模型的參數(shù),最后在測(cè)試集中采樣評(píng)估模型的性能.與GIDC方法相比,在0.1的采樣率下,峰值信噪比提升了3.29dB,結(jié)構(gòu)相似度提升了8%,表明ST的性能優(yōu)于傳統(tǒng)CNN,全局感受野可以有效提高圖像增強(qiáng)的質(zhì)量,實(shí)現(xiàn)低采樣率條件下較好的圖像增強(qiáng).

4 結(jié) 論

本文為了應(yīng)對(duì)空間非局域目標(biāo)探測(cè)問(wèn)題,降低采樣率,提高成像速度,采用基于全局注意力機(jī)制的SUNet模型,在低采樣率下實(shí)現(xiàn)單像素成像圖像增強(qiáng).提出的方法在單像素成像圖像恢復(fù)領(lǐng)域取得了較好的效果,證明ST的全局注意力機(jī)制相對(duì)于傳統(tǒng)的卷積神經(jīng)網(wǎng)絡(luò)能夠進(jìn)一步提高增強(qiáng)后圖像的質(zhì)量.最后,與其他基于深度學(xué)習(xí)的方法相比,本文的方法對(duì)測(cè)試集中的圖像能表現(xiàn)出更好的圖像增強(qiáng)性能,更適用于空間非合作目標(biāo)在低采樣率下的圖像感知與識(shí)別.最后,SUNet還可以擴(kuò)展到超分辨圖像恢復(fù)領(lǐng)域,有望實(shí)現(xiàn)單像素成像的超分辨圖像重構(gòu).