來源丨h(huán)ttps://zhuanlan.zhihu.com/p/576280023編輯丨極市平臺

本文介紹我們 NeurIPS 2022 關于 Spectral Compressive Imaging （SCI）重建的工作：

《Degradation-Aware Unfolding Half-Shuffle Transformer for Spectral Compressive Imaging》

文章：https://arxiv.org/abs/2205.10102

代碼：https://github.com/caiyuanhao1998/MST

這個github倉庫是一個針對 Snapshot Compressive Imaging 重建的工具包，集成了超過12種深度學習算法。我們之前的工作如 MST, CST, MST++, HDNet 也都在這個github倉庫中開源。本文也是我們做的 Transformer in SCI 系列的第三個工作。

MST：https://arxiv.org/abs/2111.07910

CST：https://arxiv.org/abs/2203.04845

MST++：https://arxiv.org/abs/2204.07908

HDNet：https://arxiv.org/abs/2203.02149

單曝光快照壓縮成像（Snapshot Compressive Imaging，SCI）的任務是將一個三維的數(shù)據(jù)立方塊如視頻（H×W×T）或高光譜圖像（H×W×λ）通過預先設計好的光學系統(tǒng)壓縮成一個二維的快照估計圖（H×W）從而大幅度地降低數(shù)據(jù)存儲和傳輸?shù)拈_銷。常見的單曝光快照壓縮成像系統(tǒng)有 Coded Aperture Snapshot Spectral Compressive Imaging （CASSI），如下圖所示

那么在 SCI 中一個至關重要的問題就是如何從被壓縮過后的二維快照估計圖重建出原始的三維數(shù)據(jù)，當前主流的方法大都基于深度學習，可以分為兩類：端到端（End-to-end）的方法和深度展開式（Deep Unfolding）的方法。端到端的方法直接采用一個深度學習模型，去擬合一個從 2D 快照壓縮估計圖到 3D 高光譜數(shù)據(jù)的映射。這種方法比較暴力，確實可解釋性。深度展開式方法將神經(jīng)網(wǎng)絡嵌入到最大后驗概率（Maximum A Posteriori，MAP）模型中來迭代地重建出高光譜圖像，能更好地和光學硬件系統(tǒng)適配。因此，本文主要研究深度展開式算法。當前這些方法主要有兩大問題：

• 當前的深度展開式框架大都沒有從 CASSI 中估計出信息參數(shù)用于引導后續(xù)的迭代，而是直接簡單地將這些所需要的參數(shù)設置為常數(shù)或者可學習參數(shù)。這就導致后續(xù)的迭代學習缺乏蘊含 CASSI 退化模式和病態(tài)度信息指導。
• 當前的 Transformer 中全局的 Transformer 計算復雜度與輸入的圖像尺寸的平方成正比，導致其計算開銷非常大。而局部 Transformer 的感受野又受限于位置固定的小窗口當中，一些高度相關的 token 之間無法match。

為了解決上述兩個問題，我們提出了首個深度展開式的Transformer。我們貢獻可以概括為：

• 首先，我們推導出了一個能夠感知 CASSI 退化模式與病態(tài)度的深度展開框架，它從壓縮估計圖和編碼掩膜中估計出信息參數(shù)來引導后續(xù)的迭代學習。
• 接著，我們設計了一個能夠同時捕獲局部和全局依賴關系的 Transformer 并且計算復雜度相較于全局的Transformer而言，大幅減低。
• 最終，我們將我們設計的 Transformer 嵌入到我們推導的深度展開框架中來極大提升光譜圖像重建的效果。我們的算法在使用更低參數(shù)量和更少計算量的前提之下，性能大幅度地超過了前人的方法。

我們定義向量化后的壓縮估計圖為 y , 被偏移后的輸入數(shù)據(jù)為 x， 傳感矩陣為 φ ，則 CASSI 的退化數(shù)學模型為

其中的 n 表示成像時產(chǎn)生的隨機噪聲，同樣地，它也經(jīng)過了向量化。

我們首先推導出一個 CASSI 退化模式和病態(tài)度可感知的深度展開框架，Degradation-Aware Unfolding Framework （DAUF），如圖 2 所示。它以最大后驗概率為理論基礎來進行推導。結合公式（1），我們可以得到 CASSI 的最大后驗概率能量優(yōu)化函數(shù)為：

引入輔助變量 z 之后，我們可以得到

為了得到展開式的推導，同時使迭代過程更加簡單，能夠更快地收斂，我們對公式（3）采用 Half-Quadratic Splitting （HQS）算法進行展開，得到：

我們對公式（4）中的 x 和 z 進行解耦，從而得到兩個迭代的子問題如下：

其中的 x 項有一個閉式解：

其中 I 是恒等矩陣，上述閉式解涉及到矩陣求逆，對計算機不友好。為簡化矩陣求逆運算，我們做了以下推導：

將公式（7）插入到公式（6），我們可以得到：

請注意, 在 CASSI 系統(tǒng)中  是一個對角矩陣, 定義 , 由此可得:

定義  且  表示  的第  個元素, 將公式 (9) 代入公式 (8), 可得：

返回到迭代公式（5）中，從貝葉斯概率的角度來看，z 項的求解我們可以視為一個去噪問題：

我們可以用一個嵌入的神經(jīng)網(wǎng)絡來隱式地求解公式（11）。至此，我們便完成了 x 項與 z 項的迭代求解過程。

總的來看，我們的 DAUF 可以用下面這個范式來概括：

從左到右依次為參數(shù)估計，線性映射，先驗去噪。

我們半交互式 Transformer （Half-Shuffle Transformer，HST）的整體結構如圖3 (a) 所示，采用一個 U 形網(wǎng)絡，包含 一個 Encoder，Bottleneck，Decoder。其中基本單元是 Half-Shuffle Attention Block （HSAB）。HSAB中最重要的模塊是 Half-Shuffle Multi-head Self-Attention （HS-MSA）。

HS-MSA 如圖3 (d) 所示，將輸入的 token  線性映射為：

然后按通道進行二等分：

等分后分別輸入到兩個分枝，local 分枝和 non-local 分枝，其中的 local 分枝計算 self-attention 如下：

non-local 分枝首先對  進行網(wǎng)格劃分, 再reshape, 從  到 , 然后再計算 self-attention 如下:

對兩個branch的計算結果進行融合：

最后將整個 HST 作為一個 stage 的去噪先驗網(wǎng)絡（不同 stage 的網(wǎng)絡權重不共享）插入到 DAUF 當中便得到我們的 DAUHST。

定量實驗對比如表 1 所示，我們的 DAUHST 以更低的計算量和參數(shù)量顯著超越了之前 16 種 state-of-the-art 方法。我們的方法比先前最好的方法 End-to-end 方法 CST-L 和 Deep Unfolding 方法 BIRNAT 要分別高出 2.24 和 0.78 dB。

相較于先前的 Deep Unfolding 方法，我們繪制了 PSNR - FLOPS 坐標圖比較 DAUHST 和其他 Deep Unfolding 方法的 性價比。如圖4所示。我們的方法在消耗相同計算量的情況下比先前方法要高出 4 dB。

在仿真數(shù)據(jù)集上的定性結果對比如圖5所示。左上角是RGB圖像和快照估計圖（Measurement）。下方四行圖像是不同方法重建的四個波長下的高光譜圖像。右上角的圖像是下方圖像中黃色框框內的放大圖。從重建的高光譜圖像來看，我們的方法能更好地恢復出細節(jié)內容和紋理結構，請注意對比小立方塊區(qū)域。a 和 b 曲線對應著 RGB 圖像的兩個綠色框的區(qū)域的光譜強度曲線，可以看出，我們的 DAUHST 與 Ground Truth 的曲線最為接近。

圖6 展示的是各類方法在真實數(shù)據(jù)集上的對比?？梢钥闯鲋挥形覀兊姆椒軌蛟诟鞣N波長的光譜上穩(wěn)定地重建出小花并同時抑制噪聲的生成。

本文是我們 SCI 系列代表作的第五個，也是 NeurIPS 上邊首次有 SCI 重建的工作。SCI 重建作為新興的 low-level 方法這兩年迅猛發(fā)展，希望能夠看到有更多的人能夠加入的這個 topic 的研究，畢竟新的領域有更多出成果的機會。另附上我們先前在 CVPR 2022 和 ECCV 2022 上的兩個工作 MST 和 CST 的知乎解讀鏈接：

https://zhuanlan.zhihu.com/p/501101943

https://zhuanlan.zhihu.com/p/544979161

本文僅做學術分享，如有侵權，請聯(lián)系刪文。