基于全局注意力機制的單像素成像圖像增強方法

劉 輝, 楊照華, 吳 云, 趙梓棟, 余遠金

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

0 引 言

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

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

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

1 單像素成像原理

1.1 單像素成像圖像增強算法

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

1.2 單像素成像流程

單像素成像包括光場調制、探測與重構2部分.單像素成像的成像原理如圖2所示,光線從光源發出經物體O反射到數字微鏡陣列(digital micromirror device,DMD)上,DMD預先加載的掩膜矩陣序列h會對入射光場進行空間光場調制,最終由桶探測器接收與掩膜矩陣序列一一對應的總光強序列值.第m次桶探測器探測的總光強值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為掩膜矩陣的序號,(x,y)為圖像坐標;式(2)中,n為采樣次數,β為采樣率;式(3)中,N是掩膜矩陣p行與q列的乘積,p、q為掩膜矩陣的維度,也為重構圖像的像素數,本文中p=q=64.

1.3 基于哈達瑪基的光場優化調制

在常見的單像素成像中,通常使用哈達瑪矩陣來構建掩膜矩陣序列,但使用自然序列的哈達瑪散斑在低采樣率下重建圖像會導致重構后的圖像出現重影,圖像的質量偏低.為了初步提高重構圖像的質量,可以結合排序方法如CC序、折紙序等序列進行圖像重構,本文根據空間探測實時性要求高和硬件資源有限的特點,選擇構造簡單且易于實現哈達瑪基的CC序進行光場優化調制.

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

(4)

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

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

1.4 基于差分鬼成像的圖像重構

在單像素成像的圖像重構過程中,通常使用二階關聯算法重構圖像.在空間探測應用中,由于光源功率的波動或光學探測器的不穩定性會對探測的總光場強度產生波動,為此,本文選用差分鬼成像算法如下:

(5)

式中,IDGI是重構出來的圖像,〈·〉表示對任意函數求系綜平均,r是參考光路的總光強值,S是物體光路的總光強值.

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

圖4 基于CC序的差分鬼成像實驗結果圖Fig.4 Differential ghost imaging experiment results based on CC sorting

2 基于SUNet的圖像增強

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

圖5 SUNet單像素圖像增強方案圖Fig.5 SUNet single-pixel image enhancement scheme

2.1 SUNet建模

SUNet的架構基于圖像分割模型,在此架構中應用了雙上采樣模塊來避免棋盤偽影,并且SUNet中自注意力操作無法共享內核權重,不同層中的特征使用了不同的內核值,所以相較于傳統的CNN架構在圖像增強中更為合理.

SUNet由3個模塊組成:

(1)淺層特征提取模塊

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

Fshallow=MSFE(IDGI)

(6)

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

(2)UNet特征提取模塊

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

Fdeep=MUFE(Fshallow)

(7)

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

(3)重建模塊

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

(8)

2.2 編碼器

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

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

2.3 編碼器

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

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

3 實驗結果與分析

3.1 數據集

本文使用STL-10數據集,訓練集有105 000張圖片,驗證集有8 000張圖片,從驗證集中選出少量的圖片做測試集.圖像的原始分辨率是96×96,通過雙線性插值將分辨率變為64×64.STL-10數據集一共有10個大類:飛機、鳥、汽車、貓、鹿、狗、馬、猴子、船和卡車,類別較多,有助于提高神經網絡模型的泛化性能.

3.2 模型評估指標

本文中采樣的評估標準為峰值信噪比(peak signal-to-noise ratio, PSNR)和結構相似度(structural similarity index measure, SSIM).2種指標的值越大,表明模型的恢復效果越好.MSE指標表示真值圖像與恢復圖像的均方誤差,PSNR計算公式如下:

(9)

(10)

(11)

3.3 圖像增強實驗

為驗證設計的算法有效性,在 Linux操作系統上,基于Python 環境下Pytorch深度學習工具包實現算法的仿真實驗,仿真計算機配置為Intel Core i7@4.00 GHz 處理器、16 G內存和2塊Nvidia GeForce Titan Xp顯卡.

本實驗采用基于SUNet模型的圖像增強網絡對單像素成像的圖像進行增強,此方法不需要物體的空間信息,更適合提高在外太空非合作目標超遠距離成像的圖像質量.實驗使用STL-10數據集,模型的輸入圖像大小為64×64.采樣率從0.05到0.5,步長為0.05,模型訓練次數為200次.初始學習率為0.000 2,使用Adam優化器自適應調整學習率[17],損失函數選擇均方誤差.訓練方式使用Pytorch的分布式訓練,可極大地提高模型的訓練速度,每個采樣率模型訓練時間為25 h.

限于篇幅,圖8給出了采樣率為0.3時,模型訓練的次數與PSNR和SSIM的曲線.由于神經網絡參數量巨大,在模型訓練初期,大部分參數的梯度下降方向不準確,所以模型的效果會出現突然下降又迅速回升的情況.到了模型訓練后期,大部分參數都已經收斂到了次優值,訓練的參數調節幅度小,所以曲線會變平滑.

圖8 采樣率為0.3時,模型PSNR與SSIM訓練結果圖Fig.8 At the sampling rate of 0.3, PSNR and SSIM vary with model training times

同時,ST使用滑動窗口注意力機制,相對于ViT大大減少了模型的參數量,便于模型訓練和誤差收斂,完成密集型的像素預測.表1給出了不同采樣率下不同模型在測試集上得到的PSNR.表2展示了不同采樣率下不同模型在測試集上得到的SSIM.對比模型分別為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;結構相似度從0.63提升到了0.71,提升了8%.

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

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

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

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

4 結 論

本文為了應對空間非局域目標探測問題,降低采樣率,提高成像速度,采用基于全局注意力機制的SUNet模型,在低采樣率下實現單像素成像圖像增強.提出的方法在單像素成像圖像恢復領域取得了較好的效果,證明ST的全局注意力機制相對于傳統的卷積神經網絡能夠進一步提高增強后圖像的質量.最后,與其他基于深度學習的方法相比,本文的方法對測試集中的圖像能表現出更好的圖像增強性能,更適用于空間非合作目標在低采樣率下的圖像感知與識別.最后,SUNet還可以擴展到超分辨圖像恢復領域,有望實現單像素成像的超分辨圖像重構.