[论文解读] SUNet: Swin Transformer UNet for Image Denoising

Chi-Mao Fan, Tsung-Jung Liu+1 more|arXiv (Cornell University)|Feb 28, 2022

Image and Signal Denoising Methods被引用 4

一句话总结

本文提出SUNet，一种基于Swin Transformer的U-Net架构，用于图像去噪。该模型用Swin Transformer模块替代传统卷积层，以捕捉长距离依赖关系和全局上下文信息。在基准数据集上，模型实现了最先进的性能，PSNR和SSIM指标均得到提升，同时引入的双路上采样模块有效减少了棋盘纹伪影，提升了重建质量。

ABSTRACT

Image restoration is a challenging ill-posed problem which also has been a long-standing issue. In the past few years, the convolution neural networks (CNNs) almost dominated the computer vision and had achieved considerable success in different levels of vision tasks including image restoration. However, recently the Swin Transformer-based model also shows impressive performance, even surpasses the CNN-based methods to become the state-of-the-art on high-level vision tasks. In this paper, we proposed a restoration model called SUNet which uses the Swin Transformer layer as our basic block and then is applied to UNet architecture for image denoising. The source code and pre-trained models are available at https://github.com/FanChiMao/SUNet.

研究动机与目标

• 为解决基于CNN的图像去噪模型存在的局部特征依赖性和固定卷积核应用等局限性。
• 探索Swin Transformer在低层次视觉任务中的潜力，特别是图像去噪任务，其中全局上下文建模至关重要。
• 设计一种融合Swin Transformer模块的U-Net架构，以提升特征表示能力和重建保真度。
• 通过引入结合亚像素与双线性插值方法的双路上采样模块，减轻特征上采样过程中的棋盘纹伪影。
• 在标准基准数据集上验证，基于Swin Transformer主干网络的U-Net框架实现图像去噪新SOTA性能。

提出的方法

• SUNet在U-Net的编码器和解码器中用Swin Transformer模块替代标准卷积层，以建模长距离依赖关系并捕捉全局上下文信息。
• 模型首先通过一个浅层3×3卷积层提取初始特征，随后经由主干Swin Transformer-based U-Net编码器-解码器路径进行处理。
• 引入双路上采样模块，结合亚像素与双线性插值方法，以减少棋盘纹伪影并提升特征重建质量。
• Swin Transformer模块采用移位窗口机制，在保持局部与全局感受野自注意力的同时降低计算复杂度。
• 网络在DIV2K数据集的图像块上进行端到端训练，采用均方误差（MSE）损失函数，并添加不同水平的加性白高斯噪声（AWGN）。
• 最终通过一个3×3卷积层重建输出，生成去噪后的图像。

实验结果

研究问题

• RQ1 基于Swin Transformer的U-Net架构是否能在图像去噪任务中超越现有的基于CNN和U-Net的模型？
• RQ2 将Swin Transformer集成到U-Net架构中，对特征表示和重建质量有何影响？
• RQ3 所提出的双路上采样模块是否能有效减少棋盘纹伪影，相比转置卷积上采样方法？
• RQ4 与现有U-Net变体相比，该模型是否能在参数量和FLOPs更少的情况下实现具有竞争力的性能？
• RQ5 该模型在不同噪声水平和真实世界图像数据集上的泛化能力如何？

主要发现

• 在CBSD68数据集上，当σ=50时，SUNet的PSNR达到36.79 dB，SSIM达到0.953，优于所有对比的基于CNN和U-Net的模型。
• 在Kodak24数据集上，当σ=50时，SUNet的PSNR为29.54 dB，SSIM为0.810，展现出在多样化图像集合中的强大泛化能力。
• 尽管参数量达9900万（高于多数基线模型），SUNet相比DHDN和RDUNet仍分别减少3%的FLOPs和60%的参数量，体现出显著的效率优势。
• 双路上采样模块显著减少了棋盘纹伪影，视觉对比结果表明纹理更平滑、边缘更清晰。
• 在CBSD68数据集上，当σ=50时，SUNet达到最高的SSIM（0.958），表明其在感知质量与结构保真度方面表现卓越。
• 该模型在所有噪声水平（σ=10, 30, 50）下均保持优异性能，且在PSNR和SSIM指标上持续优于先前方法。