1. 引言
医学图像分割是医学图像处理中的一项重要任务,它直接影响到疾病的早期诊断、治疗计划的制定以及患者的治疗效果评估。随着计算机视觉和深度学习技术的快速发展,基于深度神经网络的医学图像分割方法在多个医学领域,尤其是肿瘤检测和器官分割中取得了显著的成果。医学影像技术在癌症诊断中发挥着不可或缺的作用,主要包括X射线检查、电子计算机断层成像(computed tomography, CT)、超声(ultrasound)、乳腺钼靶X线摄影检查、磁共振成像(magnetic resonance imaging, MRI)以及正电子发射断层成像(positron emission tomography, PET)等[1]。U-Net网络是一种专门为医学图像分割任务设计的深度卷积神经网络,它首次由Olaf Ronneberger等人[2]在2015年提出,并迅速成为医学图像分割领域的标准模型。U-Net网络在医学图像分析中的广泛应用,尤其是在细胞分割、肿瘤检测和器官分割等任务中,表现出了良好的分割性能。
曹心姿等人[3]提出了一种基于改进的U-Net网络的Res-U-Net模型,通过引入残差结构和数据标准化,有效提高了膀胱肿瘤MRI图像分割的准确性和鲁棒性。实验结果表明,改进后的模型在分割效果上优于原始U-Net模型,具有较高的研究价值。Li X等人[4]提出了一种用于膀胱癌和膀胱壁MRI图像分割的新型深度学习模型MH2AFormer (Multi-scale Hybrid Attention with Transformer)。该模型通过引入多尺度混合注意力机制和Transformer模块,有效提高了膀胱图像分割的准确性和效率,尤其在处理复杂肿瘤形态和模糊边界时表现出色。Wang Z等人[5]提出了Mmaba-Unet,结合了U-Net网络结构和Mamba模块的新型医学图像分割网络,通过高效的利用全局上下文建模和多尺度特征融合,显著提升了分割精度。S Deari等人[6]提出的BFMD SN U-Net通过解决过拟合问题、增强特征学习和引入全局上下文信息,显著提高了视网膜血管分割的性能。胡鸿翔等人[7]提出了一种全分辨率卷积交互网络模型FRCI-Net。该模型通过多分辨率卷积交互机制,在保持完整图像分辨率的同时,有效扩展特征图的水平和垂直方向,集成相邻阶段的多尺度特征图以补充上下文信息,强化特征表示,并通过改进的残差模块学习多分辨率表示,生成高精度预测结果。然而,尽管U-Net在医学图像分割中展现了优异的性能,但是在面对复杂的医学图像时,特别是处理结构不规则和纹理模糊的区域时,仍然存在一定的局限性。因此,如何进一步提高图像的分割精度和鲁棒性,成为了当前研究的重点之一。
为了克服现有方法带来的局限性,本文提出了一种基于改进Mamba的医学图像分割模型Akmamba-Net。我们在U-Net网络模型的基础上,结合了Mamba-out模块,并引入了AKConv卷积模块,以增强模型对复杂形状和边界的捕捉能力。通过对Mamba模块的改进,优化了特征融合的策略,使得模型在处理复杂医学图像时能够更加精准地识别和分割视网膜眼底和膀胱肿瘤区域。此外,AKConv模块进一步提升了模型对细节的处理能力,从而增强了分割精度。
本文的主要贡献可以概括为以下几点:
1) 提出了一种改进的Mamba模块Mamba-out,增强了医学图像分割模型的特征融合能力,特别是在边缘区域的分割精度方面。
2) 设计了一个集成AKConv的U-Net网络,通过引入空间重采样和卷积模块,提高了模型对复杂边界和形态不规则的肿瘤区域的处理能力。
3) 在视网膜眼底和膀胱肿瘤MRI图像分割任务中,实验结果表明,改进后的模型在分割精度和鲁棒性上显著优于传统U-Net模型。
2. 网络模型
U-Net的内部结构由编码器、跳跃连接层和解码器三个主要部分组成。编码器部分通过一系列卷积层和池化层逐步提取图像特征并降低空间分辨率;跳跃连接层连接着编码器和解码器,进一步处理图像的高级特征;解码器部分通过跳跃连接将编码器的低层特征与解码器的高层特征结合,从而精确恢复图像的细节,将特征图逐渐恢复到原始的输入图像尺寸,并逐步生成分割结果。最终,输出层通过卷积操作生成分割结果,输出与输入图像相同大小的图片。
Figure 1. Akmamba-Net network architecture
图1. Akmamba-Net网络结构
本文提出了一种改进的U-Net的分割网络Akmamba-Net,在原本的U-Net模型架构中将普通卷积替换为了AKConv,通过引入自适应的卷积核尺寸来更好地捕捉输入特征的空间信息。在标准卷积操作中,卷积核的大小是固定的,而AKConv通过引入可学习的参数,使得卷积核的大小和形状能够动态调整,从而更灵活地处理不同尺度和结构的图像特征。在跳跃连接层加入了Mamba-out模块,能够有效处理图像中的复杂特征。Akmamba-Net网络结构见图1。
2.1. AKConv
AKConv模块[8]是一种改进的卷积层,提供了一种灵活的卷积机制,通过引入动态偏移和双线性插值机制,增强了卷积操作的适应性,特别是在处理图像或特征图的空间变换时(例如目标的平移或旋转)。这种方法突破了传统卷积局限于固定的窗口和采样形状,使得卷积操作可以更精准地适用于不同位置的目标和数据集。
首先,输入特征图(x)经过p_conv卷积层计算出一个偏移量(offset)。该偏移量反映了特征图中每个位置的位移量,决定了卷积操作在采样时应当采用的新坐标系。接下来,利用offset和预定义的常数生成新的空间位置(p)。这个新位置用于确定卷积操作在特征图中的采样位置,确保卷积操作可以适应输入图像中的空间变化。
在生成新的采样位置后,AKConv通过双线性插值计算每个采样点的权重。具体来说,它通过计算四个邻近像素的距离,得到每个像素的新位置的加权值,这些权重会用于将四个邻近位置的特征值加权融合,从而得到更精细的特征值。这些加权后的特征值构成了一个新的特征图(x_offset),该特征图即为输入图像经过偏移和插值后的重采样结果。
最后,重采样后的特征图经过标准的卷积层(conv)、归一化层和SiLU激活函数进行进一步处理,以提取更高层次的特征,并生成最终的输出特征图(Output)。在整个过程中,AKConv模块通过精细的空间调整和插值机制,使得网络能够更好地处理图像中的平移、旋转等空间变换,从而实现高效的特征提取。AKConv模块网络结构见图2。
Figure 2. AKConv network architecture
图2. AKConv模块网络结构
总结来说,AKConv通过三个关键步骤:计算偏移量、进行双线性插值、重采样并卷积,来实现更精细的特征提取。这种动态调整卷积采样位置的方式,使得卷积神经网络能够在面对复杂的空间变换时,仍然保持较强的适应性和表现力,提高了特征提取的准确性。
2.2. Mamba-Out
本研究提出了mamba-out模块,专门设计用于处理具有复杂时空结构的二维数据。该模块结合了卷积神经网络(CNN)和特征投影技术,能够有效提取并融合输入数据中的空间和时间信息。模块的设计目标是在处理时序数据时保持计算效率,同时保留足够的特征表达能力。以下是mamba-out模块的详细架构与工作流程。
输入Input是一个形状为(B, C, H, W)的特征图,其中B是批次大小,C是通道数,H和W分别表示特征图的高度和宽度。在输入特征图进入网络之前,首先通过线性层进行初步的处理,得到新的特征图xz。
接着,张量xz被拆分为两个部分:x和z,其中x将用于后续的卷积操作,而z用于控制模型输出的缩放。此步骤的关键在于通过对z进行控制,实现模型输出的动态调整输出的大小。
处理后的x会传入一个深度可分离卷积层conv2d,该层负责对数据进行空间特征的提取。卷积操作通过对每个输入通道应用独立的卷积核(组卷积),有效减少了计算量,并保留了重要的空间特征。卷积层的输出经过SiLU激活函数进行非线性转换,从而增强模型的表达能力。此阶段的操作通过以下公式(1)完。
(1)
网络的核心计算通过forward_core函数完成,主要通过多次线性变换、卷积操作和状态更新来处理输入图片。在这一阶段,输入的张量x会被分解为多个维度,并通过时间序列处理模块进行处理。该模块通过迭代扫描,结合状态参数Alogs和时间参数Ds,逐步更新状态信息,从而更好地学习时序特征。
通过这种方式,网络能够有效地捕捉时间依赖性和空间依赖性的信息,并在多个路径中融合这些信息。这些处理最终会生成四个张量y1、y2、y3、y4,分别代表不同状态下的输出。四个输出张量相加之后,得到最终的输出y,如公式(2)所示。
(2)
之后y会经过LayerNorm层进行标准化,使其具有零均值和单位方差。之后,输出乘以由SiLU控制的缩放因子z,动态调整输出的尺度。最终,处理后的输出通过一个线性层映射回原始的特征维度。
Mamba-out模块为处理复杂的医学特征图片提供了一种高效且灵活的解决方案。通过深度可分离卷积、状态更新和动态缩放等技术,模型能够在计算效率和表达能力之间取得平衡。实验结果表明,该模块在医学图像分割任务中表现优异,具有广泛的应用前景。Mamba-out模块网络结构见图3。
Figure 3. Mamba-out network architecture
图3. Mamba-out模块网络结构
3. 实验结果和分析
3.1. 数据集
本文采用2个数据集,分别是视网膜眼底图像公开数据集DRIVE (Digital Retinal Images for Vessel Extraction)、ISICDM2019临床数据分析挑战赛的基于磁共振成像的膀胱内外壁分割与肿瘤检测数据集,膀胱数据采集自空军军医大学(第四军医大学),已经分别对肿瘤区域和膀胱壁区域进行了标记,标签使用one-hot编码进行处理。
3.2. 实验环境和参数配置
实验是在Ubuntu22.04操作系统上进行的,使用深度学习框架Pytorch。GPU配置为单卡的NVIDIAGeForce GTX 4090,CUDA的版本为11.8。学习率为0.0005,并设置了衰减权重,衰减率为0.0001,每次训练批次(Batch)大小设置为8,优化器选取了Adam,并进行了1000个Epoch的训练周期。
3.3. 评价指标
本文采用精确率(Precision)、Dice系数和IoU对模型的分割性能进行评价。
TP:模型正确预测为正类的样本数。FP:模型错误地预测为正类的样本数。精确率高意味着模型预测为正的样本中,正确的比例较高。
(3)
Dice系数是用于衡量两个样本集合相似度的指标,常用于语义分割任务中,计算的是预测结果和真实标签之间的重叠度。
(4)
IoU也被称为Jaccard指数,主要用于度量两个区域的重叠度,常用于目标检测和图像分割任务中。IoU的定义为:
(5)
3.4. 对比试验
针对视网膜眼底图像公开数据集和膀胱肿瘤图像,为了验证提出的模型的性能,选取较为主流的分割模型Unet++ [9]、Attention-unet [10]、Unet3+ [11]、Transunet [12],对比实验结果见表1和表2。
Table 1. Comparison of experimental results
表1. 对比实验结果
Model |
Precision |
Dice |
IoU |
Unet |
95.61 |
82.22 |
71.23 |
Attention-unet |
95.32 |
80.94 |
71.15 |
Unet++ |
95.64 |
81.92 |
70.95 |
Unet3+ |
97.03 |
80.32 |
71.04 |
Transunet |
97.16 |
82.32 |
70.97 |
Ours |
97.21 |
82.57 |
71.93 |
根据表1和表2的结果表明,Akmamba-Net算法相较于Unet、Attention-unet、Unet++、Unet3+、Transunet等模型,在视网膜眼底MRI图像中Precision提升了1.6%、1.89%、1.57%、0.18%、0.11%,Dice系数提升了0.35%、1.63%、0.65%、2.25%、0.25%,IoU提升了0.7%、0.78%、0.98%、0.89%、0.96%,在膀胱肿瘤MRI图像中Precision提升了8.52%、9.04%、2.5%、3.99%、3.32%,Dice系数提升了4.31%、0.75%、9.86%、2.33%、3.55%,IoU提升了6.72%、1.13%、13.83%、3.53%、6%,更适用于医学图像分割的任务,证实了改进算法的优越性和综合性能。AKConv模块通过动态调整卷积核形状,有效捕捉了视网膜血管的细小分支和复杂拓扑结构,而Mamba-out模块通过跨层特征融合增强了血管边缘的连续性。膀胱肿瘤形态不规则且边界模糊,传统U-Net的固定卷积核难以适应此类特征。AKConv的空间重采样机制以及Mamba-out动态缩放策略则抑制了背景噪声的干扰,在Dice系数上提升了4.3%。视网膜眼底和膀胱肿瘤MRI图像相较于原始模型的分割效果见图4。
Table 2. Comparison of experimental results
表2. 对比实验结果
Model |
Precision |
Dice |
IoU |
Unet |
81.12 |
85.67 |
74.93 |
Attention-unet |
80.60 |
89.23 |
80.52 |
Unet++ |
87.14 |
80.12 |
67.83 |
Unet3+ |
85.65 |
87.65 |
78.12 |
Transunet |
86.32 |
86.43 |
75.65 |
Ours |
89.64 |
89.98 |
81.65 |
Figure 4. Akmamba-Net model segmentation performance comparison
图4. Akmamba-Net模型分割效果对比
3.5. 消融试验
为了验证本文模型性能的优越性,针对膀胱肿瘤图像,为了验证提出的模型的性能,将本文提出的Akconv模块、mamba-out模块进行了一系列的消融实验,结果见表3。
结果表明,针对膀胱肿瘤图像,Unet通过引入Akconv和mamba-out,引入Akconv后,Precision、Dice、IoU提升了6.22%、1.65%、2.31%,证明了Akconv能够提高特征提取的精度和鲁棒性,引入mamba-out后,Precision、Dice、IoU提高了7.6%、2.35%、4.08%,证明了mamba-out能够有效提取并融合输入数据中的空间和时间信息。
Table 3. Results of the ablation experiment
表3. 消融实验结果
Model |
Precision |
Dice |
IoU |
Unet |
81.12 |
85.67 |
74.93 |
+Akconv |
87.34 |
87.32 |
77.24 |
+mamba-out |
88.72 |
88.02 |
79.01 |
+Akconv +mamba-out |
89.03 |
88.97 |
80.21 |
Ours |
89.64 |
89.98 |
81.65 |
4. 结论
本文针对视网膜血管分割不佳和膀胱肿瘤图像分割边界模糊、误判等现象,提出了Akmamba-Net网络模型,创新性地提出了Akconv模块对输入的特征图进行空间调整,从而增强模型对空间信息的捕捉能力,为传统卷积操作引入了空间调整的能力,使得模型能够在处理空间复杂性时表现得更加灵活。同时引入mamba-out模块,使用了注意力机制来动态地加权输入特征,提高对关键区域的关注度,通过集成多尺度卷积和跨层信息传递,从而提升模型的分割精度。经过实验验证在视网膜眼底和膀胱肿瘤MRI图片上展现出了优越的分割性能,通过对比试验和消融实验证实了其在医学图像分割场景的应用潜力,满足医学图像分割要求。
基金项目
国家自然科学基金资助项目(11502145)。