学习 TransFuse: Fusing Transformers and CNNs for Medical Image Segmentation 医学分割

Fusing Transformers and CNNs for Medical Image Segmentation

Abstract

Medical image segmentation - the prerequisite of numerous clinical needs - has been significantly prospered by recent advances in convolutional neural networks (CNNs). However, it exhibits general limitations on modeling explicit long-range relation, and existing cures, resorting to building deep encoders along with aggressive downsampling operations, leads to redundant deepened networks and loss of localized details. Hence, the segmentation task awaits a better solution to improve the efficiency of modeling global contexts while maintaining a strong grasp of low-level details. In this paper, we propose a novel parallel-in-branch architecture, TransFuse, to address this challenge. TransFuse combines Transformers and CNNs in a parallel style, where both global dependency and low-level spatial details can be efficiently captured in a much shallower manner. Besides, a novel fusion technique - BiFusion module is created to efficiently fuse the multi-level features from both branches. Extensive experiments demonstrate that TransFuse achieves the newest state-of-the-art results on both 2D and 3D medical image sets including polyp, skin lesion, hip, and prostate segmentation, with significant parameter decrease and inference speed improvement

Introduction

CNN

• CNN在诸多医学图像分割任务中取得了优异性能，如多器官分割、肝部病变分割、脑补分割等，显示出CNN在建模特定任务特征表示的强大能力
• 缺点：CNN的一大问题是在捕获全文上下文信息方面缺乏效率，如果通过堆叠、扩大感受野的方式则需要连续的下采样-卷积操作，使得网络结构十分深，这一过程也会导致局部信息的丢失，而拒不细节信息对于密集预测任务也十分重要

Transformer

• 缺点：transformer也有自身局限，就是无法较好的建模细粒度特征，尤其是对于医学图像，细节特征十分重要,其在建模局部信息时缺少空间的归纳偏执

TransFuse

• 可以有效的捕获低层次空间特征和高级语义特征
• 无需搭建深层次网络就可以减轻梯度消失、特征无法有效复用等问题
• 模型效率和推理速度有了很大的提升，并且部署在云端或终端的效率也有很大提升

Proposed Method

two parallel branches

• CNN branch

  • 增加感受野，将特征从局部到全局

• Transformer branch

  • 从全局自注意力开始，最后恢复局部细节

• two benefits

  • 有效的发挥CNN和Transformer各自的优势，在不需要搭建深层次网络的同时，有效捕获全局信息，还能保持精确的low-level信息
  • BIfusion在提取特征的过程中，同时利用了CNN和Transformer的不同性质，从而融合的更好

Transformer branch

• HXWX3的输入切成patch（patch number= 16）;随后patch经过线性映射并展平，在送入Transformer之前还会嵌入可训练的位置信息。

• Transformer中包含L层MSA和MLP

  • $$SA(z_i)=softmax(\frac{q_i{k}^T}{\sqrt{D_h}})v$$

• Transformer encoder处理后的结果会送到Decoder，在解码器部分则采用渐进上采样的方式（PUP），类似SETR的操作

• 首先会将Z^L的输出reshape到原始二维维度，拥有D0个通道；然后使用两个连续的上采样-卷积层来恢复空间分辨率，最终获得不同大小的上采样结果，与CNN的feature map进行特征融合

CNN branch

• 使用了ResNet作为CNN的分支
• 保留前4层的输出，然后将他们与Transformer的结果融合获得融合后的特征提取

BiFusion Module

• 融合CNN和Transformer提取的特征

  • 通道注意力

    • $${\hat{t}}^i=ChannelAttn(t^i)$$

    • SE-Block

  • 空间注意力

    • $${\hat{g}}^i=ChannelAttn(g^i)$$

    • CBAM block 作为空间滤波器，增强局部细节，抑制无关区域，低层次的CNN特征会存在噪声

  • 3x3卷积

    • $${\hat{b}}^i=Conv(t^i{W}_1^i\cdot g^i{W}_2^i)$$

    • Hadamard product，矩阵点乘，对两个分支的特征之间的细粒度交互进行建模

  • 残差连接

    • $$f^i=Residual([b^i,t^i,g^i])$$

  • 利用attention-gate(AG)来生成最终的分割结果

    • $$f^{i+1}=Conv([Up(f^i),AG(f^{i+1},Up(f^i))])$$

Loss Function

• 加权mIoU损失
• BCE损失

Experiments and Results

Data Acquisition

• 息肉分割

  • Kvasir , CVC-ClinicDB , CVC-ColonDB ,EndoScene and ETIS

    • 352×352

• 皮肤损伤

  • 2017 International Skin Imaging Collaboration skin lesion segmentation dataset (ISIC2017)

    • 192×256

• 髋关节分割

• 前列腺分割

  • volumetric Prostate Multi-modality MRIs from the Medical Segmentation Decathlon

    • 320 × 320

Implementation Details

• TransFuse-S

  • ResNet-34 (R34) and 8-layer DeiT-Small (DeiT-S)

• TransFuse-L

  • Res2Net-50 and 10-layer DeiT-Base (DeiT-B)

• TransFuse-L*

  • ResNetV2-50 and ViT-B

Evaluation Results

• Results of Polyp Segmentation

  • mean Dice

    • Dice系数是一种集合相似度度量函数，通常用于计算两个样本的相似度，取值为[0,1]。

  • mIOU

  • Table1展示了息肉分割的对比结果，TransFuse-S/L与其他CNN网络相比均达到了SOTA，而且参数量比PraNet等减少了20%,实时性也更好

  • 而预训练过的TransFuse-L*的性能也优于SETR和TransUNet

• Results of Skin Lesion Segmentation

  • 在皮肤损伤分割实验中性能指标是Jaccard指数和Dice系数以及逐像素点的accuracy。
  • Table 2中，TransFuse性能优于UNet++,同时不需要任何预处理或后处理，而UNet++需要使用预训练的R34作为主干

• Results of Hip Segmentation

  • Hausdorff Distance (HD)
  • Average Surface Distance (ASD)
  • Table 3：髋关节分割的对比结果，主要需要分割腓骨、左股骨、右股骨。与UNet++和HRNet相比，TransFuse在HD和ASD这两个指标参数上均更加优秀，充分证明了本文提出的TransFuse可以有效捕获精细的结构，生成更加清晰准确的轮廓

• Results of Prostate Segmentation

  • nnUNet是目前在前列腺分割排名第一的分割网络
  • Table 4：TransFuse与nnUMet的对比结果。可以看到与nnUNet-3d相比，TransFuse-S不仅性能更好，而且参数量减少了41%,吞吐率增加了50%。

• Ablation Study

  • Table 5：并行分支的消融实验，Table 6展示的是BiFusion的消融实验。
  • 可以看到两个分支选择CNN搭配Transformer性能是最佳的，BiFusion结合空间注意力、通道注意力、计算内积都会提升性能。
    

Conclusion

评审意见

官方评审意见：https://miccai2021.org/openaccess/paperlinks/2021/09/01/496-Paper0016.html

评分：6，6，6

参考：http://t.csdn.cn/d2JR8

​