Swin-Unet【CVPR2021】

文章：Cao H, Wang Y, Chen J, et al. Swin-Unet: Unet-like Pure Transformer for Medical Image Segmentation[J]. arXiv preprint arXiv:2105.05537, 2021.
代码:link.

Abstract(创新点): 由于卷积操作的局限性，它不能很好地学习全局和大范围的语义信息。本文提出Swin-Unet，是一种类似于Unet的纯Transformer,用于医学图像分割。
标记化的图像块被输入到基于Transformer的U行编码器-解码器架构中，并带有用于全局语义特征学习的跳过连接。特别地，作者使用带有移位窗口的Swin Transformer作为编码器来提取上下文特征。 并且设计了具有补丁扩展层的基于对称Swin Tansformer的解码器来执行上采样操作，以恢复特征图的空间分辨率。

网络简介：

输入的医学图像被分成不重叠的图像补丁，每个补丁都被视为一个标记，并输入到基于Transformer的编码器中以学习深度特征表示，提取的上下文特征被具有补丁扩展此层的解码器上采样，并通过跳过连接与来自编码器的多尺度特征融合，从而恢复特征图的空间分辨率并进行进一步的分割预测。具体贡献可总结如下：

1. 基于Swin Transformer模块，我们构建了一个具有跳过连接的对称编码器-解码器架构，在编码器中，实现了从局部到全局的self-attention；在解码器中，解码器被上采样到相应像素级分割预测的输入分辨率。
2. 在不使用卷积或插值运算的情况下，开发了一个补丁扩展层，用于实现上采样和特征维度的增加。
3. 实验发现，skip connection对于Transformer也是有效的，因此最终构建了一个基于Transformer的带有跳跃连接的纯U形编解码器架构，命名为Swin-Unet。

架构总览:
Swin-Unet由encoder，bootleneck，decoder和skip connection组成。基本的Swin-Unet单元是Swin Transformer块。对于编码器，为了将输入图形转化为嵌入序列，医学图像被分成了不重叠的块，块大小为4 * 4. 通过这种划分方法，每个补丁的特征维度变为4 * 4 * 3=48. 此外，线性嵌入层用于将特征维度投影到任意维度(表示为c). 转换后的补丁标记通过几个Swin Transformer块和补丁合并层以生成分层特征表示。具体而言，补丁合并层负责下采样进而增维，而Swin Transformer负责特征表示学习。受U-Net的启发，我们设计了一个基于Transformer的对称解码器。解码器由Swin Transformer块和补丁扩展层组成。提取的上下文特征通过skip connection与来自编码器的多尺度特征融合，以补充由下采样引起的空间信息丢失。与补丁合并层相反，补丁扩展层专门设计用于上采样。补丁扩展层将相邻维度的特征图重塑为具有2 * 分辨率的上采样大特征图。最后的补丁扩展层用于执行4 * 的上采样，从而将特征图还原为输入是的分辨率(W * H) .然后在这些上采样后的特征上应用线性投影层输出像素级别的语义分割预测。

SwinTransformer block

Swin Transformer是基于移位窗口构建的。每一个Swin Transformer都是由LayerNorm(LN)层，多头自注意力模块，残差连接和具有GELU非线性的双层MLP组成.基于窗口的多头自注意力(W-MSA)模块和基于移动窗口的多头自注意力(SW-MSA)模块被分别应用于两个连续的Transformer块中，基于上述窗口划分机制，连续的Swin-Transformer可以被表述为：

self-attention的计算如下：

其中
表示query，key和value矩阵，
代表窗口中的补丁数和query或key矩阵的维度。其中B的值取自偏置矩阵
Encoder：
在编码器中，分辨率为$\frac{H}{4}\times \frac{W}{4}$ 的C维标记被馈送到两个连续的Swin Transformer块中进行表征学习，其中特征维度和分辨率保持不变。同时，补丁合并层将减少标记的数量($2\times$下采样)并将特征维度增加到原始维度的2倍。此过程将在编码器中重复三次。
patch merging layer：输入补丁被分成4部分，并由补丁合并层连接在一起。通过这样的处理，特征分辨率将被进行$2\times$的下采样。并且，由于连接操作导致特征维度增加了$4\times$, 因此在连接的特征上应用线性层将特征维度统一为原始的$2\times$.

Bootleneck:
由于Transformer太深而无法收敛，因此仅使用两个连续的Swin Transformer块来构建bootleneck以学习深度特征表示。在bootleneck中，特征维度和分辨率保持不变。

Decoder:
对应于编码器，对称解码器是基于Swin Transformer块构建的。为此，与编码器中使用的补丁合并层相比，我们使用解码器中的补丁扩展层对提取的深度特征进行上采样。补丁扩展层将相邻维度的特征图重塑为更高分辨率的特征图($2\times$上采样)，并相应地将特征维度减少到原始维度的一半。

Patch expanding layer：
以第一个补丁扩展层为例，在上采样之前在输入特征上应用线性层，将输入特征增加到原始维度的二倍($\frac{W}{32}\times \frac{H}{32}\times 8C$ $\Rightarrow$$\frac{W}{32}\times \frac{H}{32}\times 16C$). 然后，我们使用重排操作将输入特征的分辨率扩展为原来的二倍，分辨率缩小为$\frac{1}{4}$($\frac{W}{32}\times \frac{H}{32}\times 16C$ $\Rightarrow$$\frac{W}{16}\times \frac{H}{16}\times 4C$)

skip connection
与U-Net类似，skip-connection用于将来自编码器的多尺度特征与上采样特征融合。我们将浅层特征与深层特征融合在一起，以减少下采样造成的空间信息丢失。后跟一个线性层，连接特征的维度与上采样特征的维度保持相同。

Experiments

Datasets：
Synapse multi-organ segmentation dataset (Synapse): 包含了30个病例，3779张轴向腹部临床图像。其中，18个样本被划分为训练集，12个样本被划分为测试集。并且DSC和HD被用作评估指标来评估Swin-Unet对8个腹部器官(主动脉，胆囊，脾脏，左肾，右肾，肝脏，胰腺，脾脏，胃)进行评价。

Automated cardiac diagnosis challenge dataset(ACDC):ACDC数据集使用MRI扫描仪从不同患者处收集，对于每个患者的MR图片，都标记了左心房，右心房和心肌。数据集被分成70个训练样本，10个验证样本和20个测试样本。使用平均DSC评估Swin-Unet在该数据集上的性能。

评价指标相关拓展知识：

• T1：蓝色部分表示真实脑肿瘤区域(Ground Truth)
• T0：除T1蓝色外其他部分为预测的正常脑区域
• P1：红色部分为预测的脑肿瘤区域(Prediction)
• P0: 除P1红色的其他部分为预测的正常脑区域
• TP： True Positive，被判定为正样本，事实上也是正样本，即$T1\cap P1$
• TN：True Negative， 被判定为负样本，事实上也是负样本，即红色与蓝色以外的区域
• FP： False Positive， 被判定为正样本，但事实上是负样本，即红色中除了蓝色的部分
• FN： False Negative， 被判定为负样本，但事实上是正样本，即蓝色中除了红色部分

DIce系数(Dice Similarity Coefficient, DSC): Dice系数是一种集合相似度度量指标。通常用于计算两个样本的相似度，取值范围是$[0,1]$,分割结果最好时值为1，最差时值为0
$Dice(P,T)=\frac{|P1\cap T1|}{(|P1|+|T1|)/2}\to$ $Dice=\frac{2TP}{FP+2TP+FN}$

代码实现

def dice_coef(output, target):#output为预测结果 target为真实结果
    smooth = 1e-5 #防止0除
 
    if torch.is_tensor(output):
        output = torch.sigmoid(output).data.cpu().numpy()
    if torch.is_tensor(target):
        target = target.data.cpu().numpy()
 
    intersection = (output * target).sum()
 
    return (2. * intersection + smooth) / \
        (output.sum() + target.sum() + smooth)

Hausdorff_95(95% HD): Dice系数对分割出的内部填充比较敏感，而hausdorff distance 对分割出的边界比较敏感


单向： 先找最小，再找最大(先找$x_0$到$Y$中最小距离$x_0{y}_1$, 再找$y_1$到$X$中最大距离，为$y_1{x}_2$,即$y_1{x}_2$为所求的单项HD距离)
双向： 先找单向，再找最大
环境安装

pip install -i https://pypi.tuna.tsinghua.edu.cn/simple numba
pip install hausdorff

python代码实现

import numpy as np
from hausdorff import hausdorff_distance
 
# two random 2D arrays (second dimension must match)
np.random.seed(0)
X = np.random.random((1000,100))
Y = np.random.random((5000,100))
 
# Test computation of Hausdorff distance with different base distances
print("Hausdorff distance test: {0}".format( hausdorff_distance(X, Y, distance="manhattan") ))
print("Hausdorff distance test: {0}".format( hausdorff_distance(X, Y, distance="euclidean") ))
print("Hausdorff distance test: {0}".format( hausdorff_distance(X, Y, distance="chebyshev") ))
print("Hausdorff distance test: {0}".format( hausdorff_distance(X, Y, distance="cosine") ))
 
# For haversine, use 2D lat, lng coordinates
def rand_lat_lng(N):
    lats = np.random.uniform(-90, 90, N)
    lngs = np.random.uniform(-180, 180, N)
    return np.stack([lats, lngs], axis=-1)
        
X = rand_lat_lng(100)
Y = rand_lat_lng(250)
print("Hausdorff haversine test: {0}".format( hausdorff_distance(X, Y, distance="haversine") ))

Experiment results on Synapse dataset