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
- S A ( z i ) = s o f t m a x ( q i k T D h ) v SA(z_i) = softmax(\frac{q_ik^T}{\sqrt{D_h}})v SA(zi)=softmax(DhqikT)v
-
Transformer encoder处理后的结果会送到Decoder,在解码器部分则采用渐进上采样的方式(PUP),类似SETR的操作
-
首先会将Z^L的输出reshape到原始二维维度,拥有D0个通道;然后使用两个连续的上采样-卷积层来恢复空间分辨率,最终获得不同大小的上采样结果,与CNN的feature map进行特征融合
CNN branch
- 使用了ResNet作为CNN的分支
- 保留前4层的输出,然后将他们与Transformer的结果融合获得融合后的特征提取
BiFusion Module
-
融合CNN和Transformer提取的特征
-
通道注意力
-
t ^ i = C h a n n e l A t t n ( t i ) \widehat{t}^i = ChannelAttn(t^i) t i=ChannelAttn(ti)
-
SE-Block
-
-
空间注意力
-
g ^ i = C h a n n e l A t t n ( g i ) \widehat{g}^i = ChannelAttn(g^i) g i=ChannelAttn(gi)
-
CBAM block 作为空间滤波器,增强局部细节,抑制无关区域,低层次的CNN特征会存在噪声
-
-
3x3卷积
-
b ^ i = C o n v ( t i W 1 i ⋅ g i W 2 i ) \widehat{b}^i = Conv(t^iW^i_1 \cdot g^iW^i_2) b i=Conv(tiW1i⋅giW2i)
-
Hadamard product,矩阵点乘,对两个分支的特征之间的细粒度交互进行建模
-
-
残差连接
- f i = R e s i d u a l ( [ b i , t i , g i ] ) f^i = Residual([b^i,t^i,g^i]) fi=Residual([bi,ti,gi])
-
利用attention-gate(AG)来生成最终的分割结果
- f i + 1 = C o n v ( [ U p ( f i ) , A G ( f i + 1 , U p ( f i ) ) ] ) f^{i+1} = Conv([Up(f^i), AG(f^{i+1}, Up(f^i))]) fi+1=Conv([Up(fi),AG(fi+1,Up(fi))])
-
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