MST++:MultistageSpectralwiseTransformerforEfficientSpectralReconstruction

摘要&＃xff1a;现有的光谱重建(SR)主要方法集中在设计更深或更宽的卷积神经网络(CNN ),以学习从 RGB 图像到其高光谱图像(HSI)的端到端映射。这些基于 CNN 的方法实现了令人印象深刻的恢复性能&＃xff0c;同时在捕获long-range dependencies 和self-similarity方面显示出局限性。为了解决这个问题&＃xff0c;我们提出了一种新的基于Transformer的方法&＃xff0c;Multi-stage Spectral-wise Transformer(MST&＃43;&＃43;)&＃xff0c;用于高效光谱重建。特别地&＃xff0c;我们采用基于 HSI 空间稀疏和频谱自相似特性的Spectral-wise Multi-head Self-attention(S-MSA)来组成基本单元&＃xff0c;Spectral-wise Attention Block-频谱式注意块(SAB)。然后 SABs 建立单级谱式变换器(SST ),利用 U 形结构提取多分辨率上下文信息。最后&＃xff0c;我们的 MST&＃43;&＃43;由几个SST 级联而成&＃xff0c;从粗到细逐步提高了重建质量。综合实验表明&＃xff0c;我们的 MST&＃43;&＃43;明显优于其他先进的方法。在 NTIRE 2022 光谱重建挑战赛中&＃xff0c;我们的方法获得了第一名。

一、介绍

         高光谱成像在狭窄的波段&＃xff08;narrow bands&＃xff09;中记录真实世界的场景光谱&＃xff0c;其中每个波段捕获特定光谱波长的信息。与普通的 RGB图像相比&＃xff0c;HSIs 具有更多的光谱波段来存储更丰富的信息&＃xff0c;并描绘所捕捉场景的更多细节。由于这一优点 &＃xff0c; HSI 具有广泛的应用 &＃xff0c;例如医学图像处理&＃xff0c;遥感&＃xff0c;目标跟踪等等。尽管如此&＃xff0c;这种具有丰富光谱信息的 HSI 是费时的&＃xff0c;以至于光谱仪被用来沿着空间或光谱维度扫描场景。这种局限性限制了 HSIs 的应用范围&＃xff0c;尤其是在动态或实时场景中。

        解决这个问题的一个方法是开发快照压缩成像(SCI)系 统和计算重建算法从 2D 测量到三维 HSI 立方体。然而&＃xff0c;这些方法依赖于昂贵的硬件设备。为了降低成本&＃xff0c;提出了从给定的 RGB 图像重建 HSI 的光谱重建(SR)算法&＃xff0c;该图像可以容易地由 RGB 相机获得。

        传统的 SR 方法主要基于稀疏编码或相对较浅的学习模型。尽管如此&＃xff0c;这些基于模型的方法表现能力有限&＃xff0c;泛化能力差。近年来&＃xff0c;随着深度学习的发展&＃xff0c;人工智能取得了重大进展。深度卷积神经网络(CNN)已经被用于学习从 RGB 图像到 HSI 立方体的端到端映射函数。尽管已经取得了令人印象深刻的性能&＃xff0c;但是这些基于 CNN 的方法在捕获远程相关性(long-range dependencies)和光谱间自相似性方面显示出局限性。

        近年来&＃xff0c;自然语言处理(NLP)模型&＃xff0c;Transformer &＃xff0c;已在计算机视觉中得到应用&＃xff0c;并取得巨大成功。Transformer 中的multi-head self-attention(MSA)机制在建模长期依赖关系&＃xff08;modeling long-range dependencies&＃xff09;和无局部自相似性&＃xff08;no-local self-similarity&＃xff09;方面比 CNN 做得更好&＃xff0c;这可以缓解基于 CNN 的 SR 算法的局限性。

        然而&＃xff0c;直接使用标准Transformer的SR 会遇到两个主要问题。㈠Global和local Transformer记录空间区域的相互作用。然而&＃xff0c;HSI 表示在空间上是稀疏的&＃xff0c;而在频谱上是高度自相似的。因此&＃xff0c;对空间相互依赖性进行建模可能比捕获光谱间相关性的成本效益低。(ii)一方面&＃xff0c;standard globel MSA 的计算复杂度与空间维度成二次关系&＃xff0c;这是一个难以承受的巨大负担。另一方面&＃xff0c;基于局部窗口&＃xff08;local window-based&＃xff09;的 MSA 在位置特异性窗口&＃xff08;position-specific windows&＃xff09;内受到有限 receptive fields的影响。

         为了解决上述限制&＃xff0c;我们提出了第一个Transfoemer-based的框架&＃xff0c;Multi-stage Spectral-wise(MST&＃43;&＃43; ),用于从 RGB 图像进行有效的光谱重建。值得注意的是&＃xff0c;我们的 MST&＃43;&＃43;是基于先前的工作 MST&＃xff0c;它是为光谱压缩成像复原而定制的。首先&＃xff0c;我们注意到 HSI 信号在空间上是稀疏的&＃xff0c;而在频谱上是自相似的。基于这一性质&＃xff0c;我们采用Spectral-wise Multi-head Self-Attention(S-MSA)来构成基本单元——Spectral-wise Attention Block-谱式注意块(SAB)。S-MSA 将每个光谱特征图作为一个token&＃xff0c;沿光谱维度计算self-attention。其次&＃xff0c;我们的 SABs建立了我们提出的单级光谱-Single-stage Spectral-wise(SST ),它利用 U 形结构来提取多分辨率谱上下文信息&＃xff0c;这对于 HSI 恢复是至关重要的。最后&＃xff0c;我们的 MST&＃43;&＃43;由几个 SST 级联&＃xff0c;开发了一个多阶段学习方案&＃xff0c;从粗到细逐步提高重建质量&＃xff0c;从而显著提高了性能。

 本文的主要贡献如下。
        我们为SR提出了一个新颖的框架 MST&＃43;&＃43;.据我们所知&＃xff0c;这是在这项任务中首次尝试探索 Transformer 的潜力。
        我们在这个任务上验证了一系列自然图像恢复模型。针对这些问题&＃xff0c;我们提出了 Top- K 多模型集成策略来提高分辨率性能。
        定量和定性实验表明&＃xff0c;我们的 MST&＃43;&＃43;明显优于SOTA 方法&＃xff0c;同时需要更cheaper的参数和触发器。令人惊讶的是&＃xff0c;我们的 MST&＃43;&＃43;获得了 NTIRE 2022 光谱重建挑战赛第一名。

2. 相关工作

2.1. 高光谱图像采集

        用于收集 HSI 的传统成像系统通常采用光谱仪沿着空间或光谱维度扫描场景。三种主要类型的扫描仪&＃xff0c;包括语音室扫描仪、推进式扫描仪和带式顺序扫描仪&＃xff0c;通常用于采集 HSI。几十年来&＃xff0c;这些扫描仪已经广泛用于探测、遥感、医学成像和环境监测。例如&＃xff0c;推扫式扫描仪和 whiskbroom 扫描仪已经用于卫星传感器&＃xff0c;用于摄影测量和遥感。然而&＃xff0c;扫描过程通常需要很长时间&＃xff0c;这使得它不适合测量动态场景。此外&＃xff0c;成像设备通常体积太大&＃xff0c;无法插入便携式平台。为了解决这些限制&＃xff0c;研究人员开发了 SCI 系统来捕获 HSI&＃xff0c;其中 3D HSI 立方体被压缩成单个 2D 测量值.在这些 SCI 系统中&＃xff0c;编码孔径快照光谱成像(CASSI) 脱颖而出&＃xff0c;形成一个很有前途的研究方向。尽管如此&＃xff0c;到目前为止&＃xff0c;SCI 系统对于消费级应用来说仍然过于昂贵。即使是“低成本”的 SCI 系统通常也 在 10 万美元 10K 左右。因此&＃xff0c;随机共振课题具有重要的研究和实用价值。        

        NLP 模型转换器被提议用于机器翻译。近年来&＃xff0c;由于它在捕捉空间区域之间的长程相关性方面的优势&＃xff0c;它已被引入计算机视觉并受到广泛欢迎。在高层视 觉 中 &＃xff0c;变 换 在 图 像 分 类 中 得 到 了 广 泛 的 应 用 &＃xff0c;对象 检测&＃xff0c;语义分割&＃xff0c;人体姿态估计等。此外&＃xff0c;vision Transformer 还 已在低级视觉中使用.例如&＃xff0c;蔡等人提出第一个基于变压器的端到端框架 MST&＃xff0c;用于根据压缩测量值重建 HSI。林等将 HSI 稀疏性嵌入到 Transformer 中&＃xff0c;以建立用于光谱压缩成像的由粗到细的学习方案。前工作人员采用由 Swin Transformer 块进行自然图像恢复。然而&＃xff0c;据我们所知。Transformer 在光谱超分辨率方面的潜力还 没有被开发出来。这项工作旨在填补这一研究空白。

3. 方法

3.1. 网络体系结构 

        如图所示&＃xff0c;图(a)描述了建议的多级频谱式变压器 (MST&＃43;&＃43;)&＃xff0c;它由 Ns 个单级频谱式变压器(SSTs)级联而成。我们的 MST&＃43;&＃43;接受一个 RGB 图像作为输入&＃xff0c;并重建它的 HSI 副本。利用长身份映射来简化训练过程。

         图(b)显示了由编码器、瓶颈和解码器组成的 U 形SST。嵌入和映射模块是单个 conv3×3 层。编码器中的特征映射依次经历下采样操作 (一个步长 &＃xff0c;conv4×4 层)、N1 频谱关注块(sab)、下采样操作和 N2 SABs。瓶颈是由 N3 SABs。解码器采用对称架构。上采样操作是一个步长的 deconv2×2 层。为了避免下采样中的信息丢失&＃xff0c;在编码器和解码器之间使用跳跃连接。

        图(c)示出了 SAB 的组件&＃xff0c;即前馈网络(FFN)。

        图(d)展示了基于频谱的多头自关注(S-MSA)和两层规范化。 

        图(e)显示了S-MSA 的详细信息。

3.4. 集成策略

        NLP 模型转换器被提议用于机器翻译。在 NTIRE 2022 光谱重建挑战赛中&＃xff0c;我们采用了三种集成策略&＃xff0c;包括自集成、多尺度集成和Top-K 多模型集成&＃xff0c;以提高 MST&＃43;&＃43;的性能和通用性。在这一部分,我们将详细描述它们。

3.4.1 自我集成

        RGB 输入被 上 / 下 / 左 / 右 翻 转 或 旋 转 90°/180°/270°&＃xff0c;以加入网络。随后&＃xff0c;输出被转换到原始状态进行平均。

3.4.2 多尺度集成

         我们分别用大小为256×256×128 和 64×64的图片训练我们的模型 。然后对输出(整个图像)进行平 均以提高恢复质量。

3.4.3 Top-K 多模型集成

4. 实验

4.1. 资料组
        NTIRE 2022 光谱重建挑战提供的数据集包含 1000个 RGB-HSI 对。该数据集按 18:1:1 的比例分成训练、有效和测试子集。大小为 482×512 的每个 HSI 具有从400 纳米到 700 纳米的 31 个波长。为了产生相关的对应于 RGB 对应物 I ∈ RH×W×3&＃xff0c;变换矩阵 M ∈RNλ×3 被 应 用 于 地 面 真 实 HSI 立 方 体 Y ∈ RH×W×Nλ as I &＃61; Y M.  然后&＃xff0c;生成的 RGB 图像被注入散粒噪声以模拟真实的摄像机情况。

4.2.实施细节

        在 训 练 过 程 中 &＃xff0c; RGB 图像被线性重新缩放到 [0&＃xff0c;1]&＃xff0c;之后从数据集中裁剪出 128×128的RGB 和 HSI 样本对。batch size设置为 20&＃xff0c;参数优化算法选择 Adam&＃xff0c;其中β1 &＃61;0.9&＃xff0c;β2 &＃61; 0.999。学习率被初始化为 0.0004&＃xff0c;采用Cosine Annealing scheme训练300轮。通过随机旋转和翻转来增加训练数据。提出的 MST&＃43;&＃43;已经在 Pytorch 框架上实现&＃xff0c;在单个 RTX 3090 GPU 上训练一个网络大约需要 48 小时。采用预测的和地面真实的 HSI 之间的MRAE 损失函数作为目标。在我们 MST&＃43;&＃43;的实现中&＃xff0c;我们设置 Ns &＃61; 3&＃xff0c;N1 &＃61; N2 &＃61; N3 &＃61; 1&＃xff0c;C &＃61; 31。
        在测试阶段&＃xff0c;RGB 图像也是线性的重新调整到[0&＃xff0c;1]并输入网络以实现频谱恢复。我们
的 MST&＃43;&＃43; 在 RTX 3090 GPU 上 重 建 每 幅 图 像 (482 512 3 大小)需要 102.48 毫秒。
        我们采用三个评估指标来评估模型性能。第一个度量是平均相对绝对误差(MRAE ),其计算重建的和地面 真实 HSI 的所有波长之间的像素差异。

4.3. 主要结果

4.3.1 有效集上的定量结果

        我们比较了我们的 MST&＃43;&＃43;和 SOTA 方法&＃xff0c;包括两种 SCI重建方法(MST和HD-Net)、三种 SR 算法 (HSCNN&＃43; &＃xff0c;AWAN 和HRNet)和五种自 然 图 像 恢 复 模 型 (MIRNet &＃xff0c; MPRNet &＃xff0c;Restormer &＃xff0c;HINet &＃xff0c;EDSR)在有效集合上。请注意&＃xff0c;HSCNN&＃43; &＃xff0c;AWAN 和HRNet是 NTIRE 2018 的获奖者和 2020 年光谱 重建的挑战。结果列在表2。

 

        我们的 MST&＃43;&＃43;在需要最少的参数和浮点运算的同时&＃xff0c;大大优于 SOTA 方法。例如&＃xff0c;我们的 MST&＃43;&＃43;在 PSNR 实现了 3.10、7.43 和7.96 dB 的 改 进 &＃xff0c; 而 只 需 要 40.10% (1.62 /4.04) 、 5.11% 、 34.84% 的 参 数 和 8.52% (23.05 / 270.61)&＃xff0c;与 AWAN、HRNet 和 HSCNN&＃43;相比&＃xff0c;失败次数分别为 14.07%和 7.57%。

        为了直观地展示 MST&＃43;&＃43;的优越性&＃xff0c;我们在图中提供了不同算法的 PSNR-Params-FLOPS比较。 Figure1纵轴是PSNR(性能)&＃xff0c;横轴是 FLOPS(计算成本)&＃xff0c;圆半径是Params(内存成本)。可以看出&＃xff0c;我们的 MST&＃43;&＃43;占据了左上角&＃xff0c;展示了我们的方法在效率上的巨大优势.

 

4.3.2 测试集上的定量结果

表2列出了 NTIRE 2022 光谱挑战(测试集)的前 12 名&＃xff0c;其中*表示使用集合模型。令人印象深刻的是&＃xff0c;我们 的方法在 231 名参与者中赢得了冠军&＃xff0c;这表明了我们的 MST&＃43;&＃43;的优越性。

4.4. 消融研究

        我们使用有效的子集来进行消融。基础线性模型是通过从MST&＃43;&＃43;中移除 S-MSA 而得到的。

4.4.1 Self-Attention机制

        我们已经在第二节讨论了不同的自我注意机制。在本部分中&＃xff0c;我们进行消融研究以验证这些 MSA 的性能&＃xff0c;包括global MSA (G-MSA) &＃xff0c;local window-based MSA (W-MSA) &＃xff0c;Swin MSA (SW-MSA) &＃xff0c;以及通过的S-MSA.结果记录在表3a。

        

        为了公平起见&＃xff0c;使用不同 MSA 的模型的参数被设置为相同的值。值得注意的是&＃xff0c;G-MSA 的输入特性被缩减为 1 的大小&＃xff0c;以避免内存不足。可以观察到&＃xff0c;我们采用的 SMSA 实现了最显著的改进&＃xff0c;同时需要最少的存储器和计算成本。具体来说&＃xff0c;当我们分别应用 SW-MSA、WMSA、G-MSA 和 S-MSA 时&＃xff0c;性能在 MRAE 分别提高了0.0338、0.0553、0.1356 和 0.1532&＃xff0c;同时增加6.42、6.42、7.43和5.37 GFLOPS。正如 Sec. 3.3&＃xff0c;这些结果主要来源于HSI在空间上稀疏&＃xff0c;而在光谱上自相似的性质。因此&＃xff0c;捕获光谱间的相关性比建模空间区域的相关性更具成本效益。

4.4.2 Stage Number

        我们改变 MST&＃43;&＃43;的级数 Ns 来研究它的效果。结果显示在表3b。当 Ns &＃61; 3 时&＃xff0c;性能达到峰值。因此&＃xff0c;我们最终采用 3 阶段 MST&＃43;&＃43;作为我们的 SR 模型。

 

4.4.3 集成策略

在 Sec.3.4&＃xff0c;我们采用三种集成策略应对 NTIRE2022 光谱重建挑战。在这一部分中&＃xff0c;我们进行消融来研究它们的效果。在有效集上&＃xff0c;自集成、多尺度集成和 Top-K (K 设置为 5)多模型集成在 MRAE 方面分别实现了 0.015、0.033 和 0.045 的改进。

5. 未来的工作

        到目前为止&＃xff0c;还没有一个用于 SR 研究的低成本、高精度的开源基线。我们的 MST&＃43;&＃43;旨在填补这一空白。此外&＃xff0c;所有的源代码和预训练的模型在表2包括 11 个 SOTA方法是公开可用的。

6. 结论

        在本文中&＃xff0c;我们提出了第一个基于 Transformer 的框架 MST&＃43;&＃43;&＃xff0c;用于从 RGB 进行光谱重建。基于 HSI 的空间稀疏性和光谱自相似性&＃xff0c;我们采用 S-MSA 将每个光谱特征图作为一个自注意计算的令牌来组成基本单元SAB。然后 sab 建立 SST。最后&＃xff0c;我们的 MST&＃43;&＃43;由几个 SST 级联而成。MST&＃43;&＃43;采用多阶段学习方案&＃xff0c;从粗到细逐步提高重建质量。定量和定性实验表明&＃xff0c;我们的MST&＃43;&＃43;明显优于 SOTA 方法&＃xff0c;同时需要更便宜的内存和计算成本。令人印象深刻的是&＃xff0c;我们的 MST&＃43;&＃43;在 NTIRE2022 挑战赛中获得了 RGB 光谱重建的第一名。