UNETR: Transformers for 3D Medical Image Segmentation

UNETR: Transformers for 3D Medical Image Segmentation

背景：

1、具有跟踪和扩展路径的全卷积神经网络（FCNNs）（如编码器和解码器），近年来在各种医学图像分割应用中显示出了突出的地位。在这些体系结构中，编码器通过学习全局上下文表示来扮演一个不可或缺的角色，该全局上下文表示将进一步用于解码器的语义输出预测。
2、卷积层作为FCNN的主要构造块，其局部性限制了此类网络学习长时空间依赖的能力。
3、在NLP领域，基于transformer的模型已经在各种任务中实现了最先进的基准。transformer中的自我注意机制使他们能够动态地突出词序列的重要特征并学习其长时依赖关系。这一概念最近被扩展到计算机视觉，引入了视觉transformer(ViT)。在ViT中，一张图像被表示为将用于直接预测用于图像分类的类标签的图片块嵌入序列。

网络结构图：

1、整体架构：类似于u-net网络，将传统的图像分割的resnet主干网络替换成transformer，提取transformer中不同层的序列表示，并通过跳越连接将它们与解码器合并。

2、编码器 创建一个三维输入（h×W×D×C）的一维序列，通过将其划分为flatten的相同大小的非重叠图片块。使用一个线性层将平面化的图片块投影到一个K维嵌入空间中，该空间在整个transformer中保持不变。此外，为了保留所提取的图片块的空间信息，在投影图片块上加入一个一维可学习的位置嵌入，具体的学习位置可见ViT。在嵌入层之后，利用由多头自注意力机制(MSA)和多层感知器(MLP)子层组成的transformer块堆叠，MLP包括两个具有GELU激活函数的线性层，MSA块包括n个平行的自注意力(SA)头。SA块是一个参数化函数，它学习输入序列(z)中的两个元素及其query和key表征之间的相似性。

3、编码器的多个分辨率的特征与解码器合并，从transformer中提取序列特征，将其装换为统一大小的张量。利用连续的3×3×3的卷积层，然后进行批处理归一化，将大小重塑的张量从嵌入空间投射到输入空间。在编码器的bottleneck（即transformer的最后一层输出），应用一个去卷积层到变换后的特征图，以2倍的增加其分辨率。将调整大小的特征映射与前一个transformer输出的特征映射连接起来，将它们传递到连续的3×3×3卷积层中，并使用反卷积层对输出进行上采样。对于直到原始输入分辨率的所有其他后续层，重复此过程。最终输出被传递到带有softmax激活函数的1×1×1卷积层，以生成逐像素的语义预测。

4、损失函数是dice和交叉熵项的组合。

其中I是体素的数量J是类的数量Yi，j和Gi，j分别表示在体素i处的类j的概率输出和单热编码的groundtruth。

实验结果：

数据集
为了涵盖各种对象和图像模态，采用我5倍交叉验证数据拆分的实验，采用了来自MSD挑战[22]的任务1（脑肿瘤MRI分割）和任务9（脾CT分割）的数据集。 对于任务1，将具有神经胶质瘤分割坏死/活动性肿瘤和水肿的groundtruth标签的484个多模式多站点MRI数据（FLAIR，T1w，T1gd，T2w）的整个训练集用于模型训练。 任务1的分辨率/间距统一为1.0×1.0×1.0 mm3。 对于任务9，使用带有脾脏注释的41个CT量。 任务9中卷的分辨率/空间范围为0.613×0.613×1.50 mm3至0.977×0.977×8.0 mm3 在预处理期间，将所有体积重新采样到1.0mm的各向同性体素中。

对于具有MRI图像的任务1，使用z分数归一化对体素强度进行了预处理。 对于具有CT图像的任务9，图像的体素强度根据总前景强度的第5个和第95个百分位数归一化为[0,1]范围。 此外，任务1的问题被公式化为具有4通道输入的3类分割任务，而任务9被公式化为具有单通道输入的二进制分割任务（前景和背景）。 对于任务1和任务9，我们分别以[128,128,128]和[96,96,96]的体积随机采样输入图像。 前景/背景的随机色块以1：1的比例进行采样。
实验结果

Table1：脑肿瘤分割任务的交叉验证结果。先划分成五份数据集交叉验证，对于每个分割，提供了三个类的平均骰子分数。DSC1、DSC2和DSC3分别表示整个肿瘤(WT)、增强肿瘤(ET)和肿瘤核心(TC)的平均骰子得分。

Table 2.脾脏分割任务的交叉验证结果。对于每个分割，我们提供了前景类的平均骰子分数。

Fig. 2. (a) Ground Truth. Outputs of : (b) UNETR. © SegResNet. (d) UNet.
分析：
1、在训练时，没有使用任何预先训练过的变压器模型，因为预训练没有显示出任何性能改善。
2、定量评估：在表1中，比较了UNETR与基于CNN的网络在脑肿瘤分割任务中的性能。在所有语义类中，UNETR平均比最接近的基线高出2.28%。特别是，UNETR在分割肿瘤核心(TC)方面的表现要好得多。在表2中，比较了UNETR与基线的性能在脾脏分割的任务中的表现。类似地，UNETR的性能比最接近的基线高出至少1.11%。
3、定性评估：在图2展示了UNETR的模型的分割输出以及其他基线的定性结果。对于脑肿瘤的分割，UNETR的模型在捕获肿瘤的细粒度结构细节方面显示出了更好的性能。

结论

在本文中，引入了一种新的基于transformer的医学图像3D语义分割架构，并提出将其重新表述为一个一维序列对序列的预测。我们提出使用纯transformer来提高模型学习长期依赖关系并在多个尺度上有效捕获全局上下文表示的能力。在CT和Mr图像模式中验证了UNETR在MSD数据集的体积脑肿瘤和脾脏分割任务上的有效性，我们的基准在这些模式上一致显示了良好的性能