SwinUNet2022

1. 概述

本文提出了一种以$Swin$变压器层为基本块的$SUNet$恢复模型，并将其应用于$UNet$架构中进行图像去噪。

2. 背景

图像恢复是一种重要的低级图像处理方法，可以提高其在目标检测、图像分割和图像分类等高级视觉任务中的性能。在一般的恢复任务中，一个被损坏的图像Y可以表示为：
$$Y=D(X)+n\text{\hspace{1em}(1)}$$
其中$X$是一个干净的图像，$D(\cdot ）$表示退化函数，$n$表示加性噪声。一些常见的恢复任务是去噪、去模糊和去阻塞。

2.1 CNN局限性

虽然大多数基于卷积神经网络(CNN)的方法都取得了良好的性能，但卷积层存在几个问题。首先，卷积核与图像的内容无关（无法与图像内容相适应）。使用相同的卷积核来恢复不同的图像区域可能不是最好的解决方案。其次，由于卷积核可以看作是一个小块，其中获取的特征是局部信息，换句话说，当我们进行长期依赖建模时，全局信息就会丢失。

3. 结构

3.1 UNet

目前，UNet由于具有层次特征映射来获得丰富的多尺度上下文特征，是许多图像处理应用中著名的架构。此外，它利用编码器和解码器之间的跳跃连接来增强图像的重建过程。UNet被广泛应用于许多计算机视觉任务，如分割、恢复[。此外，它还有各种改进的版本，如Res-UNet，Dense-UNet，Attention-UNet[和Non-local-UNet。由于具有较强的自适应骨干网，UNet可以很容易地应用于不同的提取块，以提高性能。

3.2 Swin Transformer

Transformer模型在自然语言处理(NLP)领域取得了成功，并具有良好的竞争性能，特别是在图像分类方面。然而，直接使用Transformer到视觉任务的两个主要问题是：

（1）图像和序列之间的尺度差异很大。由于Transformer需要参数量为一维序列参数的平方倍，所以存在长序列建模的缺陷。

（2）Transformer不擅长解决实例分割等密集预测任务，即像素级任务。然而，Swin Transfomer通过滑动窗口解决了上述问题，降低了参数，并在许多像素级视觉任务中实现了最先进的性能。

3.3 SUNet

所提出的Swin Transformer UNet(SUNet)的架构是基于图像分割模型，如上图所示。SUNet由三个模块组成：

（1）浅层特征提取；

（2）UNet特征提取；

（3）重建模块

浅层特征提取模块：

对于有噪声的输入图像$Y\in R^{H\times W\times 3}$，其中H，W为失真图像的分辨率。我们使用单个3×3卷积层$M_{SFE}(\cdot )$获取输入图像的颜色或纹理等低频信息。浅特征$F_{shallow}\in R^{H\times W\times C}$可以表示为:
$$F_{shallow}=M_{SFE}(Y)\text{\hspace{1em}(2)}$$
其中，C是浅层特征的通道数，在后一个实验部分中，我们都设置为96.

UNet 特征提取网络：

然后，将浅层特征$F_{shallow}$输入UNet特征提取$M_{UFE}(\cdot )$，UNet用来提取高级、多尺度深度特征$F_{deep}\in R^{H\times W\times C}$：
$$F_{deep}=M_{UFE}(F_{shallow})\text{\hspace{1em}(3)}$$
其中，$M_{UFE}(\cdot )$是带有Swin变压器块的UNet架构，它在单个块中包含8个Swin Transformer层来代替卷积。Swin Transformer Block(STB)和Swin Transformer Layer(STL)将在下一小节中进行详细说明。

重建层：

最后，我们仍然使用3×3卷积$M_R(\cdot )$从深度特征$F_{deep}$中生成无噪声图像$\hat{X}\in R^{H\times W\times 3}$，其公式为：
$$\hat{X}=M_R(F_{deep})\text{\hspace{1em}(4)}$$
注意，$\hat{X}$是以噪声图像$Y$作为SUNet的输入得到的，其中$X$是（1）中Y图像的原高分率图像。

3.4 Loss function

我们优化了我们的SUNet端到端与规则的$L1$像素损失的图像去噪：
$$L_{denoise}=||\hat{X}-X|{|}_1\text{\hspace{1em}(5)}$$

3.5 Swin Transformer Block

在UNet提取模块中，我们使用STB来代替传统的卷积层，如下图所示。STL是基于NLP中的原始Transformer Layer。STL的数量总是2的倍数，其中一个是window multi-head-self-attention(W-MSA)，另一个是shifted-window multi-head self-attention(SW-MSA)。

STL的公式描述：
$$\begin{gathered} {\hat{f}}^L=W-MSA(LN(f^{L-1}))+f^{L-1} \\ {f}^L=MLP(LN({\hat{f}}^L))+{\hat{f}}^L \\ {\hat{f}}^{L+1}=SW-MSA(LN(f^L))+f^L \\ {f}^{L+1}=MLP(LN({\hat{f}}^{L+1}))+{\hat{f}}^{L+1}\text{\hspace{1em}(6)} \end{gathered}$$
其中，$LN(\cdot )$表示为层归一化，$MLP$是多层感知器，它具有两个完全连接的层，同时后面跟一个线性单位(GELU)激活函数。

3.6 Resizing module

由于UNet具有不同的特征图尺度，因此调整大小的模块（例如，下样本和上样本）是必要的。在我们的SUNet中，我们使用$patchmerging$，并提出$dualup-sample$分别作为下样本和上样本模块。

3.6.1 patch merging

对于降采样模块，该文将每一组2×2相邻斑块的输入特征连接起来，然后使用线性层获得指定的输出通道特征。我们也可以把这看作是做卷积操作的第一步，也就是展开输入的特征映射。

3.6.2 Dual up-sample

对于上采样，原始的Swin-UNet采用patch expanding方法，等价于上采样模块中的转置卷积。然而，转置卷积很容易面对块效应。在这里，我们提出了一个新的模块，称为双上样本，它包括两种现有的上样本方法（即Bilinear和PixelShuffle），以防止棋盘式的artifacts。所提出的上采样模块的体系结构如下图所示。

4. 结果

评估指标：

为了进行定量比较，我们考虑了峰值信噪比(PSNR)和结构相似度(SSIM)指数度量。

训练集：

采用DIV2K作为训练集，一共有900张高清图片。我们对每个训练图像随机裁剪100个大小为$256\times 256$的斑块，并对$800$张训练图像从$\sigma =5$到$\sigma =50$的$patch$中随机添加AWGN噪声。至于验证集，我们直接使用包含100张图像的测试集，并添加具有三种不同噪声水平的AWGN，$\sigma =10、\sigma =30和\sigma =50$。

测试集：

对于评估，我们选择了CBSD68数据集，它有68张彩色图像，分辨率为768×512，以及Kodak24张数据集，由24张图像组成，图像大小为321×481。

在表1中，我们对去噪图像进行了客观的质量评价，并观察到以下三件事：

（1）该文的SUNet具有竞争性的SSIM值，因为Swin-Transformer是基于全局信息（q，k，v可以提取全局信息），使得去噪图像拥有更多的视觉效果。

（2）与基于unet的方法(DHDN、RDUNet)相比，该文所提出的SUNet模型中参数（↓60%）和FLOPs（↓3%）较少，在PSNR和SSIM上仍保持良好的得分

（3）与基于cnn的方法(DnCNN，IrCNN，FFDNet)相比，该文得到了其中最好的PSNR和SSIM结果，以及几乎相同的FLOPs。虽然该文的模型的参数最多(99M)，但它是由于自注意操作不能共享核的权值造成的。

4. 总结

• 提出了一种基于图像分割的双unet模型的双变换网络进行图像去噪。
• 该文提出了一种双上样本块结构，它包括亚像素方法和双线性上样本方法，以防止棋盘伪影。实验结果表明，该方法优于转置卷积的原始上样本。
• 该文的模型是第一个结合Swin变压器和UNet进行去噪的模型。

Reference: Swin-unet: Unet-like pure transformer for medical image segmentation