导读

 

近年来，基于学习的方法表现出了良好的性能，并主导了主流研究方向。然而，现有的基于CNN的方法在捕获长程相关性和非局部自相似性方面存在局限性。本工作提出了一种能将光谱表征嵌入到Transformer的重建方法：CST。

本文介绍我们 ECCV 2022 关于 Snapshot Compressive Imaging 重建的工作：

《Coarse-to-Fine Sparse Transformer for Hyperspectral Image Reconstruction》

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

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

这个github仓库是一个针对 Snapshot Compressive Imaging 重建的工具包，集成了11种深度学习算法。

1. 简介

快照压缩成像系统如图1所示。左边为待成像的场景，也即三维光谱型号（空间维度长和宽，通道维度是不同波段的光谱）。它通过预先设计好的光路，首先被编码孔径掩膜进行调制，然后被三棱镜进行散射，在探测器上不同的空间位置进行成像，这些像叠加在一起之后便得到一个二维的快照估计图，如右下角 Measurement 所示。如此一来，原先输入的三维光谱数据（x，y，λ）便被压缩成了二维数据（x，y），从而降低数据存储和传输的负担。

图1 快照压缩成像系统

然而压缩估计图并不能直接应用，我们还需要将其复原成高光谱图像才可以进一步分析处理。通常来说，进行高光谱重建的方法可以分为四类：

（1）基于先验模型的传统方法。这一类方法会预先手工设计一些图像先验，如 total variation，low rank property 等。然而这类方法的泛化性差，并且每次使用前需要调整参数，花费大量时间。

（2）端到端的深度学习方法。这类方法直接采用一个深度学习模型，去拟合一个从 2D 快照压缩估计图到 3D 高光谱数据的映射。这类方法目前主要基于卷积神经网络，没有适配光谱表征在空间维度呈现出的稀疏性。并且卷积神经网络在捕获非局部依赖关系（non-local / long-range dependences）上有明显的短板。

（3）迭代式的深度学习方法。这一类方法运用迭代公式，将卷积神经网络嵌入到每一个迭代中作为去噪网络。

（4）即插即用的方法。这类方法将预先训练好的一个深度学习网络插入到每一个迭代始终，无需微调直接应用。

本文主要研究第（2）类方法，主要贡献点可以概括为如下：

• 提出了一种能将光谱表征嵌入到Transformer的重建方法。名字是 Coarse-to-Fine Sparse Transformer (CST)
• 提出一种检测密集光谱表征区域的方法，光谱感知筛选机制 ，Spectrum-Aware Screening Mechanism (SASM)
• 提出一种基于哈希来聚合相关光谱表征的多头自注意机制，Spectra-Aggregation Hashing Multi-head Self-Attention (SAH-MSA)。
• 在仿真数据集上，我们的 CST 系列模型用了更少的参数量取得了更高的结果。在真实数据上，效果更逼真。

2. 方法2.1 网络的整体结构

我们的 CST 的整体结构如图2（a）所示。首先将二维快照估计图 Y 滑动截取为初始的光谱图像 H 并与偏移的掩膜 M 进行 concate。网络的第一部分是一个稀疏度估计器，它估计出光谱表征比较密集的区域，然后输出一个

图2 CST的算法流程图

稀疏度图，第二阶段根据这个稀疏度图筛选出光谱密集区域，集中对这些区域进行计算。两个阶段均采用U-Net。

2.2 光谱感知的筛选机制

为了使稀疏度估计器能够实现我们想要的“筛选出光谱密集区域”的功能，我们对它输出的稀疏度图进行监督。我们设立的目标是重建光谱图与真值光谱图之间的差异沿通道的均值，如下式所示：

这很好理解，我们设想在高光谱成像的场景中，背景往往很黑暗，这些背景区域几乎没有信息，很好重建。那么难以重建的部分便是光谱密集的区域。基于此，我们设立的损失函数如下：

其中 MsM_sM_s 为估计的稀疏度图。那么整体的损失函数为：

其中 X' 为网络重建的光谱图，X* 为真值。λ 为权重系数，平衡两个损失函数。

在进行光谱密集区域筛选时，我们对稀疏度图设置一个阈值，高于这一阈值的区域就被选出。

2.3 基于哈希聚合相关光谱表征的多头自注意机制

CST 的第二阶段基于第一阶段输出的稀疏度图来进行光谱图像重建，其基本组成单元如图1（b）所示。有两条支路，顶端支路是一个恒等连接，底端支路首先基于稀疏度图的筛选，只对选出的区域进行计算，其组件中包含一个基于哈希聚合相关光谱表征的多头自注意机制（SAH-MSA），如图1（c）所示。以前的多头注意力机制（MSA）总是对所有的 token 进行计算，有些 token 毫无关联，这种计算方式效率低下。我们的 SAH-MSA 正是解决这一问题。首先，我们通过一个哈希映射函数，给每一个 token 算出一个分数，公式如下：

然后，我们根据 token 的分数从高到低排序，划分成不同的 buckets，如下：

则 SAH-MSA 的输出为：

其中，每一个 head 的自注意力由下式计算得到：

同时每一个 head 中的每一个 bucket 都由 multi-round 机制计算得到

其中每一个round的全总如下所示

3. 实验3.1 定量实验对比

在仿真数据上的定量实验的结果如下标所示

我们的 CST 系列采用更少的计算量达到了更高的性能。特别地，与我们 CVPR 2022的工作 MST 的比较如下图3所示。

图3 CST vs. MST3.2 定性实验对比

CST 与其他方法在仿真数据集和真实数据集上的定性结果对比分别如图4 和图5 所示。

图4 CST与其他方法在仿真数据集上的定性结果对比图5 CST与其他方法在真实数据集上的定性结果对比3.3 可视化分析

我们对稀疏度估计器输出的稀疏度图进行可视化，结果如图6所示。

图6 稀疏度图可视化与对比4. 回顾

ECCV 2022 的这次工作是 Transformer 应用于 Snapshot Compressive Imaging 的第二个工作。它采用与我们CVPR 2022 的工作 MST 不同的技术方案，实现了更低计算量，更高精度的高光谱重建效果。期待未来能看到更多的人从事 Snapshot Compressive Imaging 的研究。

另附上 MST 与 MST++ 的知乎链接：

phantom：[CVPR 2022 & NTIRE 冠军] 首个高光谱图像重建Transformer

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

*博客内容为网友个人发布，仅代表博主个人观点，如有侵权请联系工作人员删除。