Sam2Rad 带可学习提示的图像分割模型 !

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

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

picture.image

picture.image

像Segment Anything模型(SAM)这样的基础模型在医学图像分割中的应用需要高质量的 manually生成的提示,但这些提示的生成耗时且需要医学专业知识。即使使用稀疏提示(如框、点或文本),以及密集提示(如SAM或其医学图像相关变种的细调,如MedSAM)、SAM 2等,仍无法在超声(US)图像中分割骨头,因为在显著的领域转移下。作者提出了一种新的提示学习方法,将SAM、其最近的一个变种SAM 2和其他变种 Adapt到分割US图像中的骨质区域,无需人工提示。

作者引入了一个提示预测网络(PPN)带有轻量级交叉注意力模块,以增强现有的提示编码器,并直接预测来自图像编码器的提示嵌入。PPN输出感兴趣区域的边界框和 Mask 提示以及256维的提示嵌入。

作者的新框架还允许手动提示,可以在Mask解码器的输入中与学习得到的提示合并。PPN和Mask解码器可以使用参数高效的微调方法端到端训练。为了保留SAM的广泛世界知识,作者希望保持SAM的所有模块冻结,只训练提示预测网络。作者通过冷冻SAM的所有参数并实现与微调Mask解码器相当的对骰子得分的有效性来示范提示预测网络的有效性。作者的模型可以在没有人监管的情况下自由使用,半自主地与人类反馈,或完全手动使用,如原始SAM。

在半自主模式中,模型预测 Mask ,如果结果不令人满意,用户可以提供额外的提示(框、点或笔刷)来完善预测。自主设置可以用于实时应用,而半自主设置适用于数据标注,或用于主动学习框架。作者已在3个肌肉骨骼US数据集上测试了提出的模型Sam2Rad,这些数据集包括手腕(3822张图像)、肩袖(1605张图像)和臀部(4849张图像)。

如果没有Sam2Rad,所有SAM 2变体都无法在没有人工框启用的零样本泛化中分割肩部US,而SAM 2在其他US数据集的分割精度优于SAM。作者的模型Sam2Rad改进了所有SAM基本网络在所有数据集上的性能,无需手动提示。改进的程度因髋或手腕的Dice得分(3-7%的增加)而到肩部的33%(从49%到82%的Dice得分提升)。

值得注意的是,Sam2Rad可以在10张标记图像的情况下进行训练。Sam2Rad可以与任何SAM架构兼容,并可用于自动分割。

代码可以在https://github.com/aswahd/SamRadiology

1 Introduction

医学图像中结构的自动化分割是辅助诊断和治疗规划过程的必要步骤。在过去十年中,深度学习模型的发展显著加速了自动医学图像分割的进步。传统上,分割任务通过专门的网络分别处理 。因此,这种方法在更广泛的临床环境中应用有限,因为它需要大量的数据来训练每个任务,并忽略了不同分割目标之间的相关性。同时,深度学习的扩展定律表明,增加训练数据集的大小和模型的参数数量可以提高准确性 从而形成基础模型。这些模型通常在包含广泛一般信息的广泛数据集中进行训练,并封装了可以相对轻松地适应特定应用的世界知识谱。基础模型已经被证明在带有很少标记图像的下游任务中表现良好,对 adversarial 攻击更具鲁棒性,对抗域漂移更具有弹性,并且可以在零样本环境中有效工作。因此,基础模型在医学图像中变得越来越受欢迎,其中大规模的标记数据集通常很稀缺。

2 Related Works

Foundation Models

CLIP[16]是由OpenAI公司推广的一种广泛使用的foundation模型。这些模型通过在大规模数据集(例如数百万张文本图像对)上进行深度学习训练,可以有效地在不同下游任务间进行泛化。这些模型通常通过在大数据集上使用自监督学习方法进行预训练。例如,CLIP就是在数百万张文本与图像对上进行训练,以创建图像和文本的联合嵌入空间。它已经被广泛应用于带文本提示的图像分类、目标检测、图像分割以及各种自然语言处理任务(NLP)。

Segment Anything Model (SAM)[17]是一种新兴的foundation模型,用于自然图像分割。它提供了一个框架,可以在一次分割任务中处理多个分割提示,例如点、边界框、遮挡符或文本。SAM已经激发了大量的后续工作,导致开发了用于医学图像分割的通用模型,如。最新版本的SAM2[15],专门设计用于图像和视频分割。当在帧中给出提示时,它可以跟踪目标在视频帧之间的运动。值得注意的是,SAM2比其前身快6倍。它包括一个记忆注意机制,允许当前帧的预测受之前帧图像嵌入和预测的条件化影响。

然而,在医学图像领域应用SAM的研究发现其存在特定未见过的医学模态的局限性。为了应对这一挑战,人们已经尝试在医学领域的大型精选数据集上对SAM进行微调,从而得到医学foundation模型,例如。然而,这些模型的性能仍然受其训练模态的限制,且严重依赖提示的质量。这意味着当这些模型应用于完全不同的数据分布(如超声)时,其表现力可能不佳。解决这些局限性可能需要针对特定领域的微调,以实现最佳性能,同时保持广泛的应用性。此外,这些foundation模型需要手动提示,而在没有必要医学专业知识的情况下,这在实践中是困难的。因此,自动生成适当提示的深度学习方法已经成为研究的重要焦点。

Prompt Learning

传统的深度学习模型训练方法通常涉及两种情况:从头训练模型或先使用无标签数据预训练模型,接着使用特定任务的标注数据进行微调。相比之下,基于提示的学习模型可以通过用户提供的提示自动适应各种任务。这些模型利用提示引入的领域知识来指导其性能,使任务适应更加高效和灵活,而无需进行大量特定任务的训练。

SAM在自然图像的零样本泛化方面表现出色。然而,在推理阶段,它仍需要提示(如点、框、文本或 Mask )来分割图像中的目标,即使是推理阶段。提示可以由人类提供(例如,在目标周围绘制框或输入文本),也可以由独立模型预测。基于SAM图像编码器生成的上下文信息,Grounded SAM [39]使用Grounding DINO [40]为图像中的所有目标生成边界框,然后用于提示SAM。这种方法的局限性在于,它需要一个从头训练的独立图像编码器模型。

最近,自动方法如PerSAM [41]和Matcher [42]被提出用于为自然图像生成边界框提示。这些方法假设前景和背景 Patch 是不相似的,并使用参考图像和测试图像之间的 Patch 相似性来生成提示。然而,在US图像中,前景和背景 Patch 往往很难区分,导致不相关的 Patch 之间具有很高的余弦相似性。作者提出了一种新方法,使用SAM的图像编码器学习提示,而不是依赖其他独立模型如Grounding DINO [40]。作者开发了一个提示预测网络(PPN),它输入SAM图像编码器提取的图像特征,并预测目标区域的适当提示。

相反,作者提出了一种轻量级的注意力机制来学习预测目标区域在图像中的位置的非线性函数。提示预测网络使用轻量级注意力模块,基于SAM图像编码器获得的图像特征预测目标区域的可选提示。

据作者所知,这是首次报告SAM2(SAM的最新版本)的提示学习工作。见图4比较作者的提示预测网络与其他提示技术。

picture.image

Parameter-efficient Fine-tuning

一些研究已经表明,即使有手动提示,SAM仍无法扩展到医学图像。在作者遇到的问题中,如图2所示,超声图像(US)也出现了类似的问题。直接解决这个问题的一种方法是调整所有模型的参数--普通的 fine-tuning,但这种方法通常是不最优的,常常需要大量的标注数据,有时与参数高效的 fine-tuning(PEFT)方法相比表现较差。

例如,CoOp [43] 报告了在 CLIP 图像编码器上普通 fine-tuning后分类准确率降低了40%。PEFT 是一种更有效的方式适配基础模型到下游任务。PEFT 方法如 adapter 和低秩适配(LoRA) [44] 都是设计用来适配基础模型到下游任务,最小程度上改变模型的参数。PEFT 方法在标注数据有限时尤为有用。adapter [27] 在基础模型的 transformer 层中并行插入小型模块,只微调这些新的参数。LoRA [23] 是一种其他 PEFT 方法,仅微调模型的一个小子集参数。为了避免在有手动提示时,SAM 或 SAM2 在图像分割中失败,可以冻结图像编码器参数同时使用 LoRA 微调 mask 解码器。研究表明,将 PEFT 与 SAM 结合不仅降低了计算成本,还减少了过拟合并提高了整个分割的准确度 [25]。

picture.image

Key Contributions

作者提出了一种新的PPN,它能直接预测N个256维的嵌入,同时提供边界框坐标和 Mask 提示。这种方法使网络可以学习一个提示的抽象表示,而不是像方形这种稀疏表示。此外,N可以设置为任意值,与其他方法相比,通常是预测每个目标一个边界框。研究表明,将SAM适配到医学影像通常需要多个点的提示。在PPN中,N可以代表任意数量的点或其他与目标相关的信息。例如,对于难以分割的目标,如薄叠加的的组织,网络可以通过预测多个点提示或使用其他表示来进行自适应。

3 Methodology

作者提出了一种名为SAM2Rad的创新框架,用于在不同提示下分割医学图像。本文主要针对将其应用于超声成像。SAM2Rad包括一个提示预判网络(生成器)(PPN),可预测候选边界框 Proposal 、 Mask 提示和 256维提示。类似于Mask R-CNN的目标检测,其中边界框回归头使用 Backbone 网络的特征来预测坐标,PPN的边界框回归头可以被视为目标检测。它利用SAM图像编码器产生的特征来为分割目标区域生成适当的提示,类似于Mask R-CNN的目标 Proposal 网络(RPN)。

Sam2Rad

SAM2Rad框架包括一个预训练的SAM/SAM2编码器、 Mask 解码器、提示编码器和提示预测器网络(如图1所示)。

picture.image

图像编码器:图像编码器从输入图像中提取特征。它基于视觉 Transformer (ViT)架构[45]并使用masks autoencoder(MAE)[46]进行预训练。对于大小为的输入图像,图像编码器输出大小为的嵌入。

Mask 解码器: Mask 解码器将图像嵌入和提示嵌入转换为最终的 Mask 。

提示编码器:提示编码器处理稀疏提示(点和对角线)和密集提示( Mask )。对于点,将其编码为256维嵌入,结合位置和学习的嵌入。对于对角线,使用对角线位置编码和对角线学习的嵌入。 Mask ()通过一个下采样卷积块转换为的嵌入。

为了使学习的提示与手工提示的SAM训练对齐,提示预测器网络输出边界框坐标、 Mask 提示和个256维嵌入。如图4展示了作者的提示预测器网络与其他提示技术 comparison。

在下一节中,作者将描述提示预测器的的设计。

3.1.1 Prompt predictor network (prompt learning)

提示(可学习)嵌入 和参考图像的嵌入 之间的相似性。

picture.image

给定尺寸为 的图像 ,SAM图像编码器提供嵌入 。 中的每个点对应图像中的一个 Patch 。 PerSAM [41] 和 Matcher [42] 使用这个特征图来计算余弦相似性。如介绍中所述1,US图像往往呈现出前景和背景 Patch 之间显著的相似性,导致无关 Patch 之间的余弦相似性较高。因此,作者提出了一种轻量级注意力机制来学习预测给定 中的 Patch 嵌入的非线性函数以预测 中目标区域的位置,该 Patch 嵌入可以用 中的位置。为此,作者将可学习位置嵌入( Token )表示为 ,其中 是可学习 Token 的数量, 是通道数。

雅克布·L·杰雷姆科博士是加拿大的CIFAR AI主席,他的学术时间由埃德蒙顿的医疗影像顾问公司(MIC)提供。作者感谢TD准备金、IC-IMPACTS、One Child Every Child、类风湿关节炎协会以及阿尔伯塔省的紧急战略临床网络为作者的研究提供了财务支持,为临床扫描提供了支持,并为作者提供了计算资源,包括高性能GPU、用于训练和测试作者的深度学习模型。

总结

在这篇论文中,介绍了Sam2Rad,这是一个新颖的提示学习方法,它显著提高了SAM/SAM2及其变体在超声图像分割方面的性能。通过消除对手动提示的需求并提高各种数据集上的分割精度,解决了可提示模型的关键局限性。Sam2Rad在零样本泛化方面优于所有SAM和SAM2变体,特别是在具有挑战性的肩部超声病例中。

该模型在Dice分数方面取得了显著的改进,对于髋部和手腕数据集,Dice分数提高了2-7%,对于肩部数据,从49%增加到82%,提高了33%。值得注意的是,Sam2Rad可以仅用10张标记的图像进行训练。通过利用SAM的图像编码器特征,网络无缝地整合到SAM的架构中,以产生高质量的提示。

提供三种模式:自主模式、半自主模式(带有人工智能)和完全手动模式。与SAM架构的无缝整合和与任何SAM变体的兼容性使Sam2Rad成为医学成像中自动分割的强大和多功能的工具。

参考

[1].Sam2Rad: A Segmentation Model for Medical Images with Learnable Prompts.

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

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