ColorMamba:面向基于Mamba的高质量NIR到RGB光谱转换

ColorMamba: Towards High-quality NIR-to-RGB Spectral Translation with Mamba

https://github.com/AlexYangxx/ColorMamba

https://arxiv.org/pdf/2408.08087

Abstract

Translating NIR to the visible spectrum is challenging due to cross-domain complexities. Current models struggle to balance a broad receptive field with computational efficiency, limiting practical use. Although the Selective Structured State Space Model, especially the improved version, Mamba, excels in generative tasks by capturing long-range dependencies with linear complexity, its default approach of converting 2D images into 1D sequences neglects local context. In this work, we propose a simple but effective backbone, dubbed ColorMamba, which first introduces Mamba into spectral translation tasks. To explore global long-range dependencies and local context for efficient spectral translation, we introduce learnable padding tokens to enhance the distinction of image boundaries and prevent potential confusion within the sequence model. Furthermore, local convolutional enhancement and agent attention are designed to improve the vanilla Mamba. Moreover, we exploit the HSV color to provide multi-scale guidance in the reconstruction process for more accurate spectral translation. Extensive experiments show that our ColorMamba achieves a 1.02 improvement in terms of PSNR compared with the state-of-the-art method.

由于跨域复杂性，将近红外（NIR）光谱转换为可见光谱具有挑战性。当前的模型在平衡广泛感受野与计算效率方面存在困难，这限制了其实际应用。尽管选择性结构化状态空间模型（尤其是其改进版本Mamba）通过以线性复杂度捕获长距离依赖关系在生成任务中表现出色，但其将二维图像转换为一维序列的默认方法却忽略了局部上下文。

本文提出了一种简单但有效的主干网络，称为ColorMamba，它首次将Mamba引入光谱转换任务中。为了探索全局长距离依赖性和局部上下文以实现高效的光谱转换，引入了可学习的填充标记来增强图像边界的区分度，并防止序列模型内部潜在的混淆。此外，还设计了局部卷积增强和代理注意力机制来改进原始Mamba。而且，利用HSV颜色在重建过程中提供多尺度指导，以实现更精确的光谱转换。

大量实验表明，与最先进的方法相比，本文的ColorMamba在峰值信噪比（PSNR）方面实现了1.02的改进。

“What I cannot create, I do not understand.”

—— Richard P. Feynman, 1988

Introduction

背景与动机：

近年来，研究人员致力于超越简单的观察，以达到对视觉数据深刻理解的目标，这推动了生成艺术的发展。早期的方法如变分自编码器（VAE）和生成对抗网络（GAN）已在多种下游任务中展现出令人瞩目的性能，如图像超分辨率和灰度图像上色。

在本文中，研究者们聚焦于一个具体但重要的生成任务：近红外（NIR）到可见光（RGB）的光谱转换。NIR光谱（780nm-1000nm）虽然与可见光谱（380nm-780nm）相邻，但人类眼睛无法直接看到NIR图像。NIR成像已广泛应用于物体检测、夜间视频监控和遥感等领域。然而，NIR图像的光谱响应与人类的感知经验和计算机视觉系统所熟悉的可见光光谱反射大相径庭。

Figure 1: Visual effect display compared to three methods: Restorer Zamir et al. (2022), DRSformer Chen et al. (2023a), and CoColor Yang et al. (2023a).

研究现状与挑战：

当前NIR-to-RGB光谱转换方法主要集中在学习像素级的映射关系，并采用U-Net作为骨干网络进行密集预测任务。然而，这些方法大多基于卷积神经网络（CNN），其感受野是局部的，难以捕获全局上下文信息。

最近，基于Transformer架构的方法因其固有的全局感受野在图像翻译任务中表现出色，并超越了CNN框架的性能。然而，尽管Transformer模型在扩大感受野方面取得了显著进步，但它们在计算效率上仍存在挑战。

Mamba模型的引入与改进：

近期，结构化状态空间序列模型（S4）及其改进版本Mamba因其在线性复杂度下捕获长距离依赖性的能力而在生成任务中表现出色。然而，原始的Mamba模型默认将2D图像转换为1D序列，忽视了局部上下文信息。

针对上述问题，研究者们提出了ColorMamba模型，这是首次将Mamba引入光谱转换任务的工作。ColorMamba旨在通过引入可学习的填充令牌、局部卷积增强和代理注意力机制，同时探索全局长距离依赖性和局部上下文信息，以实现高效的光谱转换。

ColorMamba的主要贡献：

提出了基于Mamba的NIR-to-RGB光谱转换骨干网络ColorMamba，能够同时建模全局长距离依赖性和捕获局部上下文。

通过局部卷积增强、代理注意力和新的扫描策略，改进了标准Mamba模型，以形成视觉状态空间块（VSSB），有效解决了局部上下文忽视的问题，并提升了Mamba在2D图像上的性能。

提出了HSV颜色预测子网络，利用颜色先验在重建过程中提供多尺度指导，以实现更准确的光谱转换。

实验结果表明，ColorMamba在PSNR指标上相比当前最先进的方法提升了1.02，表明ColorMamba为光谱转换研究提供了一个强大且有前途的基础架构。

Figure 2: Status scanning strategy. We inject learnable padding tokens between two adjacent tokens that do not share approximate spatial correlation to enhance boundary distinction and prevent potential confusion within the sequence model.

Methodology

1. Preliminaries

Vision State Space Module (VSSM)

Vision State Space Module是一种通过递归计算将2D图像作为展平的一维序列来处理的模块，它能够利用全局感受野来捕捉长程依赖关系。然而，标准的Mamba模型在处理图像时采用的默认扫描策略会不经意间将空间上接近的像素放置在一维数组中的不同位置，导致所谓的“局部上下文忽视”现象，即相邻像素之间的空间相关性没有得到充分保留。

解决局部上下文忽视问题

为了解决局部上下文忽视问题，作者提出了在状态空间模型的扫描序列中插入可学习的填充令牌（padding tokens）。这些填充令牌被放置在两个在空间中不相关的相邻令牌之间，以增强图像边界的区分度，并防止序列模型中的潜在混淆。

具体来说，给定输入特征图 F∈R^{1×H×W×N}，首先将其填充为 Fp​∈R^{1×(H+2)×(W+2)×N}。然后，通过变换二维空间信息为一组四个一维序列，每个序列包含 (H+2)(W+2) 个元素，来展开特征图。这些序列通过四种不同的扫描路径重新组织，包括从左上角到右下角的对角线方向（左到右、上到下、右到左、下到上），以有效捕捉特征图的空间连续性。

序列处理与重构

重构后的序列形式化为​，其中 n=4 是序列的数量，L=(H+2)(W+2) 是每个序列的长度。根据连续线性时不变系统，作者使用线性常微分方程（ODE）来映射输入序列到输出序列，并通过离散化过程将其应用于二维视觉信号。Mamba利用零阶保持（ZOH）规则对ODE进行离散化，并通过递归形式对重构的序列进行处理，最终合并所有序列特征以得到输出映射 y。

公式化描述

• 连续时间系统：通过隐式潜在状态 h(t)∈R^N 映射输入 x(t)∈R 到输出 y(t)∈R：

  

  其中 A,B,C,D 是权重矩阵，N 是状态大小。

• 离散化：利用零阶保持规则对ODE进行离散化，得到离散化版本的状态更新和输出方程：

  

  其中 Δ 是时间尺度参数。

• 递归形式：基于重构的序列 ​​，使用递归形式更新隐藏状态和输出：

  

  然后，合并所有序列特征得到输出映射 y，并将其裁剪回原始尺寸。

2. Overall Architecture

架构概述

ColorMamba的整体架构旨在处理单色NIR图像（x_nir​∈R^{H×W×1}）作为输入，并输出彩色化的NIR图像（y_rgb​∈R^{H×W×3}）。为了实现这一目标，ColorMamba集成了多个关键组件，包括HSV颜色预测子网（GB​）、RGB重建网络（GA​）、融合块（Fusion Block）以及跨注意力块（Cross-Attention Block）。

HSV颜色预测子网（GB​）

HSV颜色预测子网（GB​）旨在提供稳健且动态的颜色先验，以辅助主RGB重建网络（GA​）在不同尺度上进行颜色信息的传递。这一子网通过预测HSV颜色空间中的颜色信息，为后续的RGB重建提供多尺度的颜色指导。

RGB重建网络（GA​）

RGB重建网络（GA​）是ColorMamba的核心部分，负责将输入的NIR图像转换为RGB图像。该网络通过结合HSV颜色预测子网提供的颜色先验和NIR图像本身的纹理信息，生成高质量的彩色图像。

融合块（Fusion Block）

为了保留和增强NIR输入图像中丰富的纹理信息，ColorMamba采用了融合块（Fusion Block）。该块利用拉普拉斯算子从NIR图像中分离出纹理特征，并将这些特征与HSV颜色预测子网产生的颜色图进行结合。这一步骤通过SPADE Resnet Block（Sun and Jung, 2020）实现，显著提高了不同区域颜色信息的准确性。

跨注意力块（Cross-Attention Block）

为了实现NIR特征图和纹理丰富的HSV颜色图的和谐融合，ColorMamba引入了跨注意力块（Cross-Attention Block）。这一机制促进了由生成器GA​产生的颜色预测与复杂纹理增强的HSV颜色图之间的无缝融合，从而生成最终的输出图像。

架构流程

1. 输入处理：首先，将单色NIR图像作为输入。
2. HSV颜色预测：通过HSV颜色预测子网（GB​）预测HSV颜色空间中的颜色信息。
3. RGB重建：利用RGB重建网络（GA​）结合HSV颜色先验和NIR图像的纹理信息，生成初步的RGB图像。
4. 纹理融合：通过融合块（Fusion Block）将NIR图像的纹理特征与HSV颜色图进行结合。
5. 跨注意力融合：利用跨注意力块（Cross-Attention Block）实现NIR特征图和HSV颜色图的和谐融合。
6. 输出：生成最终的彩色RGB图像。

Figure 3: The pipeline of ColorMamba. The model consists of two generative networks: (a) the RGB Reconstruction Network (GA) and (b) the HSV Color Prediction Sub-network (GB). (c), (d), (e), (f) illustrate details of the Shallow Feature Extraction layer, Visual State Space block (VSSB), Vision State Space Module (VSSM) and 2D-selective Scan Module (2D-SSM), respectively.

3. RGB Reconstruction Network

RGB Reconstruction Network（记作 GA​）是ColorMamba框架中的核心网络之一，主要负责从输入的近红外（NIR）图像生成对应的RGB图像。为了实现这一目标，网络设计考虑到了长距离依赖关系的建模和局部上下文信息的保持。以下是RGB Reconstruction Network的几个关键组件：

视觉状态空间块（Visual State Space Blocks, VSSBs）

• 基本构建：RGB Reconstruction Network使用Visual State Space Blocks（VSSBs）作为其构建块。VSSBs在Mamba模型的基础上进行了扩展，引入了局部卷积增强和代理注意力机制，以更有效地处理图像数据中的局部上下文和全局长距离依赖。
• 局部卷积增强：VSSBs利用局部卷积操作来增强网络对图像局部细节的捕捉能力，这对于提高重建图像的准确性至关重要。
• 代理注意力：为了进一步提升全局依赖的建模能力，VSSBs引入了代理注意力机制，帮助网络更好地理解全局结构和场景上下文。

网络结构

• 输入与输出：RGB Reconstruction Network的输入是预处理后的NIR图像，输出是生成的RGB图像。
• 融合机制：在重建过程中，网络会将来自HSV Color Prediction Sub-network的颜色信息与NIR图像的纹理特征进行融合。这种融合利用了空间自适应的归一化残差块（SPADE Resnet Block），旨在提高色彩在不同区域的准确性。
• 插值操作：由于网络处理的是扁平化的一维序列，生成的特征图需要恢复到原始尺寸。网络使用适当的插值操作来实现这一点，以确保最终输出图像的尺寸与输入图像一致。

4. HSV Color Prediction Sub-network

HSV Color Prediction Sub-network（记作 GB​）是ColorMamba框架中的另一个关键组件，其设计目标是利用HSV颜色空间来提供多尺度的颜色先验信息，从而辅助RGB Reconstruction Network生成更准确的RGB图像。

HSV颜色空间的优势

• 色相（Hue）：描述了颜色的基本属性，有助于模型识别出输入图像中的主要颜色成分。
• 饱和度（Saturation）：反映了颜色的强度或纯度，对于表现颜色的丰富度至关重要。
• 亮度（Value）：代表了颜色的明亮程度，对于区分图像中的亮区和暗区很有帮助。

网络结构与功能

• 多尺度指导：HSV Color Prediction Sub-network在不同的尺度上预测HSV颜色值，以提供多层次的颜色先验信息。这种多尺度策略有助于模型在多个级别上捕捉和传递颜色信息。
• 特征融合：生成的HSV颜色图随后与NIR图像的纹理特征进行融合，利用SPADE Resnet Block等技术来实现特征的深度融合和颜色的准确再现。
• 颜色映射：HSV颜色预测的最终目标是将这些颜色先验信息映射到RGB颜色空间，以便与RGB Reconstruction Network的输出进行对齐和整合。

5. Visual State Space Block

在ColorMamba模型中，为了克服标准Mamba模型在处理2D视觉任务时可能遇到的局部上下文忽略问题，我们引入了Visual State Space Block（VSSB）作为关键构建块。VSSB不仅继承了Mamba模型在捕获全局长距离依赖方面的优势，还通过增加局部卷积增强和代理注意力机制来改进模型对局部上下文的建模能力。

局部卷积增强

局部卷积增强通过在VSSB中集成卷积层来捕捉图像的局部空间信息。传统的Mamba模型在处理图像时，将其视为扁平化的一维序列，这可能会破坏像素之间的空间相关性。为了弥补这一缺陷，我们在Mamba块之前和之后添加了卷积层。这些卷积层能够提取图像的局部特征，并将这些特征作为Mamba处理的输入和输出的一部分，从而增强了模型对局部上下文的敏感性。

代理注意力

除了局部卷积增强外，VSSB还引入了代理注意力机制来进一步改进模型对关键区域的关注。代理注意力是一种高效的注意力机制，它能够在不显著增加计算复杂性的情况下，使模型更加关注于输入中的重要区域。在VSSB中，代理注意力被用于调整Mamba块的输出，以确保模型能够更准确地捕捉到对光谱转换任务至关重要的信息。

新的扫描策略

为了增强图像边界的区分度并防止序列模型中的潜在混淆，我们提出了一种新的扫描策略。具体来说，我们在状态空间模型的扫描序列中，在两个相邻但不具有近似空间相关性的令牌之间插入可学习的填充令牌。这些填充令牌有助于Mamba块更准确地解释图像边缘，从而增强了模型的空间意识和序列数据处理的完整性。

6. Objectives

在ColorMamba模型中，定义了一系列优化目标来指导模型的训练过程。这些目标旨在确保生成的RGB图像不仅在颜色上接近真实图像，而且在纹理和细节上也能保持高度一致。

重构损失

首先，使用了重构损失（Reconstruction Loss）来衡量生成的RGB图像与真实RGB图像之间的差异。具体来说，我们采用了均方误差（MSE）作为重构损失的度量标准，以确保生成的图像在像素级别上与真实图像保持一致。

对抗损失

为了进一步提高生成图像的真实性，引入了对抗损失（Adversarial Loss）。对抗损失通过训练一个判别器来区分生成的RGB图像和真实的RGB图像，从而促使生成器生成更加逼真的图像以欺骗判别器。对抗损失的使用有助于捕获真实图像中的复杂纹理和细节信息，从而提高生成图像的质量。

HSV颜色预测损失

为了充分利用HSV颜色空间提供的多尺度指导，定义了HSV颜色预测损失（HSV Color Prediction Loss）。该损失函数衡量了生成图像的HSV颜色表示与真实图像的HSV颜色表示之间的差异。通过优化HSV颜色预测损失，可以确保生成的RGB图像在颜色分布和饱和度等方面与真实图像保持一致。

总体优化目标

综上所述，ColorMamba模型的总体优化目标可以表示为上述三个损失函数的加权和：

其中，λ_rec​、λ_adv​和λ_fea分别是重构损失、对抗损失和HSV颜色预测损失的权重系数，用于平衡不同损失函数对总体优化目标的影响。

Results

Figure 4: Visual comparison of different methods on testing datasets. From left to right are SST Yan et al. (2020), NIR-GNN Valsesia et al. (2020), MFF Yan et al. (2020), ATCGAN Yang and Chen (2020), Restormer Zamir et al. (2022), DRSformer Chen et al. (2023a), MPFNet Yang et al. (2023b), CoColor Yang et al. (2023a), and MCFNet Zhai et al. (2024).

Figure 5: Visual examples of deficiencies. Our ColorMamba generates some oversaturated images compared to ground truths.