SAM-VMNet 用于冠状动脉造影血管分割的深度神经网络 !

点击下方卡片,关注 「AI视界引擎」 公众号

( 添加时备注:方向+学校/公司+昵称/姓名 )

picture.image

冠状动脉疾病(CAD)是心血管领域最为普遍的一种疾病,也是全球范围内导致死亡的重要因素之一。冠状动脉造影(CTA)图像被视为冠心病诊断的权威标准,通过对CTA图像进行血管分割和狭窄检测,医生能够更精确地诊断冠心病。

为了结合基础模型和领域特定模型的优势,并在有限的训练样本下实现高精度、全自动的分割与检测,我们提出了一种全新的架构SAM-VMNet。

该模型结合了MedSAM强大的特征提取能力和VM-UNet的视觉状态空间模型具有线性复杂度的优势,使其具有比Vision Transformer更快的推理速度,更强的数据处理能力,实现了对CTA图像更高的分割准确率和稳定性。

实验结果表明:SAM-VMUnet架构在CTA图像分割任务中表现卓越,分割准确率高达98.32%,敏感性高达99.33%,显著优于现有的其他模型,具有更强的领域适应性。

对CTA图像分割任务的综合评估表明,SAM-VMNet精确地提取了血管主干和毛细血管,展示了其在血管分割任务上的巨大潜力和广泛的应用场景,也为进一步狭窄检测奠定了坚实的基础。

  1. 介绍

冠状动脉疾病是最常见的心血管疾病,并且是世界上最常见的死亡原因之一。在世界卫生组织(WHO)发布的最新数据,心血管疾病是四大非传染性慢性疾病之一,也是死亡人数增加最多的疾病[1]。该疾病的特征为血管中斑块的积累,称为动脉粥样硬化(Atherosclerosis),它导致血管狭窄和硬化,从而造成组织或脏器的缺血性改变,增加了心绞痛、心肌梗死以及其他心血管事件的风险 [2]。

冠状动脉造影(CA)在临床上被认为是诊断冠状动脉疾病(CAD)的“金标准”[3],可以显示出左冠状动脉或右冠状动脉的骨干及其分支血管,可以了解血管有无狭窄病灶存在。由于冠状动脉血管造影图像较为复杂,很难像脑部、四肢血管造影图像一样得到清晰的减影图像,因此从冠脉造影图像提取出清晰、完整的血管树结构是至关重要的。

目前,动脉造影图像分割有许多经典的方法如:阈值分割法[4],Canny算子边缘检测法[5],区域生长法[6],以及基于跟踪的方法[7-9]。但由于动脉造影图像通常对比度低且背景复杂,传统算法在实际应用中可能难以满足临床需求。因此,近年来提出了许多改进方法和新技术,例如深度学习方法,以提高动脉造影图像分割的准确性和鲁棒性。卷积神经网络(CNN)利用深度学习模型自动提取特征,通过跳跃连接保留高分辨率特征 [10]。其中基于卷积神经网络的UNet[11]已经广泛用于生物医学图像分割任务中。Segment Anything Model(SAM)[12]作为基于Vision Transformer[13]进行特征提取的图像分割大模型具有强大的通用性以及泛化能力,其中MedSAM[14]在SAM基础上采用一百万张医学图像进行训练,更好地适应了医学图像的特点。在医学图像分割中展示了巨大潜力。最近,Mamba模型[15]作为新兴的深度学习序列架构被提出,它通过选择性状态空间模型来改进传统的状态空间模型,提高了处理长序列数据的效率。基于Mamba模型提出的VM-Unet[16]处理大规模医学图像数据时更为高效。然而,最新的深度学习架构在处理小血管、严重狭窄的血管或图像质量较差的情况下仍然表现不佳[17-18]。本研究提出了一种新的深度学习框架:SAM-VMNet,通过结合VM-UNet和MedSAM的优点,实现了冠状动脉造影图像的高质量自动分割。主要过程是利用SAM中的Transformer模块,提取冠状动脉造影图像中的全局和局部特征;再采用Mamba模型中的选择性状态空间模型,增强对长序列数据的处理能力,提高分割精度。通过对比实验以及消融实验,SAM-VMNet在分割精度和计算效率方面均优于基础模型和其他流行的深度学习模型。

2 Method

2.1 Architecture Overview

picture.image 模型图(1)

(图1:SAM-VMUet框架示意图)

我们所提出的SAM-VMUnet框架的体系结构如图1所示,该模型结合了MedSAM和VM-UNet的技术优势。具体而言,SAM-VMUnet具有两个并行的编码器,其中MedSAM是医学图像分割领域被广泛使用的大模型,VM-UNet是基于视觉状态空间块的全新架构。这使得它不仅能利用MedSAM强大的特征提取能力,还发挥VM-UNet在建立远距离依赖关系的同时保持线性复杂度的优势,加快推理速度。由于MedSAM需要强提示,我们首先训练了一个简单的VM-UNet,对图像进行粗分割,随后在掩码图像上等间距选取10个点作为提示反馈给MedSAM,从而自动获取MedSAM编码器输出的特征向量。来自两个编码器的特征向量被串联,然后送入VM-UNet的解码器,最终得到输出的掩码图像。MedSAM作为即插即用的插件,其参数在训练过程中被冻结,只有VM-UNet的权重在训练过程中不断更新。

2.2 VM-UNet

Vision Mamba Uet(VM-UNet)是第一个纯基于State Space models(SSM)的医学图像分割模型。该模型采用非对称的编码器-解码器结构,并引入了视觉状态空间(VSS)块来捕捉广泛的上下文信息,从而在医学图像分割任务中表现出色。

picture.image

(图2:(a)VM-UNet总架构; (b)VSS块是VM-UNet的主要构建块,而SS2D是VSS块中的核心操作)

图(a)展示了VM-UNet的总架构,具体来说,VM-UNet包括Patch Embedding layer, an encoder, a decoder, a Final Projection layer, and skip connections。Patch Embedding层将输入图像划分为大小为的不重叠的Patch,随后将图像映射为维。这个过程将产生图像嵌入。最后,在将输入编码器进行特征提取之前,我们使用Layer Normalization对其进行归一化。编码器由四个阶段组成,在前三个阶段的末尾应用Patch合并操作,以降低输入特征的高度和宽度,同时增加通道数量。四个阶段分别使用个VSS区块,每个阶段的通道计数为。解码器同样被分成四个阶段。在最后三个阶段的开始,采用patch expanding操作来减少特征通道的数量,增加特征通道的高度和宽度。在四个阶段中,使用个VSS块(非对称架构),每个阶段的通道计数为。在解码器之后,使用Final Projection层来恢复特征的大小以匹配分割目标。具体来说,通过patch expansion进行4次上采样来恢复特征的高度和宽度,然后通过投影层来恢复通道数量。对于跳连接,直接采用加法操作,因此不会引入任何额外参数。

VSS块来源于VMamaba,是VM-UNet的核心模块,如图(b)所示。输入经过Layer Normalization后,分成两个分支。 在第一个分支中,输入经过一个线性层和一个激活函数。在第二个分支中,输入通过线性层、深度可分离卷积和激活函数进行处理,然后输入到2D-Selective-Scan (SS2D)模块中进行进一步的特征提取。随后,使用Layer Normalization对特征进行规范化,然后合并两条路径。最后,使用线性层混合特征,并将此结果与残差连接相结合,形成VSS块的输出。

VM-UNet 最基本的 Binary Cross-Entropy and Dice loss (BceDice loss)和Cross-Entropy and Dice loss (CeDice loss)作为损失函数用于二类和多类分割任

务:

其中:为样本总数,为类别总数;和分别表示真实标签和预测;是一个指标,如果样本属于类,则等于1,否则等于0;

为模型预测样本i属于类别c的概率;和分别表示真实值和预测值;为损失函数的权值,默认值均为1。

2.3 Parallel network architecture

MedSAM是首个专为通用医学图像分割设计的基础模型,该模型基于Vision Transformer架构,包含三个主要组成部分:图像编码器、提示编码器和掩码解码器。并通过在超过100万张公开医疗图像数据集上进行微调,达到了多个医疗图像分割任务的最佳性能,该架构已被证明具有极强的特征提取能力。

为了结合MedSAM强大的特征提取能力,我们首先训练了一个简单的VM-UNet对图像进行粗分割,以获得适用于MedSAM的边界提示符。因此,SAM-VMUnet采用了两条并行的网络结构来提取图像特征。具体而言,图像输入经过分支1的VM-UNet得到粗分割图像,在此图像上等间距选择10个点作为MedSAM的prompt,从而获得特征向量。

图像输入经过分支2的VM-UNet编码器获得特征向量。对用1×1卷积和平均池化层进行下采样,使的维度与维度保持一致,进一步实现特征融合。使用注意力机制进行加权的特征融合。这一过程合理利用了MedSAM的提示编码器,从而生成了高质量的特征向量。

融合向量接着通过VM-UNet的解码器进行上采样,最终输出预测结果。在反向传播的训练过程中,仅分支2中的VM-UNet权重不断更新。

MedSAM和和分支1的简单VM-UNet的权重则被冻结,这使它们可以作为即插即用的特征提取插件进行使用。

3、实验设置

本研究中使用了多个数据集,其中包括一个私有数据集和公开的ARCADE数据集。私有数据集经中国贵航集团三0二医院批准,一共包含25张带注释的冠脉造影图像分割数据集,其中15张用于训练,10张用于测试。ARCADE数据集[19]包含包含1200张冠状动脉血管树图像,按照训练集(1000张)和验证集(200张)进行划分。每张训练图像都有根据Syntax Score方法[20]标注的26个不同区域的注释。

关于数据集使用,首先我们将私有数据集用于粗分割的 VM-UNet,基于vmamba_small_e238_ema预训练的权重上进行了微调。我们将ARCADE数据集用于主干网络SAM-VMNet的学习;在MedSAM模型中,我们使用了medsam_vit_b作为初始权重来进行feature的获取。

在网络的训练中,我们基于Pytorch实现,并在一个NVIDA GeForce 4090显卡上进行训练。设置了200个epoch,seed为42,batch为32,dorp速率为0.2,初始学习率为0.001,采用AdamW优化策略进行优化,损失函数同VM-UNet,使用二元交叉熵和Dice损失。

在评价指标中,我们选择了mLoU、dsc、accuracy、specificity以及sensitivity。

mLoU 是一种常用于评估图像分割任务的性能指标。首先,计算预测区域与真实区域的交集与它们的并集之比,计算每个类别的交并比之后,取这些值的半均作为 mLoU。mLoU 值越高,表示分割结果与真实情况的吻合度越高。计算公式如下:

dsc 可以用于评估两个样本的相似度。计算方式为两倍的预测结果和真实结 果交集的大小除以两者的总大小。计算公式如下:

accuracy 可以显示模型正确识别出所有类别的能力。计算方式是所有正确分 类的像素数除以总像素数。计算公式如下:

specificity 衡量的是所有负样本中,模型正确识别为负的比例,值越高,表示模型很少将负样本误判为正样本;sensitivity 衡量的是所有正样本中,模型正确识别为正的比例,值越高,表示模型很好地捕捉到了所有正样本。计算公式如

下:

4、结果

为了评估SAM-VMNet对于冠脉造影图像的分割性能,我们选取了医学图像分割流行的深度学习网络:UNet[21]、UNet++[22]、Transunet[23]、MALUNet[24]、Transfuse[25]、missformer[26]与我们提出的网络以及原始的VM-UNet进行对比,采用相同的数据集以及预处理策略,计算分割指标如表1所示。

表1:模型对比的指标

ModelmIoUAccSpeSenF1
UNet51.53%96.56%52.26%53.25%55.23%
UNet++58.56%98.08%76.47%71.43%73.87%
Transunet43.64%97.62%75.69%51.41%58.89%
MALUNet55.61%98.01%99.15%68.01%71.47%
Transfuse57.37%97.69%98.19%85.11%72.10%
missformer23.81%96.93%68.63%27.04%36.80%
VM-UNet54.45%97.73%98.86%69.87%70.51%
Proposed network63.03%98.32%99.33%73.43%77.33%

表中显示了在1500数量的冠状动脉图像数据上的模型分割性能。在所有的模型中,Proposed network在mIoU、Acc, Spe, F1评估指标优于其他模型,这显示了我们的Proposed network具有显著竞争力的性能,我们的模型性能超过了基于Transformer的Transunet,以及具有多种注意力模块的MALUNet。实验结果也表明了SSM模型在医学图像分割领域具有优越性。

研究结果表明,Proposed network在冠状动脉图像分割上有极好的表现,通过将预训练的MedSAM编码器集成到VM-Unet框架中,有效提高了模型的泛化能力和性能。proposed network既可以利用MedSAM强大的特征提取能力,也发挥了VM-Unet捕获上下文信息的能力。在冠状动脉图像分割的实验中显示出了优于其他模型的优势。超过Transunet, MALUnet和 Transfuse等先进模型。与SAM相比,Proposed network可以使用VM-Unet针对医学图像语义分割进行优化,提高了图像分割的准确性。我们的研究表明,通过多头设计,Proposed network在端到端训练期间同时学习形状先验,而不再需要训练单独的网络来编码图像信息,结合多种分割框架的模型可以有效提高医学图像分割的准确性,不同模块提取的特征融合,进一步增加了特征的丰富度,使得模型可以更好地进行分割任务。

5、讨论

研究结果证明了SAM-VMNet在医学图像分割方面的优越性能。通过将预训练的SAM编码器集成到VM-UNet框架中,SAM-VMNet利用了MedSAM在医学图像上强大的特征提取能力,同时受益于VM-UNet的长距离建模能力。在公开的冠脉血管数据集上,SAM-VMNet在各种任务中始终保持最高的准确性,超越了高级模型如nnUNet、UNet++和TransUNet,并且保持着最高的特异性,这表明我们的模型在识别非目标区域方面非常有效。

与nnSAM[27]相比,我们通过一个微调的VM-UNet进行粗略分割,为MedSAM提供提示点,从而获得特征向量,而不是嵌入。这种方法更好地利用了MedSAM的强大特征提取能力,更有利于提取感兴趣的血管区域。我们的研究证明,结合了Transformer和Mamba分割框架的SAM-VMNet可以有效地进一步提高医学图像分割的精度。此外,由于VM-UNet保持线性复杂度,能够降低运算的消耗,从而提高计算效率。

6、局限

尽管SAM-VMNet在医学图像分割中表现出色,但仍存在一些局限性需要进一步探讨和改进。

首先,SSM的记忆本质上是有损的,因此它不如注意力机制的无损记忆。Mamba在处理短序列方面无法展现它的优势,所以可能会损失图像的一些局部特征或小范围的结构,但是这一领域恰好是注意力机制表现出色的地方。

其次,在提示点的选择上,虽然使用微调的VM-UNet进行粗略分割并生成提示点的方法在我们的研究中表现良好,但这种方法的效果可能会受到初始分割结果的质量影响。如果初始分割结果不够准确,可能会影响后续MedSAM的特征提取效果,从而影响最终分割结果。

此外,尽管我们的方法在冠脉血管数据集上表现出色,但其在其他类型的医学图像分割任务中的泛化能力仍需进一步验证。不同类型的医学图像具有不同的特征和挑战,因此,需要在更广泛的医学图像数据集上进行评估,以全面验证SAM-VMNet的通用性和稳定性。

7、总结

我们提出了一个全新的网络架构:SAM-VMNet,该架构充分结合了MedSAM的强大特征提取能力与VM-UNet在建立远距离依赖关系同时保持线性复杂度的能力。通过这种组合,SAM-VMNet实现了对冠脉造影图像的高质量分割,适用于冠脉血管多分支的情况。

我们的实验结果表明,SAM-VMNet在医学图像分割任务中表现出色,在准确性和特异性上超过了当前的先进模型。未来,我们计划进一步优化该模型,并在更广泛的医学图像分割任务中进行验证,以期提升其通用性和实用性。

点击上方卡片,关注 「AI视界引擎」 公众号

0
0
0
0
评论
未登录
暂无评论