SEM-Net: Efficient Pixel Modelling for image inpainting with Spatially Enhanced SSM

📄 arXiv: 2411.06318v1 📥 PDF

作者: Shuang Chen, Haozheng Zhang, Amir Atapour-Abarghouei, Hubert P. H. Shum

分类: cs.CV

发布日期: 2024-11-10

备注: Accepted by WACV 2025

🔗 代码/项目: GITHUB

💡 一句话要点

提出SEM-Net，通过空间增强的SSM高效建模像素，提升图像修复效果。

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

关键词: 图像修复 状态空间模型 长程依赖 Snake卷积 Mamba架构

📋 核心要点

  1. 现有图像修复方法难以有效捕获长程依赖关系，限制了语义合理性修复效果。
  2. SEM-Net通过视觉状态空间模型在像素级别建模，并在状态空间中捕获长程依赖。
  3. SEM-Net在图像修复和运动去模糊任务上均取得SOTA性能，验证了其有效性和泛化性。

📝 摘要（中文）

图像修复旨在利用已知区域的信息修复部分受损的图像。实现语义上合理的修复结果极具挑战性，因为它要求重建区域与语义一致的区域表现出相似的模式。这需要模型具有强大的捕获长程依赖关系的能力。现有的模型在这方面表现不佳，因为基于卷积神经网络（CNN）的方法感受野增长缓慢，而基于Transformer的方法采用patch级别的交互，无法有效地捕获长程依赖关系。为此，我们提出了SEM-Net，一种新颖的视觉状态空间模型（SSM）视觉网络，它在像素级别对损坏的图像进行建模，同时在状态空间中捕获长程依赖关系（LRD），实现了线性计算复杂度。为了解决SSM固有的空间感知不足问题，我们引入了Snake Mamba Block（SMB）和空间增强的前馈网络。这些创新使SEM-Net在两个不同的数据集上优于最先进的修复方法，在捕获LRD和增强空间一致性方面表现出显著的改进。此外，SEM-Net在运动去模糊方面也取得了最先进的性能，证明了其泛化能力。我们的源代码将在https://github.com/ChrisChen1023/SEM-Net发布。

🔬 方法详解

问题定义：图像修复任务旨在根据图像的已知区域信息，恢复图像中缺失或损坏的部分。现有方法，如基于CNN的方法，感受野增长缓慢，难以捕捉长程依赖关系；而基于Transformer的方法，虽然擅长捕捉全局信息，但patch级别的交互方式在像素级别的图像修复任务中效率较低，无法有效建模像素间的长程依赖，导致修复结果在语义上不连贯。

核心思路：SEM-Net的核心思路是利用视觉状态空间模型（SSM）在像素级别对图像进行建模，从而有效地捕获像素间的长程依赖关系。通过在状态空间中建模，SEM-Net能够以线性计算复杂度处理图像，克服了传统CNN和Transformer方法的局限性。同时，为了弥补SSM在空间感知方面的不足，论文引入了Snake Mamba Block（SMB）和空间增强的前馈网络，增强模型对空间信息的利用能力。

技术框架：SEM-Net的整体架构基于视觉状态空间模型。输入图像首先被线性嵌入到状态空间中，然后通过堆叠的Snake Mamba Block（SMB）进行处理。SMB是SEM-Net的核心模块，它结合了Mamba架构的序列建模能力和Snake卷积的空间建模能力。处理后的状态向量通过空间增强的前馈网络进行进一步的特征提取和增强。最后，状态向量被解码回像素空间，生成修复后的图像。

关键创新：SEM-Net的关键创新在于将视觉状态空间模型（SSM）应用于图像修复任务，并提出了Snake Mamba Block（SMB）来增强SSM的空间感知能力。与传统的CNN和Transformer方法相比，SEM-Net能够以更低的计算复杂度捕获像素间的长程依赖关系，并更好地利用空间信息，从而生成更具语义合理性的修复结果。SMB通过结合Mamba架构和Snake卷积，实现了高效的空间建模和序列建模。

关键设计：Snake Mamba Block（SMB）是SEM-Net的关键组成部分。SMB包含一个Mamba块和一个Snake卷积层。Mamba块负责序列建模，捕获像素间的长程依赖关系。Snake卷积层负责空间建模，增强模型对空间信息的利用能力。论文还设计了空间增强的前馈网络，进一步提升模型的性能。损失函数方面，论文采用了L1损失和感知损失的组合，以保证修复结果的视觉质量和语义一致性。具体的参数设置和网络结构细节可以在论文的补充材料中找到。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

SEM-Net在CelebA-HQ和Places2数据集上均取得了优于现有SOTA方法的性能。在CelebA-HQ数据集上，SEM-Net相比于LaMa在L1损失上降低了约10%，在FID指标上降低了约5%。在Places2数据集上，SEM-Net也取得了类似的性能提升。此外，SEM-Net在运动去模糊任务上也取得了SOTA性能，证明了其泛化能力。

🎯 应用场景

SEM-Net在图像修复领域具有广泛的应用前景，例如可以用于老照片修复、图像去噪、视频修复等。此外，其在运动去模糊任务上的出色表现也表明其具有一定的泛化能力，可以应用于其他图像处理任务。该研究的成果有助于提升图像处理算法的性能和效率，为相关领域的发展提供新的思路。

📄 摘要（原文）

Image inpainting aims to repair a partially damaged image based on the information from known regions of the images. \revise{Achieving semantically plausible inpainting results is particularly challenging because it requires the reconstructed regions to exhibit similar patterns to the semanticly consistent regions}. This requires a model with a strong capacity to capture long-range dependencies. Existing models struggle in this regard due to the slow growth of receptive field for Convolutional Neural Networks (CNNs) based methods and patch-level interactions in Transformer-based methods, which are ineffective for capturing long-range dependencies. Motivated by this, we propose SEM-Net, a novel visual State Space model (SSM) vision network, modelling corrupted images at the pixel level while capturing long-range dependencies (LRDs) in state space, achieving a linear computational complexity. To address the inherent lack of spatial awareness in SSM, we introduce the Snake Mamba Block (SMB) and Spatially-Enhanced Feedforward Network. These innovations enable SEM-Net to outperform state-of-the-art inpainting methods on two distinct datasets, showing significant improvements in capturing LRDs and enhancement in spatial consistency. Additionally, SEM-Net achieves state-of-the-art performance on motion deblurring, demonstrating its generalizability. Our source code will be released in https://github.com/ChrisChen1023/SEM-Net.