Mamba-UIE: Enhancing Underwater Images with Physical Model Constraint

📄 arXiv: 2407.19248v2 📥 PDF

作者: Song Zhang, Yuqing Duan, Daoliang Li, Ran Zhao

分类: cs.AI

发布日期: 2024-07-27 (更新: 2024-07-31)

🔗 代码/项目: GITHUB

💡 一句话要点

Mamba-UIE：提出基于物理模型约束的Mamba网络用于水下图像增强

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 水下图像增强 物理模型约束 Mamba网络 长程依赖 图像重构

📋 核心要点

  1. 现有水下图像增强方法难以兼顾长程依赖建模和计算效率，且缺乏有效的物理模型约束，导致图像真实感不足。
  2. Mamba-UIE通过分解图像成分并施加重构一致性约束，实现了对水下图像增强过程的有效物理约束。
  3. Mamba-UIE网络基于线性复杂度状态空间模型，在建模长程依赖的同时，保留了CNN骨干网络以恢复局部特征和细节，实验结果表明其优于现有方法。

📝 摘要（中文）

水下图像增强(UIE)中，卷积神经网络(CNN)在建模长程依赖方面存在固有局限性，难以有效恢复全局特征。Transformer虽然擅长建模长程依赖，但其计算复杂度随图像分辨率的增加呈平方增长，效率面临严峻挑战。此外，大多数监督学习方法缺乏有效的物理模型约束，导致生成图像的真实感不足和过拟合。为了解决这些问题，我们提出了一种基于物理模型约束的水下图像增强框架Mamba-UIE。具体而言，我们将输入图像分解为四个组成部分：水下场景辐射、直接透射图、反向散射透射图和全局背景光。这些分量根据修正后的水下图像形成模型重新组装，并在重构图像和原始图像之间应用重构一致性约束，从而实现对水下图像增强过程的有效物理约束。为了解决Transformer处理长序列时计算复杂度过高的问题，我们引入了基于线性复杂度状态空间模型的Mamba-UIE网络。通过在卷积块中加入Mamba，可以在通道和空间层面对长程依赖进行建模，同时保留CNN骨干网络以恢复局部特征和细节。在三个公共数据集上的大量实验表明，我们提出的Mamba-UIE优于现有的最先进方法，在UIEB数据集上实现了27.13的PSNR和0.93的SSIM。

🔬 方法详解

问题定义：水下图像增强旨在恢复水下图像的清晰度和色彩，但现有方法在建模长程依赖、计算效率和物理真实性方面存在不足。CNN难以捕捉全局信息，Transformer计算复杂度高，而缺乏物理模型约束容易导致图像失真和过拟合。

核心思路：Mamba-UIE的核心思路是将水下图像增强问题置于物理模型约束下，通过分解图像成分并利用修正后的水下图像形成模型进行重构，从而保证增强结果的物理合理性。同时，利用Mamba结构高效地建模长程依赖，提升全局特征恢复能力。

技术框架：Mamba-UIE框架首先将输入图像分解为水下场景辐射、直接透射图、反向散射透射图和全局背景光四个组成部分。然后，利用Mamba网络对这些分量进行增强和估计。最后，根据修正后的水下图像形成模型，将这些分量重新组合成增强后的图像。为了保证增强结果与原始图像的一致性，框架还引入了重构一致性约束。

关键创新：Mamba-UIE的关键创新在于将Mamba结构引入水下图像增强任务，并结合物理模型约束。Mamba结构能够以线性复杂度建模长程依赖，克服了Transformer的计算瓶颈。物理模型约束则保证了增强结果的物理合理性，避免了图像失真和过拟合。

关键设计：Mamba-UIE网络在卷积块中加入了Mamba结构，从而在通道和空间层面对长程依赖进行建模。网络采用U-Net结构，以保留局部特征和细节。损失函数包括重构损失、感知损失和对抗损失，以保证增强结果的质量和真实感。重构一致性约束通过最小化重构图像和原始图像之间的差异来实现。

📊 实验亮点

实验结果表明，Mamba-UIE在三个公共数据集上均取得了优于现有方法的性能。在UIEB数据集上，Mamba-UIE的PSNR达到了27.13，SSIM达到了0.93，显著优于其他最先进的方法，证明了其在水下图像增强方面的有效性。

🎯 应用场景

Mamba-UIE在水下机器人视觉、海洋生物研究、水下考古、水下工程检测等领域具有广泛的应用前景。它可以提高水下图像的清晰度和可读性，帮助研究人员更好地观察和分析水下环境，从而促进相关领域的发展。

📄 摘要（原文）

In underwater image enhancement (UIE), convolutional neural networks (CNN) have inherent limitations in modeling long-range dependencies and are less effective in recovering global features. While Transformers excel at modeling long-range dependencies, their quadratic computational complexity with increasing image resolution presents significant efficiency challenges. Additionally, most supervised learning methods lack effective physical model constraint, which can lead to insufficient realism and overfitting in generated images. To address these issues, we propose a physical model constraint-based underwater image enhancement framework, Mamba-UIE. Specifically, we decompose the input image into four components: underwater scene radiance, direct transmission map, backscatter transmission map, and global background light. These components are reassembled according to the revised underwater image formation model, and the reconstruction consistency constraint is applied between the reconstructed image and the original image, thereby achieving effective physical constraint on the underwater image enhancement process. To tackle the quadratic computational complexity of Transformers when handling long sequences, we introduce the Mamba-UIE network based on linear complexity state space models. By incorporating the Mamba in Convolution block, long-range dependencies are modeled at both the channel and spatial levels, while the CNN backbone is retained to recover local features and details. Extensive experiments on three public datasets demonstrate that our proposed Mamba-UIE outperforms existing state-of-the-art methods, achieving a PSNR of 27.13 and an SSIM of 0.93 on the UIEB dataset. Our method is available at https://github.com/zhangsong1213/Mamba-UIE.