基于Swin-Unet改进的医学图像分割算法
Improved Medical Image Segmentation Algorithm Based on Swin-Unet
DOI: 10.12677/airr.2024.132037, PDF, HTML, XML, 下载: 13  浏览: 38  科研立项经费支持
作者: 康家荣*, 王 元#:广西民族师范学院,数理与电子信息工程学院,广西 崇左;邵鹏飞:浙江万里学院,信息与智能工程学院,浙江 宁波
关键词: 神经网络医学图像分割算法Neural Network Medical Image Segmentation Algorithm
摘要: 在过去的十几年时间里,基于CNN的神经网络在医学图像分割领域取得了突破性进展。尤其是以unet为代表的u形网络架构和跳跃连接被广泛应用于一系列的医学图像任务。由于CNN的内在局限性,不能够很好的获取到全局和远程语义信息交互。由于腹部器官复杂,容易发生形变、边缘模糊、体积小等原因导致分割比较困难。因此在Swin-Unet的基础上改进,首先末端编码器与首个解码器之间引入多尺度模型提取模块,增强不同形状大小信息特征提取。其次将最后两个编码器和解码器的swin transformer Block引入残差机制来缓解模型深度带来的梯度弥散现象。并且在最后的编码器末端引入ASSP模块获取多尺度细节信息。最后,在跳跃连接中引入通道注意力机制(CAM),可以让模型强化重要信息特征通道,弱化特征不相关通道,最后达到有效提高模型分割精度和准确度的效果。
Abstract: In the past ten years, CNN-based neural networks have made breakthrough progress in the field of medical image segmentation. In particular, the U-shaped network architecture and skip connections represented by unet are widely used in a series of medical image tasks. Due to the inherent limitations of CNN, global and remote semantic information interaction cannot be well obtained. Due to the complexity of abdominal organs, which are prone to deformation, blurred edges, and small size, segmentation is difficult. Therefore, based on Swin-Unet, a multi-scale model extraction module is introduced between the terminal encoder and the first decoder to enhance the extraction of information features of different shapes and sizes. Secondly, the residual mechanism is introduced into the swin transformer Block of the last two encoders and decoders to alleviate the gradient dispersion phenomenon caused by the depth of the model. And the ASSP module is introduced at the end of the final encoder to obtain multi-scale detailed information. Finally, introducing the channel attention mechanism (CAM) in the skip connection allows the model to strengthen important information feature channels and weaken feature irrelevant channels, and finally achieve the effect of effectively improving the model segmentation precision and accuracy.
文章引用:康家荣, 邵鹏飞, 王元. 基于Swin-Unet改进的医学图像分割算法[J]. 人工智能与机器人研究, 2024, 13(2): 354-362. https://doi.org/10.12677/airr.2024.132037

1. 介绍

精确的描述病变在疾病预防后的评估、疾病量化和治疗效果评价体系当中起着关键的作用。通过专业医护人员手动进行医学图像分割可以准确可靠的支撑起病理的诊断。但是在标准、智能临床技术当中精确和稳定的自动分割能够起到基石的作用。自动分割提供了高效率全天候的优质医疗服务,大大节约了劳动成本,减少误诊率。因此智慧临床医疗对于基于人工智能的医学图像自动分割技术有很大的依赖性。

过去的十几年时间里,基于卷积网络在分割任务中大放异彩取得了非常瞩目的成就,其中典型的代表有全连接神经网络(FCN)、U-net [1] 、注重跳越连接的增强UNet++ [2] 等。伴随何凯明等人推出的残差概念,ResUnet [3] 在此期间出现,主要是为了防止过度卷积带来的特征图谱的梯度弥散,而保证最差结果为保留的上一层特征图。

结合两种受U型编码和解码影响的架构,提出的DAResUnet [4] 中包含了残差块(Res-Block)和双头注意力两个模块。在u型结构中条约连接为解码器提供了初始特征,弥补了编解码器之间的语义差距。然而,感知领域卷积的局限性和偏差会影响分割的准确性,无法建立全局上下文的远程依赖关系。

对于transformer [5] 最早出现在2017年的NLP领域中的序列到序列的模型开发当中,并且得到了迅速的发展,从NLP到视觉领域几乎无处不在。例如:在transformer视觉领域的开篇之作VIT [6] 通过将一幅图像分成多个小图像模块,然后将这些小图像模块的先验嵌入序列反馈给网络作为transformer的输入,获得了一个很好的效果,Transformer进一步提高了分割精度。在VIT的启发下,TransUNet [7] 进一步将VIT和U形网络在医学图像分割领域的优势结合起来。具体来说,引入带有transformer的编码器对图像进行处理,采用了传统的CNN和跳跃连接进行精确的上采样特征恢复。Swin-unet [8] 将unet和swin-transformer block模块像融合,取得了良好的分割改良效果。

上述基于医学图像分割的工作当中,对u-net结构网络和transformer架构的改进,取得了一定程度的进步,但是也存在一定的局限性:

1) 在传统的u-net医学图像分割模型中,单一的卷积处理导致全局上下文语义和远程交互信息不足。

2) 在transformer引入到视觉领域之后,transformer-block中的注意力机制便开始注入到u形网络当中,然后依然存在精确定位和细节特征的不足。

3) Transformer在拥有强大的特征提取能力的同时能够保留全局信息。为了进一步激活Transformer的潜力,如何将CNN和transformer更高效的融合是一个挑战。

4) 对于U-net模型而言跳跃连接是一个非常重要的部分,跳跃连接弥合了编码器和解码器之间的语义差异。优化这些跳跃连接是目前的工作之一,使用解码器能够检索和返回更精确的特征映射,从而增强模型的鲁棒性。

鉴于上述问题,我们本次提出了一种改进型swin-unet。本论文着重于强调U-net与swin-transformer-block模块的集成,为了提高swin-transformer的特征精度,我们引入Resswin-transformer模块,防止梯度爆炸。

考虑到编码器–解码器和跳过连接的结构,我们使用通道注意力模块(CAM)对三层跳过连接中编码器传递的特征进行特征优化,这有助于解码器–编码器结构进一步减少语义划分,从而产生统一的特征表示。这种融合方法可以充分利用不同方面、全局和局部特征,使它们相互补充,从而增强模型的医学图像分割性能。

与传统的u型结构一致,编码器用于获取上下文特征和跳过连接,解码器用于统一特征的融合。与大多数基于cnn的编码器不同,我们在编码中加入了swin-Transformer模块,以进一步提取远距离的上下文特征,从而大大提高编码能力。

通过引入ASPP [9] 模块以利用不同膨胀的空洞卷积并行捕获图像的不同尺度特征,获得不同的感受野特征,并将其进行融合,充分考虑了上下文信息,同时提高了网络对细节特征的提取能力。

为了提高解码器的性能,我们进一步完善了跳过连接所传达的特征。由于这些增强,解码和医学图像分割能力都得到了显著增强。

2. 相关工作

最近,注意力机制在U-net架构中越来越受欢迎 [1] 。例如,Attention U-net采用注意机制来增强胰腺的定位和分割性能 [10] ;DAResUnet在U-net中集成了双重注意和残余机制 [4] ;注意Res-UNet探索软注意替代硬注意 [11] ;Sa-unet在U-net中引入了空间注意机制 [12] 。在此基础上,TransUNet创新地将Transformer和U-net元素结合在一起 [7] 。在TransUNet的基础上,TransU-Net++将注意力机制集成到跳过连接和特征提取中 [13] 。swin-unet [8] 通过用swin – transformer [14] 替换U-net中的每个卷积块来改进。DS-TransUNet提出将tif模块合并到跳跃连接中以改进模型 [15] 。AA-transunet利用块注意力模型(CBAM)和深度可分离卷积(DSC)进一步优化TransUNet [16] 。TransFuse使用双重注意力Bifusion block和AG融合了CNN和Transformer两个不同部分的特征 [17] 。许多注意力机制已经被添加到U-net和TransUNet模型中,但需要进一步的探索。与之前的方法不同,我们的实验将通道注意机制和swin-Transformer block和ASSP模块引入了传统的方法u形编码器–解码器和跳过连接,产生了有希望的结果。

3. 改进的Swin-Unet腹部多图像分割方法

3.1. 架构概述

本文所提出的改进型Swin-Unet的思想架构如图1所示。Swin-Unet主要由编码器、瓶颈、解码器和跳跃连接四个部分组成。主要改进部分如下所示。

Figure 1. Improved swin-unet architecture diagram

图1. 改进型swin-unet架构图

① 将ResSwin-Transformer Block替换到Swin-Unet第三第四个编码器-四个解码器中Swin Transformer Block框架中的新架构。这种方法s高了编码器–解码器结构的灵活性和功能性,从而提高了医学图像分割任务的性能。

② 通过在最后一个编码器尾部添加ASPP,该机制可以增强解码器层的特征提取能力,捕获图像的不同尺度特征和感受野特征将其进行融合,充分考虑了上下文信息和图片细节特征的提取能力。丰富了ResSwin-Transformer层中编码器的功能U-net结构。

③ 我们通过在每一层中加入通道注意力残差模块(CAM Block)来增强跳过连接的有效性,这是一种经过消融研究证实的改进,它可以更准确地向解码器传递特征并提高图像分割性能。

我们提出的改进的Swin-Unet方法在多个医学图像数据集上取得了最先进的性能,证明了我们的方法的有效性及其对推进医学图像分割的贡献。

3.2. 数据预处理

由于医学图像具有隐私性,人工标注的过高成本导致收集的数据集合比较少,但是对于卷积神经网络,如果训练诗句样本少容易出现过拟合,使得训练的模型不能很好的分割,假如丰富数据集的数量和样式,那么训练之后的深度学习网络将进一步提高其分割鲁棒性和精度。

在图像数据集扩展的过程中,常用的方式有旋转、缩放、椒盐噪声、高斯噪声、亮度处理镜像、降低像素模糊处理、伽马增强。在本次论文中还添加了平移处理和随机剪切等方式对数据集合进行增强处理。扩充了训练数据的数量。

此外,本次论文的数据集合都属于CT图片数据,CT图片数据的每个层级像素灰度值反映出该切片的某些物理属性并且是定量的,因此要对数据进行归一化处理,保证数据集的标准。

3.3. 空洞空间金字塔池化

大部分医学图像分割是基于编码器–解码器的u形网络结构,通过下采样达到多尺度特征提取、上采样恢复特征定位,但是很容易丢失重要的特征语义信息,上下文语义信息不能远程交互,在精度和准确度上不够完善。DeepLab V2 [9] 方法中首次提出ASPP模块,该模块能够捕获不同尺度的感受野特征并进行融合,充分考虑了上下文信息从而提升细节特征能力提取的能力。

ASPP模块包含了5条并行分支,这五条分支分别为 ① 1个1 × 1的卷积 ② 空洞率为3的3的3 × 3空洞卷积 ③ 空洞率6的3的3 × 3空洞卷积 ④ 空洞率为12的3的3 × 3空洞卷积 ⑤ 全卷积的池化。通过5分支特征结果进行concat拼接最后通过1 × 1卷积得到相同尺寸大小的输出。空洞卷积的使得在不增加模型参数数量的同时增大了特征图像的感受野 [18] 。最后通过双线性插值来还原图像至原始尺寸大小。ASPP的模块的整体结构如图2所示。本次改进Swin-Unet中将该模块添加到最后编码器的末端来充分提取多尺度信息和增大感受野,帮助解码器更好的恢复细节信息。

Figure 2. Atrous convolution multi-size pooling module-ASSP block diagram

图2. 空洞卷积多尺寸池化模块-ASSP框图

3.4. 残差Swin Transformer模块

网络模型层数不断的增加伴随而来的是模型过度拟合和网络退化等问题,这些问题限制了网络模型精度和准确度的提升。He等人提出ResNet模型 [19] 来缓解过拟合和退化问题,

主要思想是通过一条从数据起始端到输出端前一节点的快捷连接,将输入与输出进行叠加,在不增加数量的前提下,有效缓解了网络梯度弥散的现象,丰富了语义特征,改善了识别精度和准确度。

引入残差思想叠加swin-transformer Block模块,用一条快捷有向连接用于swin-Transformer block的头尾连接,通过恒等映射阻止梯度消失的现象。在本次论文中将每一对带有残差连接的W-MSA和SW-MAS模块对整合成一个Res-Swin Transformer Block模块。具体模块结构如图3所示。

Figure 3. ResSwin-Transformer-Block structure diagram

图3. ResSwin-Transformer-Block 结构图

3.5. 残差注意力机制

为了弥补上下采样丢失的信息、丰富细节特征,在编码器个解码器同层水平方向添加跳跃连接,使得编码器低层次的特征能够融合解码器高层次的特征。然而低层次特征当中夹杂大量冗余的信息,又缺乏有效的特定语义特征,可能会影响到医学图像分割的精确度。

为此,本文将残差通道注意力机制融入到跳跃连接中,可以强化主要的通道特征区域,减弱冗余的特征通道,从而提高腹部多器官CT图像分割的准确率。

每一个通道的特征图谱都可以认为是一个特定物体类的信号特征响应。不同的语义通道相互关联,我们可以通过强调通道相互依赖的特征映射和改进特定的特征表示来构建一个通道注意力模块用于增强重要的通道信息弱化不重要的特征通道信息有助于图像分割的精度。

通道注意力模块结构如图X所示。我们首先将原始特征A∈RC×H×W展开为RC×N,然后将A与其转置进行矩阵乘法。随后,我们应用一个softmax层,得到通道关注图X∈RC×C

x i j = exp ( A i A j ) i = 1 c exp ( A i A j ) (1)

这里,xji测量第i个通道对第j个通道的影响。接下来,我们在X和a的转置之间执行矩阵乘法,将结果重塑为RC×H×W。然后,我们将结果乘以一个尺度参数β,并与a执行逐元素求和运算,得到最终输出E∈RC×H×W

E j = β i = 1 N ( x j i A i ) + A j (2)

和α一样,β也是通过训练习得的。CAM在提取通道特征时,每个通道的最终特征是所有通道与原始特征的加权和,从而赋予CAM强大的通道特征提取能力。

当然对于跳跃连通道注意力机制CAM,在没有增加模型参数量的同时,为了保留原来文本信息并且防止模型出现过拟合问题,在跳越连接的输入到输出之间映入残差机制,构成残差注意力机制模块(CAM),如图4所示。

Figure 4. Channel Attention Module-CAM

图4. 通道注意力模块-CAM

4. 实验

4.1. 数据集和评价指标

本实验的数据主要来自公开数据集Synapse通过本论文数据增强方法处理之后的数据和本地区某三甲医院临床医学腹部扫描图片数据,总共包含了四十例患者数据共计4500张。主要的评估参数指标有HD (Hausdorff Distance)距离和DSC (Dice Coefficient)两个,评估了CT图像的整体评分和8个单独器官(主动脉、胆囊、左肾、右肾、肝脏、脾脏、胰腺和胃)的评分,从数据上展示模型性能。

4.2. 实验细节

本文实验均在python3.8,pytorch 2.2,cuda11.8,单个GPU为24 GB显存的NVIDIA RTX3090显卡。我们使用随机梯度下降法SGD训练这个模型,设置学习率0.01,训练的epoch为150,默认batch大小为48,动量设置为0.9,权重衰减le-4。本次实验的数据集为增强版Synapse和某三甲医院临床腹部多器官CT图片。使用数据集用于训练和测试数量上的分布为3:1,也就是百分之七十五用于训练,百分之二十五用于测试,这样可以保证模型因为训练不足而过拟合。本次训练的损失函数为Loss = 1/2 × Cross – EntropyLoss + 1/2 × DiceLoss。

4.3. 实验结果

将本文提出的方法跟近几年主流方法在synapse数据集的分割结果指标对比分析,实验结果如表1所示,实验分割效果如图5所示。

Table 1. Experimental DSC results of different methods on Synapse data set

表1. 不同方法在Synapse数据集的实验DSC结果

Figure 5. Segmentation results of this paper’s method and Swin-Unet on the Synapse data set

图5. Synapse数据集上本论文方法和Swin-Unet的分割结果

5. 结语

本次论文提出的方法基于Swin-Unet改进而来,重新设计了跳跃连接的模块,引入通道注意力机制加强了有用通道信息特征的提取,抑制了冗余的通道信息。在编码器的swin-transformer block中添加的res残差模块抑制了模型的梯度消失。在编码器的ASPP模块充分挖掘了多尺度语义和细节信息,完善了特征提取的细节。经过重新设计的网络,利用腹部多器官CT扫描图像进行模型分割训练和测试,其中在自己收集和增强版Synapse数据集上测出的参数平均指标分别为DSC为80.09%、HD为26.22 mm,充分证明本次网络架构的可行性和正确性。

基金项目

浙江省教育厅一般科研项目(Y202044861);

广西高校中青年教师科研基础能力提升项目2022KY0766。

NOTES

*第一作者。

#通讯作者。

参考文献

[1] Ronneberger, O., Fischer, P. and Brox, T. (2015) U-Net: Convolutional Networks for Biomedical Image Segmentation. Medical Image Computing and Computer-Assisted Intervention, MICCAI 2015: 18th International Conference, Munich, 5-9 October 2015, 234-241.
https://doi.org/10.1007/978-3-319-24574-4_28
[2] Zhou, Z., Rahman Siddiquee, M.M., Tajbakhsh, N., et al. (2018) Unet : A Nested U-Net Architecture for Medical Image Segmentation. Deep Learning in Medical Image Analysis and Multimodal Learning for Clinical Decision Support: 4th International Workshop, DLMIA 2018, and 8th International Workshop, ML-CDS 2018, Held in Conjunction with MICCAI 2018, Granada, 20 September 2018, 3-11.
https://doi.org/10.1007/978-3-030-00889-5_1
[3] Diakogiannis, F.I., Waldner, F., Caccetta, P., et al. (2020) ResUNet-a: A Deep Learning Framework for Semantic Segmentation of Remotely Sensed Data. ISPRS Journal of Photogrammetry and Remote Sensing, 162, 94-114.
https://doi.org/10.1016/j.isprsjprs.2020.01.013
[4] Shi, Z., Miao, C., Schoepf, U.J., et al. (2020) A Clinically Applicable Deep-Learning Model for Detecting Intracranial Aneurysm in Computed Tomography Angiography Images. Nature Communications, 11, 6090.
https://doi.org/10.1038/s41467-020-19527-w
[5] Vaswani, A., Shazeer, N., Parmar, N., et al. (2017) Attention Is All You Need. Advances in Neural Information Processing Systems. arXiv preprint arXiv: 1706.03762.
[6] Dosovitskiy, A., Beyer, L., Kolesnikov, A., et al. (2020) An Image Is Worth 16x16 Words: Transformers for Image Recognition at Scale. arXiv: 2010.11929.
[7] Chen, J., Lu, Y., Yu, Q., et al. (2021) Transunet: Transformers Make Strong Encoders for Medical Image Segmentation. arXiv preprint arXiv: 2102.04306.
[8] Cao, H., Wang, Y., Chen, J., et al. (2022) Swin-Unet: Unet-Like Pure Transformer for Medical Image Segmentation. European Conference on Computer Vision, Springer Nature Switzerland, Cham, 205-218.
https://doi.org/10.1007/978-3-031-25066-8_9
[9] Chen, L.C., Papandreou, G., Schroff, F., et al. (2017) Rethinking Atrous Convolution for Semantic Image Segmentation. arXiv:1706.05587.
[10] Oktay, O., Schlemper, J., Folgoc, L.L., et al. (2018) Attention U-Net: Learning Where to Look for the Pancreas. arXiv:1804.03999.
[11] Maji, D., Sigedar, P. and Singh, M. (2022) Attention Res-UNet with Guided Decoder for Semantic Segmentation of Brain Tumors. Biomedical Signal Processing and Control, 71, Article ID: 103077.
https://doi.org/10.1016/j.bspc.2021.103077
[12] Guo, C., Szemenyei, M., Yi, Y., et al. (2021) SA-UNet: Spatial Attention U-Net for Retinal Vessel Segmentation. 2020 25th International Conference on Pattern Recognition (ICPR), Milan, 10-15 January 2021, 1236-1242.
https://doi.org/10.1109/ICPR48806.2021.9413346
[13] Jamali, A., Roy, S.K., Li, J., et al. (2023) TransU-Net : Rethinking Attention Gated TransU-Net for Deforestation Mapping. International Journal of Applied Earth Observation and Geoinformation, 120, Article ID: 103332.
https://doi.org/10.1016/j.jag.2023.103332
[14] Liu, Z., Lin, Y., Cao, Y., et al. (2021) Swin Transformer: Hierarchical Vision Transformer Using Shifted Windows. Proceedings of the IEEE/CVF International Conference on Computer Vision, Montreal, QC, 10-17 October 2021, 10012-10022.
https://doi.org/10.1109/ICCV48922.2021.00986
[15] Lin, A., Chen, B., Xu, J., et al. (2022) DS-TransUnet: Dual Swin Transformer U-Net for Medical Image Segmentation. IEEE Transactions on Instrumentation and Measurement, 71, 1-15.
https://doi.org/10.1109/TIM.2022.3178991
[16] Yang, Y. and Mehrkanoon, S. (2022) AA-TransUNet: Attention Augmented TransUNet for Nowcasting Tasks. 2022 International Joint Conference on Neural Networks (IJCNN), Padua, 18-23 July 2022, 1-8.
https://doi.org/10.1109/IJCNN55064.2022.9892376
[17] Zhang, Y., Liu, H. and Hu, Q. (2021) Transfuse: Fusing Transformers and Cnns for Medical Image Segmentation. Medical Image Computing and Computer Assisted InterventionMICCAI 2021: 24th International Conference, Strasbourg, 27 September-1 October 2021, 14-24.
https://doi.org/10.1007/978-3-030-87193-2_2
[18] 姚庆安, 张鑫, 刘力鸣, 等. 融合注意力机制和多尺度特征的图像语义分割[J]. 吉林大学学报(理学版), 2022, 60(6): 1383-1390.
[19] He, K., Zhang, X., Ren, S., et al. (2016) Deep Residual Learning for Image Recognition. Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, Las Vegas, NV, 27-30 June 2016, 770-778.
https://doi.org/10.1109/CVPR.2016.90