GT U-Net:一种UNet型的Group Transformer用于齿根分割
Abstract
为了进行精确的根管治疗,一个基本步骤是需要对抠像X光图像的齿根进行分割,因为齿根的位置对根管治疗至关重要。但是齿根的界限很模糊使得齿根分割十分具有挑战性。
本文提出一种新的端到端的类UNet的Group Transformer网络(GT UNet)用于齿根分割,该网络保留了UNet的基本架构但是将每一组encoder和decoder都用一个Group Transformer代替;利用group和bottleneck可以大大减少计算量和计算成本。
因此本文的GT UNet是一个CNN-Transformer混合架构,并且不需要预训练。为了进一步优化。本文还进一步提出一种形状敏感的傅里叶描述器(FD)损失函数来利用形状这一先验知识。实验结果显示本文提出的GT-UNet在齿根分割和视网膜分割上均达到了SOTA。
Section I Introduction
全球约有7.43亿人患有严重的牙周炎,牙周炎是世界公认第六大常见疾病。牙周炎的常规治疗方法就是根管治疗,如果处置不当会引起一系列后果。因此齿根的形状对评估只管重要,但是鉴于以下原因使得齿根分割十分具有挑战性:
(1)齿根边界十分模糊,其周围组织强度预期相似,如Fig 1(a)所示;
(2)X光中齿根会被其他骨骼或组织遮挡,参考Fig1(b);
(3)X光图片可能存在过曝或欠曝,参考Fig1©。
为了解决上述问题,Zhao等人提出一种两阶段的基于注意力的分割方法,通过自动捕捉牙齿区域来缓解强度分布相近的问题;Lee等人则采用在mask R-CNN上微调的方法;但是上述方法并不能有效解决分割边界模糊的问题。
Chen提出的MSLPNet通过使用多尺度的结构相似性损失来提升牙齿分割效果,还有的使用UNet+DFM来学习像素体现的方向性,以此来约束分割结果。虽然这些工作取得了一定的提升,但是都是基于CNN做的,受限于CNN的局部性无法较好的处理全局特征。
Transformer可以较好的捕捉长程依赖关系缓解这一问题。比如TransUNet提出使用基于Transformer的编码器用来分割,encoder是一个12层的Transformer;但是ViT需要依赖在大规模数据集上预训练的结果,如果数据集有限效果就不那么尽如人意。BoTNet则是将卷积与Transformer结合来作为分割backbone,但是又受限于Transformer的计算复杂度,BoTNet只能将ResNet最后几层的部分卷积替换成了Transformer。
为了缓解上述问题,本文提出了GT UNet,是一种CNN-Transformer混合的架构但是不依赖于预训练,通过group 结构和bottleneck结构有效减少了整体计算量。持此之外还提出了FD loss可以充分利用形状这一先验知识。
本文的贡献总结如下:
(1)GT Unet保持了UNet的网络架构,通过引入Transformer缓解convolution的局限性;
(2)本文还设计了一个分组结构和瓶颈结构,从而减少Transformer的计算量,使其适用于分割这种密集预测情况;
(3)对于齿根分割任务,本文还提出一种形状敏感的傅里叶损失来解决边界模糊的问题。
Section II Method
Part 1 U-Net like Group Transformer Network
Fig 2展示了GT UNet的整体结构,是一个U型网络,使用了group Transformer。对于Transformer中的self-attention结构使用的是多头注意力,头数设置为4.
Group Transformer由 skip connectiong,grouping module,3x3卷积,MHSA,merging module组成。其中skip connection是为了防止梯度消失,同时保留low-level信息。MHSA需要n^2d的内存,因此使用了group module和3x3卷积来减少计算量。
Grouping Structure and Bottleneck Structure
Transformer最先被用于NLP领域,但鉴于其可以有效的通过non-local运算捕获长程依赖,在CV领域也引起了人们极大的研究热情。但是鉴于医学图像和自然语言、自然图像还存在很大差别,很难直接将Transformer用于医学图像。第一自然语言中的单词都是有限的,不像图像像素会随着图像尺寸指数增长。
因此本文结合医学图像的特点设计了group Transformer。
对于HxWxC的输入,MHSA的计算量是:
而Group Transformer的计算量是:
其中φ是channel scaling factor(通道缩放因子),位于瓶颈部分。而输入的hw则取决于图像尺寸和具体任务。
随着网络不断加深,感受野逐渐扩大通过提取不同GT单元的特征可以获得全局的依赖关系。
Multi-Head Self-Attention
Fig 2左侧还展示了MHSA的详细结构,Rh,Rw表示的是采用相对距离的高度和宽度位置编码,最终的注意力计算结果是:qkT+qrT。
Hybrid Structure of Convolution and Transformer
Transformer可以有效提取全局特征,凡是缺少局部性和全局不变性,这是CNN的特点。
因此3x3卷积不仅在本文的瓶颈部分使用,还与Transformer互为补充。
Part 2 Shape-sensitive Fourier Descriptor Loss Function
一般齿根的形状都比较相似,因此在损失函数中添加形状信息可以更好的指导网络模型进行齿根的分割。
(xm,ym)表示齿根的边界,共N个像素点,则其形状可以表示成以下复数:
傅里叶变换后表示为:
傅里叶描述是对某封闭形状的定量表示,包含起点、尺度、位置和旋转信息,因此可以用来衡量预测的边界和GT之间的差异:
本文原始的损失函数是交叉熵损失函数,此外还添加了位置敏感的FD loss,最终的损失函数表述为:
对BCD和FD施加同样的权重,β由Z(k)数量级决定,本文设置为10。
Section III Experiment
Part 1 Tooth Root Segmentation
数据集 齿根数据集包含248张X光图像 提供了3位专家的分割结果。
与GT UNet进行对比的是UNet,Attention UNet和TransUNet。
训练时使用了4块2080Ti
数据增强有:Random Crop,Random Rotation,Axial Flipping
训练相关:Adam Optimizer,training epoch = 200 batch_size = 12
输入图像 256 x 256
Transformer的h x w = 8 x 8
评估指标:
Accu,Se,Sp,JS, Dice socre
Table 1展示了本文与其他网络框架做齿根分割的精度对比,可以看到本文的GT UNet获得了最高的Acc,JS,Dice分数。
Fig 3可视化了3张分割结果,尤其之后两列可以看到FD loss对最终的分割节骨有显著提升,说明FD loss可以有效的借助形状这一先验知识,在周围组织类似或边界模糊时进行精确的分割。
Part 2 DRIVE segmentation
为了进一步验证GT UNet的有效性本文还在DRIVE数据集上进行了测试。
数据集 40张彩色眼底图像 分辨率 565 x 584 testing:training = 20:20
Transformer h x w = 4 x 4 φ = 2
由于FD loss是处理具有相似形状的分割任务,因此没有用在DRIVE分割上
。
并且本文原图是切成patch后作为输入的,patch size = 64
Table 2展示了DRIVE数据集上GT UNet与其他网络的精度对比,可以看到至少比其他网络在Acc上提升了0.5% F1提升了2.5% ,SE提升了3.2%。
Fig 4展示了部分分割结果,可以看到绝大部分血管得到了精确分割,除了一小部分微小血管。
Section IV Conclusion
本文提出一种端到端的类U型的Group Transformer结构(GT UNet)用于医学图像分割任务。
GT UNet将CNN与Transformer结合使用,从而不需要依赖于预训练。同时通过瓶颈结构和grouping结构的设计显著减少了Transformer的计算复杂度。
此外提出的FD Loss可以有效借助形状这一先验知识来提升边界模糊但形状相似的分割任务。
在齿根分割和血管分割任务上GT UNet都显示了优异的分割性能,展示了在其他医学图像分割任务的良好的应用前景。