[Transformer]A Multi-Branch Hybrid Transformer Network for Corneal Endothelial Cell Segmentation

from MICCAI2021

Abstract

角膜内皮细胞分割在细胞密度、变异系数等临床指标诊断中至关重要。但是由于角膜内皮的不均匀反射和受试者的运动导致细胞边缘模糊难以分割，细腰更多的细节和上下文信息来解决这一问题。由于卷积的感受野有限因此现有的深度学习方法难以充分利用全局上下文，会丢失许多细节信息。本文提出了一种基于Transformer和体边缘分支的多分支的混合Transformer框架（MBT-Net，Multi-Branch hybrid Transformer Network）。
首先借助卷积模块关注局部特征提取，在通过Transformer和残差连接建立空间、通道和层的长程依赖关系。此外借助体边缘分支来促进局部一致性，提供边缘位置信息。在 自采集数据集TM-EM3000和公共数据集Alisarine上均超过了SOTA结果。

Section I Introduction

角膜内皮细胞异常可能与许多角膜和系统性疾病有关，定量测定角膜内皮细胞密度、变异系数和六角形性对临床诊断具有重要意义。细胞分割是量化上述参数的关键步骤，但是手动分割十分费时费力并且不稳定，因此一种准确、全自动的角膜内皮细胞分割算法对于提升诊断效率和准确性至关重要。
精确分割的主要挑战在于模糊的细胞边界，示例参见Fig 1，需要借助更过的细节和上下文信息来减轻这一问题.
UNet通过收缩路径中的高分辨率特征与上采样输出相结合来实现精确定位。UNet++通过嵌套的密集的skip connection进一步优化来减少encoder和decoder之间特征映射的语义差异。Fabij´ anska首次将UNet用于角膜内皮细胞的分割任务，Fu等人则结合ROI的先验知识和临床参数设计了并行的深度网络。然而由于连续卷积-下采样的感受野有限，不能充分利用上下文信息仍然遗漏了许多细节。
Transformer被证明是建立长程依赖的一种有效方法，Vaswani等人首次提出用于语言翻译的Transformer结构，通过注意力机制建立时间、空间和不同层次之间的全局依赖性。Prajit等人则探索了Transformer在视觉分类任务中的应用，使用自注意力层替换ResNet中的卷积层，但仍然丢失了一部分全局信息。Wang等人通过使用两个独立的轴向注意力快来减少参数和计算量，同时可以在更大范围内计算注意力。

之前有的研究借助边缘信息获得了更好的分割结果，Chen等人提出深度轮廓感知网络。使用多任务学习框架学习物体和轮廓的互补信息，提升了中间特征的识别能力。Chen等人通过学习CNN中间特征的参考边缘图来改进网络输出。Ding等人提出将边界作为一个附加的语义类别，引入单向无环图的边界特征扩展模块引入边界约束，促进类内一致性。

为了保留更多的局部信息，更好的利用全局上下文，本文提出多分支的混合Transformer框架——MBT-Net。首先使用一个混合的残差Transformer特征提取模块来充分发挥卷积和Transformer的优势，提取局部细节和全局语义信息。其中卷积模块负责局部纹理特征的提取，Transformer和残差连接负责建立对空间】时间和level的长程依赖关系。此外还借助边缘和细胞体信息来辅助分割，边缘分支提供精确的边缘位置信息，提升局部一致性。实验结果表明MBT-Net效果优于其他SOTA方法，在两个角膜内皮数据集上取得了更好的性能。

Section II Method

MBT-Net的网络结构参见Fig 2，首先借助Transformer编解码模块从内皮细胞图像提取特征F，每个卷积层包含两个残差快，卷积核大小为3x3.每个Transformer模块包含两个模块，核大小为1x48.
随后将特征解耦为主体和边缘两部分，同时将来自Conv-e1的边缘信息融合进edge feature.最后将边缘特征Ef,主体特征Bf和特征图F融合后做最终的分割预测。三个分支都得到明确的监督。
整个处理流程中卷积层侧重问题特征的提取，保留更多细节；残差连接和Transformer负责在更广泛的空间、通道和层中更充分利用全局上下文信息。边缘视角有助于保留边缘细节，body视角则有助于保证局部一致性。低分辨率的特征dx会与ex特征级联相加后做进一步精炼。



Part 1 Residual Transformer Block

Fig 3展示了残差Transformer模块，包含两个1x1卷积来控制通道数，以及一个高度轴和宽度轴的轴向多头注意力模块，这大大减少了计算量，并且这种设计可以使得Transformer直接学习特征映射图大小范围内的全局空间信息。
轴向自注意力的计算表示为：


对于给定的输入特征图x，通过线性映射获得q,k,v矩阵，然后计算p-o任意两个位置之间的相关性，N1xm表示的是以位置欧威中心的局部1xm的区域，yo则是位置ode注意力输出。

此外，编解码器中的所有模块都使用的残差连接，这样可以让所有输入信号从低层传递高高层，优化信息交互。前向信息传播过程表述为：

Part 2 Body，Edge,Final Branches

为了更好的分割细胞边缘需要借助细胞体信息和边缘信息。body branch提供细胞的整体轮廓和一致性信息，edge branch提供边缘定位信息，来提升细节分割精度。本文将混合Transformer的编解码网络提取到的特征F解耦为Fboft和Fedge两部分，这一步通过卷积层实现。同时将来自encoderConv-e1的低层信息融合到Fedge中，这里进行的是维度变换操作。

最终上述三类特征图融合在一起做最终的细胞分割预测。
这三个分支都受到显式监督，Fig 4展示了三种mask，最终的分割图gt来源于数据集中专家标注的结果，edge mask则是最终分割mask通过canny算子提取到的边缘信息，body mask通过最终分割mask进行反色再进行高斯模糊获得。

Part 3 Loss Function

损失函数表示为：


三个参数分别调整body loss,edge loss和final loss的权重，依次为0.5,0.5,1.2,每一项损失计算的都是交叉熵损失函数。

Section III Experiments

Part 1 Datasets

TE-EM3000 包含184张角膜内皮细胞图像，大小为266x480.为了减少病变和伪影的干扰本文从每张图像选择一个192x192大小的oatch，本文的划分为train:val:test=155:10:19



Alizarine 包含30只猪眼的角膜内皮图像，大小为768x576，本文每张选取10个大小为192x192的patch一共300个patch，train:val=260:40



实施细节：输入大小192x192 lr=2e-4 epoch=100 batch_size=1 pytorch框架

Part 2 Comparison with SOTA methods

本文与LinkNet,DinkNet,UNet,UNet++，TransUNet进行了对比，评价指标有DICE系数，F1 ，Sen,Spe,其中DICE系数和F1score是最能衡量分割性能的指标。




对比结果参考Table 1，可以看到MBT-Net在两类数据集上都达到了总体最优的性能，TM-EM3000数据集上DICE和F1分别为0.747和0.747,本文能取得提升的主要原因在于：

（1）通过Transformer建立了编码器和解码器之间的长程依赖关系；


（2）对整个特征图送入Transformer层减少了语义信息的丢失
；


（3）edge 分支鼓励网络学习更多的一般特征，并提供边缘定位信息。

从Fig 5可以看到在TM-EM3000数据集在左侧细胞边界清晰的情况下，不同方法的分割性能没有太大差异，但是在右侧光照不均匀和边界模糊的情况下本文获得了最佳分割结果，更符合实际情况。

在Alisarine数据集中没有不模糊的边界因此所有分割方法性能都比较好，但是在细节上有不同程度的损失，导致白色箭头部分细胞边缘分割不连续，本文就没有这种情况。




Part 3 Albation Study

Table 2展示了消融实验的结果，本文会逐渐将原网络中的编码器解码器从内到外替换为Transformer结构，可以看到一开始模型捕获了更多语义信息，性能也得到了提高，但是随着越来越多的浅层卷积层被替换为Transformer层模型开始丢失局部信息，导致性能下降，因此本文的2-2-TR在局部细节和全局上下文之间实现了最好的平衡。

本文还研究了transformer和body branch对模型性能的影响，当什么也不使用时性能为0.720,加上body branch后性能提升至0.737，再加上Transformer后性能提升至0.741，分别使得性能提升了2.7%和1.4%。

Section IV Conclusion

本文提出一种用于角膜内皮细胞分割的多分支混合Transformer模型，结合了卷积和Transformer的优势，利用body-edge-branch来促进拒不一致性，提供边缘信息，达到了SOTA结果，尤其在边界模糊的情况下。消融实验的结果表明，设计良好的Transformer结构和body-edge-branch有助于性能的提升。