跨境派

跨境派

跨境派,专注跨境行业新闻资讯、跨境电商知识分享!

当前位置:首页 > 综合服务 > 电商平台 > VM-UNet

VM-UNet

时间:2024-04-30 14:10:47 来源:网络cs 作者:康由 栏目:电商平台 阅读:

标签:

首次探索纯粹基于SSM的模型在医学图像分割中的潜在应用。

在医学图像分割领域,基于CNN和Transformer的模型都已经被广泛研究。然而,CNN在长距离建模能力上存在局限,而Transformer则受到其二次计算复杂度的制约。最近,以Mamba为代表的态空间模型(SSMs)成为了一种有前景的方法。它们不仅在建模长距离交互方面表现出色,同时保持了线性计算复杂度。

在本文中,利用态空间模型,作者提出了一个U形架构模型,用于医学图像分割,命名为Vision Mamba UNet(VM-UNet)。具体来说,引入了视觉态空间(VSS)块作为基础块以捕捉广泛的上下文信息,并构建了一个非对称的编码器-解码器结构。作者在ISIC17、ISIC18和Synapse数据集上进行了全面的实验,结果表明VM-UNet在医学图像分割任务中表现出竞争力。据作者所知,这是第一个基于纯SSM模型的医学图像分割模型。作者旨在建立 Baseline ,并为未来开发更高效、有效的SSM-based分割系统提供宝贵的洞见。

代码:https://github.com/JCruan519/VM-UNet

1 Introduction

自动化医学图像分割技术帮助医生更快地进行病理诊断,从而提高了患者护理的效率。近来,基于CNN和基于Transformer的模型在各种视觉任务中展示了卓越的性能,特别是在医学图像分割方面。UNet作为基于CNN模型的代表,以其结构简单和强大的可扩展性而闻名,后续许多改进都是基于这种U形架构。

TransUnet,作为基于Transformer模型的先驱,首次在编码阶段采用视觉Transformer(ViT)[13]进行特征提取,并在解码阶段使用CNN,展示了其在获取全局信息方面的显著能力。随后,TransFuse融合了ViT与CNN的并行架构,同时捕捉局部和全局特征。此外,Swin-UNet将Swin Transformer与U形架构结合,首次引入了一个纯粹的基于Transformer的U形模型。

然而,基于卷积神经网络(CNN)的模型和基于Transformer的模型都具有固有的局限性。基于CNN的模型受到局部感受野的限制,这大大阻碍了它们捕捉远程信息的能力。这通常会导致提取的特征不充分,从而影响分割效果不佳。尽管基于Transformer的模型在全局建模方面表现出色,但自注意力机制在图像尺寸方面需要二次复杂度,导致了较高的计算负担,特别是对于像医学图像分割这类需要密集预测的任务。这些模型当前的不足促使作者开发一种新的医学图像分割架构,这种架构能够捕捉强烈的远程信息,并保持线性的计算复杂度。

近期,状态空间模型(SSMs)在研究行人中引起了极大的关注。在经典SSM研究的基础上,现代SSM(例如,Mamba)不仅建立了长距离依赖性,而且对于输入大小展现出线性复杂性。此外,基于SSM的模型在许多领域都得到了大量研究,包括语言理解,通用视觉等。特别是,U-Mamba最近引入了一种新颖的SSM-CNN混合模型,首次将其应用于医学图像分割任务中。SegMamba在编码器部分整合了SSM,而在解码器部分仍然使用CNN,为3D脑肿瘤分割任务提出了一个SSM-CNN混合模型。尽管上述工作已经将SSM用于医学图像分割任务,但基于纯SSM模型的性能尚未被探索。

受到V Mamba [20]在图像分类任务中成功的启发,本文首次引入了视觉猛蛇UNet(VM-UNet),这是一种纯粹的基于SSM的模型,旨在展示其在医学图像分割任务中的潜力。

具体来说,VM-UNet由三个主要部分组成:编码器、解码器和跳跃连接。编码器使用来自V Mamba的VSS块进行特征提取,并伴有图像下采样的像素合并操作。相反,解码器则由VSS块和像素扩展操作组成,以恢复分割结果的大小。对于跳跃连接部分,为了强调最原始的基于纯SSM模型在分割性能上的表现,作者采用了最简单的加性操作形式。

在器官分割和皮肤病变分割任务上进行了综合实验,以展示基于纯SSM模型在医学图像分割中的潜力。具体来说,作者在Synapse,ISIC17,和ISIC18数据集上进行了大量实验,其结果指出VM-UNet可以达到具有竞争力的性能。此外,值得注意的是,VM-UNet代表了基于纯SSM分割模型的最基本形式,因为它不包括任何特别设计的模块。

本文的主要贡献可以概括如下:

作者提出了VM-UNet,标志着首次探索纯粹基于SSM的模型在医学图像分割中的潜在应用。

在三个数据集上进行了全面实验,结果表明VM-UNet具有相当大的竞争力。

作者为医学图像分割任务中纯粹基于SSM的模型建立了一个 Baseline ,提供了有价值的见解,为更高效、更有效的SSM-based分割方法的发展铺平了道路。

2 Preliminaries

3 Methods

在本节中,作者初步介绍了VM-UNet的整体结构。随后,作者将详细阐述其核心组件,即VSS块。最后,作者将描述在训练过程中使用的损失函数。

Vision Mamba UNet (VM-UNet)

如图1(a)所示,VM-UNet的整体架构被呈现出来。具体来说,VM-UNet包括一个Patch嵌入层、一个编码器、一个解码器、一个最终投影层以及跳跃连接。与之前的方法[9]不同,作者没有采用对称结构,而是使用了一种非对称设计。

同样地,解码器被组织成四个阶段。在最后三个阶段的开始,使用了一个块扩展操作来减少特征通道的数量并增加高度和宽度。在四个阶段中,作者使用了[2, 2, 2, 1] VSS块,每个阶段的通道数分别为[8C, 4C, 2C, 4]。在解码器之后,采用了一个最终的投影层来将特征的大小恢复到与分割目标相匹配。具体来说,通过块扩展执行了4倍的上采样来恢复特征的高度和宽度,然后通过一个投影层来恢复通道数量。

对于跳跃连接,采用了简单直接的加法操作,没有花哨的装饰,因此不会引入任何额外的参数。

VSS block

VSS块源自VMamaba的结构,是VM-UNet的核心模块,如图1(b)所示。在经过层归一化之后,输入被分为两个分支。在第一个分支中,输入通过一个线性层,然后是一个激活函数。在第二个分支中,输入经过一个线性层、深度可分离卷积和一个激活函数的处理,然后被送入2D-选择性扫描(SS2D)模块进行进一步的特征提取。

随后,特征通过层归一化进行标准化处理,并与第一个分支的输出进行逐元素乘法,以合并两条路径。最后,使用线性层混合特征,并将此结果与残差连接相结合,形成VSS块的输出。在本文中,默认采用SiLU作为激活函数。

SS2D由三个部分组成:扫描扩展操作、S6模块和扫描合并操作。如图2(a)所示,扫描扩展操作沿四个不同方向(从左上到右下,从右下到左上,从右上到左下,以及从左下到右上)展开输入图像成序列。这些序列随后由S6模块进行处理以提取特征,确保来自各个方向的信息得到彻底扫描,从而捕捉到多样的特征。接着,如图2(b)所示,扫描合并操作对来自不同方向的序列进行求和合并。

四个方向,将输出图像恢复到与输入图像相同的大小。S6模块源自Mamba,在S4之上引入了一种选择机制,通过根据输入调整SSM的参数。这使得模型能够区分并保留相关信息,同时过滤掉不相关的内容。S6模块的伪代码在算法1中给出。

Loss function

VM-UNet的引入旨在验证基于纯SSM模型的在医学图像分割任务中的应用潜力。因此,作者分别针对二值和多类分割任务,仅使用最基本的二元交叉熵和Dice损失(BceDice损失)以及交叉熵和Dice损失(CeDice损失)作为损失函数,如方程式5和6所示。

4 Experiments

在本节中,作者对VM-UNet在皮肤病变和器官分割任务上进行了全面的实验。具体来说,作者在ISIC17、ISIC18和Synapse数据集上评估了VM-UNet在医学图像分割任务上的性能。

Datasets

ISIC17和ISIC18数据集: 国际皮肤成像协作2017和2018挑战数据集(ISIC17和ISIC18)是两个公开可用的皮肤病变分割数据集,分别包含2,150和2,694张带分割 Mask 标签的皮肤镜图像。按照之前的工作,作者将数据集以7:3的比例分为训练集和测试集。

具体来说,对于ISIC17数据集,训练集由1,500张图像组成,测试集由650张图像组成。对于ISIC18数据集,训练集包括1,886张图像,而测试集包含808张图像。对于这两个数据集,作者提供了在几个指标上的详细评估,包括平均交并比(mIoU),Dice相似性系数(DSC),准确度(Acc),敏感性(Sen)和特异性(Spe)。         whaosoft aiot http://143ai.com

Synapse多器官分割数据集(Synapse):Synapse是一个公开的多器官分割数据集,包含30个腹部CT病例,共3,779张轴向腹部临床CT图像,包括8种腹部器官(主动脉、胆囊、左肾、右肾、肝脏、胰腺、脾脏、胃)。遵循之前的工作的设置,其中18个病例用于训练,12个病例用于测试。对于这个数据集,作者报告了Dice相似性系数(DSC)和95% Hausdorff距离(HD95)作为评估指标。

Implementation details

遵循之前的工作,作者将ISIC17和ISIC18数据集中的图像大小调整为256×256,将Synapse数据集中的图像大小调整为224×224。为防止过拟合,采用了数据增强技术,包括随机翻转和随机旋转。对于ISIC17和ISIC18数据集,使用BceDice损失函数,而对于Synapse数据集,采用CeDice损失函数。作者将批量大小设置为32,并采用AdamW优化器,初始学习率为1e-3。使用CosineAnnealingLR作为调度器,最大迭代次数为50,最小学习率为1e-5。训练周期设置为300。对于VM-UNet,作者用VMamba-S的权重初始化编码器和解码器的权重,VMamba-S是在ImageNet-1k上预训练的。所有实验都在单个NVIDIA RTX A6000 GPU上完成。

Main results

作者对比了VM-UNet与一些最先进模型的性能,并在表1和表2中展示了实验结果。

对于ISIC17和ISIC18数据集,VM-UNet在mIoU、DSC和Acc指标上超过了其他模型。对于Synapse数据集,VM-UNet也取得了具有竞争力的表现。例如,VM-UNet在DSC和HD95指标上分别超过了Swin-UNet——这是第一种纯粹的基于Transformer的模型——1.95%和2.34mm。这些结果证明了基于SSM的模型在医学图像分割任务中的优越性。

Ablation studies

在本节中,作者使用ISIC17和ISIC18数据集对VM-UNet的初始化进行了消融实验。作者分别用来自VMamba-T和VMamba-S的预训练权重来初始化VM-UNet。如表3所示,实验结果表明,更强大的预训练权重显著提高了VM-UNet在下游任务中的性能,这表明VM-UNet在很大程度上受到预训练权重的影响。

5 Conclusions and Future works

结论:在本文中,作者首次引入了一种基于纯SSM的医疗图像分割模型,并提出VM-UNet作为 Baseline 。为了发挥基于SSM模型的潜能,作者使用VSS块构建了VM-UNet,并用预训练的VMamba-S初始化其权重。在皮肤病变和多器官分割数据集上进行的全面实验表明,基于纯SSM的模型在医疗图像分割任务中具有高度竞争力,并在未来值得深入探索。

未来工作:

基于SSM机制的特性,设计更适合分割任务的模块。

VM-UNet的参数数量约为3000万,这为通过人工设计或其他压缩策略简化SSM提供了机会,从而加强它们在现实医疗场景中的应用性。

考虑到SSM在捕捉长序列信息方面的优势,进一步研究在高分辨率下的分割性能将具有价值。

探索SSM在其他医疗成像任务中的应用,如检测、配准和重建等。

本文链接:https://www.kjpai.cn/news/2024-04-30/163745.html,文章来源:网络cs,作者:康由,版权归作者所有,如需转载请注明来源和作者,否则将追究法律责任!

版权声明:本文内容由互联网用户自发贡献,该文观点仅代表作者本人。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如发现本站有涉嫌抄袭侵权/违法违规的内容,一经查实,本站将立刻删除。

文章评论