DS-TransUNet: Dual Swin Transformer U-Net for Medical Image Segmentation
论文:https://arxiv.org/abs/2106.06716
Introduction
医学图像分割是一个重要而又具有挑战性的研究问题,涉及到临床应用中的许多常见任务,如息肉分割、病变分割、细胞分割等。而医学图像分割是医学图像处理与分析领域中一个复杂而关键的环节,在计算机辅助临床诊断系统中起着重要的作用。其目的是通过半自动或自动的过程对医学图像中具有特殊意义的部分进行分割,提取相关特征,为临床诊断和病理研究提供可靠的依据,协助医生做出更准确的诊断。
随着U-Net的普及,许多新颖的模型被提出,如UNet++、Res-UNet、Attention U-Net、DenseUNet、R2U-Net、KiU-Net和UNet 3+,它们都是专为医学图像分割而设计的,并具有表达性能。
Encoder
在模型的整体结构中,作者使用U型架构。对于编码器,使用 Swin Transformer 提取特征。如下图所示,将输入的医学图像首先分割成 H s × H s \frac{H}{s}×\frac{H}{s} sH×sH 的非重叠patch,其中s为patch的大小。每个patch被当作一个“ token ”,通过线性嵌入层投射到维度C。由于patch是通过卷积运算得到的,所以这里不需要额外的位置信息。
这些 patch token 被正式送入Swin Transformer,它包含四个阶段,每个阶段包含一定数量的 Swin Transformer Block,包括窗口MSA (W-MSA)和移动窗口MSA (SW-MSA)。为了产生 层次表示 , token 的数量将随着网络的深入而减少;在前三个阶段,输入特征经过 Swin Transformer Block 变换后经过 patch merge layer 降低特征分辨率,增加维数。具体来说, patch merge layer 将 相邻2×2的每一组patch的特征串联起来,然后对所连接的通道维度特征应用一个线性层。这将减少 token 的数量2× 2 = 4,分辨率的2×下采样,输出维数增加2。因此,四个阶段的输出分辨率分别为 H s × H s \frac{H}{s}×\frac{H}{s} sH×sH 、 H 2 s × H 2 s \frac{H}{2s}×\frac{H}{2s} 2sH×2sH、 H 4 s × H 4 s \frac{H}{4s}×\frac{H}{4s} 4sH×4sH 和 H 8 s × H 8 s \frac{H}{8s}×\frac{H}{8s} 8sH×8sH ,维度分别为 C 、 2 C 、 4 C 和 8 C C、2C、4C和8C C、2C、4C和8C。
Decoder
如下图所示,解码器主要由三个阶段组成。与U-Net及其变体不同,模型的每个阶段不仅包括上采样(最近上采样)和跳跃连接,而且还包括 Swin Transformer Block 。具体来说,编码器中第4阶段的输出作为解码器的初始输入。在解码器的每一阶段,将输入特征向上采样2次,然后与同一阶段编码器对应的跳跃连接特征图连接。在那之后,输出被送入 Swin Transformer Block 。我们选择这个设计是因为
- 它可以充分利用编码器和上采样的特性
- 它可以在解码器中建立长期依赖和全局上下文交互,以获得更好的解码性能。
经过以上三个阶段,我们可以得到分辨率为 H 4 × H 4 \frac{H}{4}×\frac{H}{4} 4H×4H 的输出。直接使用4次上采样算子会丢失很多浅层特征,因此,我们通过串联两个块对输入图像进行降采样,得到了分辨率为 H × W H×W H×W和 H 2 × H 2 \frac{H}{2}×\frac{H}{2} 2H×2H的低阶特征,其中每个块包含一个3×3卷积层,一组归一化层和ReLU。所有这些输出特征将被用来通过跳跃连接得到最终的掩模预测。
多尺度特征表示
虽然 self-attention 可以有效地建立 patches 之间的长期依赖关系,但对 patch 进行划分时忽略了每个 patch 内部像素级的内在结构特征,从而导致边缘、线等浅层特征信息的丢失。大尺度可以更好地捕获粗粒度特征,而