Beyond Adapting SAM:Towards End-to-EndUltrasound Image Segmentation via AutoPrompting自动提示实现端到端超声图像分割
今天看到一篇关于医学超声图像分割的研究工作AutoSAMUS,这是一个基于SAM的通用模型,专门用于超声图像分割,并进一步扩展为端到端的自动分割模型AutoSAMUS。主要工作包括:
-
SAMUS模型:
-
继承了SAM的ViT图像编码器、提示编码器和掩码解码器。
-
通过引入并行的CNN分支和跨分支注意力模块(CBA)来补充局部信息。
-
使用特征适配器和位置适配器将SAM从自然图像域适应到超声图像域。
-
-
AutoSAMUS模型:
-
通过引入自动提示生成器(APG),替代手动提示编码器,实现自动生成任务相关的提示嵌入。
-
使得SAMUS能够在端到端的方式下进行自动分割。
-
-
实验验证:
-
构建了一个包含约3万张图像和6.9万个掩码的大规模超声数据集US30K。
-
通过广泛的对比实验,证明了SAMUS和AutoSAMUS在通用性和泛化能力上优于最先进的任务特定和基于SAM的基础模型。
-
展示了AutoSAMUS在端到端自动分割任务中的优越性能。
-
这里主要记录阅读的记录,感兴趣的话可以参考一下,如果想要知己阅读原文,可以来这里,如下所示:

摘要
端到端的医学图像分割对于以任务特定模型为主导的计算机辅助诊断具有重要价值,但通常存在泛化能力差的问题。随着分割任何模型(SAM)在通用图像分割领域的最新突破,大量工作致力于将SAM适应于医学成像,但仍面临两个主要问题:1)没有适当适应的情况下性能严重下降和泛化能力有限;2)依赖于准确的手动提示进行交互的半自动分割。在这项工作中,我们首先提出了SAMUS,这是一个专门为超声图像分割定制的通用模型,并进一步将其扩展为端到端的方式,称为AutoSAMUS。具体来说,在SAMUS中,引入了一个并行的CNN分支,通过跨分支注意力补充局部信息,并联合使用特征适配器和位置适配器将SAM从自然域适应到超声域,同时降低训练复杂度。通过引入自动提示生成器(APG)来替代SAMUS的手动提示编码器,自动生成提示嵌入,实现了AutoSAMUS。我们收集了一个综合的超声数据集,包含约3万张图像和6.9万个掩码,涵盖六个对象类别,用于验证。广泛的对比实验证明了SAMUS和AutoSAMUS相对于最先进的任务特定和基于SAM的基础模型的优越性。我们相信,基于自动提示的SAM模型有潜力成为端到端医学图像分割的新范式,值得进一步探索。代码和数据可在这里,如下所示:

关键词:SAM 自动提示 基础模型 医学图像分割
1 引言
医学图像分割是一项关键技术,用于在医学图像中识别和突出显示特定的器官、组织和病变,是计算机辅助诊断系统的重要组成部分[1]。许多深度学习模型已被提出用于医学图像分割,展示了巨大的潜力[2, 3]。然而,这些模型是为特定对象量身定制的,当应用于其他对象时需要训练新的模型参数,导致在临床应用中处理多样任务时非常不便。
分割任何模型(SAM)作为一种通用的视觉分割基础模型,因其对各种对象的出色分割能力和强大的零样本泛化能力而备受赞誉[4]。根据用户提示,包括点、边界框和粗略掩码,SAM能够分割相应的对象。因此,通过简单的提示,SAM可以轻松适应各种分割应用。这种范式使得多个独立的医学图像分割任务可以集成到一个统一的框架中(即通用模型),极大地促进了临床部署[5]。
尽管构建了迄今为止最大的数据集(即SA-1B),但由于缺乏可靠的临床标注,SAM在医学领域的表现迅速下降[5]。一些基础模型通过在医学数据集上调整SAM来适应医学图像分割[6, 8]。然而,与SAM一样,它们在特征建模之前对输入图像进行无重叠的16×分块,这破坏了识别小目标和边界的关键局部信息,使得它们难以分割形状复杂/线状、边界模糊、尺寸小或对比度低的临床对象。此外,这些基于SAM的模型需要手动提供任务相关的提示来生成相应的掩码,导致半自动分割流程。这种范式在处理某些临床任务时不够灵活。
在本文中,我们首先提出了SAMUS,将SAM强大的特征表示能力转移到医学图像分割领域,然后将其扩展为自动版本(即AutoSAMUS),以灵活处理各种下游分割任务。具体来说,SAMUS继承了SAM的ViT图像编码器、提示编码器和掩码解码器,并对图像编码器进行了定制设计。首先,我们通过减少所需的输入尺寸来缩短ViT分支的序列长度,以降低计算复杂度。然后,开发了特征适配器和位置适配器,以微调ViT图像编码器从自然域到医学域的适应。为了补充ViT图像编码器中的局部(即低级)信息,我们引入了一个并行的CNN分支图像编码器,与ViT分支并行运行,并提出了一个跨分支注意力模块,使ViT分支中的每个分块能够吸收CNN分支的局部信息。构建了一个名为US30K的大规模超声数据集,用于全面训练和评估SAMUS的有效性。在获得训练好的SAMUS后,期望其能够以端到端的方式运行于下游特定任务。因此,我们通过引入具有可学习任务标记的自动提示生成器(APG)来替代SAMUS的手动提示编码器,生成任务相关的提示嵌入,从而将SAMUS扩展为AutoSAMUS。实验结果表明,SAMUS优于最先进的(SOTA)任务特定和通用分割方法。更重要的是,基于训练好的SAMUS,在特定任务上调整AutoSAMUS可以实现端到端的自动分割,并比SOTA任务特定方法获得更好的分割性能。这表明开发基于自动提示的SAM模型作为新的端到端分割范式具有潜力。
2 相关工作
将SAM适应于医学图像分割。 SAM在自然图像中表现出色,但在一些医学图像分割任务中表现不佳,特别是在对象形状复杂、边界模糊、尺寸小或对比度低的情况下[5]。为了弥合这一差距并使SAM有效适应医学图像领域,一些方法通过应用视觉调整技术来调整SAM。具体来说,MedSAM通过冻结提示编码器,专注于调整图像编码器和掩码解码器来训练SAM[6]。SAMed在图像编码器上应用基于低秩(LoRA)的策略,以较低的计算成本调整SAM,使其更适用于医学图像分割[7]。MSA在ViT图像编码器和掩码解码器上采用下-ReLU-上适配器引入医学信息[8]。与当前的基于SAM的通用模型相比,提出的SAMUS更注重补充局部特征并实现端到端的自动分割。
SAM中的提示。 原始的SAM在精确的空间提示(如点、边界框和掩码)驱动下生成任务相关的掩码。为了自动获取这些空间提示,一些方法引入了单独的输入相关网络。具体来说,AdapterShadow、SAC和UV-SAM分别使用EfficientNet、U-Net和SegFormer生成粗略掩码以生成空间提示[9, 10, 11]。这些单独网络引入的参数与任务特定方法相当,使得通用模型变得笨重。Polyp-SAM++使用Grounding DINO从文本提示生成边界框提示[12]。Adaptive SAM和SP-SAM通过CLIP将文本提示编码为提示嵌入[13, 14]。尽管这些方法可以有效利用文本信息,但在医学场景中缺乏文本-图像数据来调整在自然场景中训练的文本编码器。SurgicalSAM提出了一种基于原型的类提示编码器,用于生成密集和稀疏的提示嵌入[15]。Auto-prompting SAM通过构建上下全卷积层开发了一个自动提示编码器[16]。这些提示编码器与掩码解码器深度耦合,使得通过多目标学习为基于SAM的模型构建鲁棒的特征表示变得困难。相比之下,提出的APG是一个轻量级且独立的模块,高度可扩展到其他基于SAM的基础模型。

图1:提出的SAMUS的概述。APG表示将SAMUS扩展为AutoSAMUS的扩展模块。
3 方法
SAMUS的整体架构。 如图1所示,SAMUS的整体架构继承自SAM,保留了提示编码器和掩码解码器的结构和参数,无需任何调整。相比之下,图像编码器经过精心修改,以解决局部特征不足和计算内存消耗过大的挑战,使其更符合临床需求。主要修改包括减少所需的输入尺寸、重叠分块嵌入、在ViT分支中引入适配器、添加CNN分支以及引入跨分支注意力(CBA)。具体来说,输入空间分辨率从1024×1024像素缩小到256×256像素,由于变换器中输入序列较短,GPU内存成本大幅降低。重叠分块嵌入使用与SAM中分块嵌入相同的参数,但其分块步幅为原始步幅的一半,很好地保留了分块边界的信息。ViT分支中的适配器包括一个位置适配器和五个特征适配器。位置适配器用于适应由于输入尺寸较小而导致的较短序列的全局位置嵌入。第一个特征适配器跟随重叠分块嵌入,以对齐输入特征与预训练ViT图像编码器所需特征分布。其余的特征适配器附加在全局变换器的残差连接上,以微调预训练图像编码器。在CNN分支方面,它与ViT分支并行,通过CBA模块为后者提供互补的局部信息,CBA模块以ViT分支特征为查询,与CNN分支特征建立全局依赖关系。需要注意的是,CBA仅集成在每个全局变换器中。最后,两个分支的输出结合为SAMUS的最终编码图像嵌入Fi。

图2:自动提示生成器的详细信息。
ViT分支中的适配器。 为了促进训练好的图像编码器(即ViT分支)的泛化能力,使其适应较小的输入尺寸和医学图像领域,我们引入了位置适配器和五个特征适配器。这些适配器可以有效调整ViT分支,同时仅需要很少的参数。具体来说,位置适配器负责调整位置嵌入以匹配嵌入序列的分辨率。它首先通过步幅和核大小为2的最大池化对位置嵌入进行下采样,以达到与嵌入序列相同的分辨率。然后,应用核大小为3×3的卷积操作来调整位置嵌入,进一步帮助ViT分支更好地处理较小的输入。所有特征适配器具有相同的结构,包括三个组件:下线性投影、激活函数和上线性投影。每个特征适配器的过程可以表示为:

CNN分支。 CNN分支由一系列卷积-池化块顺序连接组成。具体来说,输入首先通过一个单独的卷积块,然后通过三个卷积-池化块进行处理。然后,CNN分支中的特征图与ViT分支的特征图具有相同的空间分辨率。在CNN分支的其余部分,这样的单独卷积块按顺序重复四次。每个单独的卷积块包含一个核大小为3×3的卷积,每个卷积-池化块包含一个步幅和核大小为2的最大池化和一个单独的卷积块。更多细节如图1所示。CNN分支的这种极简和轻量级设计是为了防止训练过程中过拟合。
表2:基础模型在分割甲状腺结节(DDTI)、乳腺癌(UDIAT)和心肌(HMC-QU)上的泛化能力比较。

跨分支注意力。 跨分支注意力(CBA)模块在CNN分支和ViT分支之间建立桥梁,进一步补充ViT分支缺失的局部特征。对于来自ViT分支Fv和CNN分支Fc的一对特征图,单头跨分支注意力可以表示为:




图3:在可见数据集(蓝色标记)和未见数据集(橙色标记)上评估的SAMUS与任务特定方法的比较。
4 实验
数据集。 为了全面评估SAMUS的有效性,我们构建了一个名为US30K的大规模超声数据集,如补充材料中所总结,包含来自七个公开可用数据集的数据,包括TN3K[17]、DDTI[18]、TG3K[19]、BUSI[20]、UDIAT[21]、CAMUS[22]和HMC-QU[23]。TN3K和TG3K的数据按照TRFE[17]划分为训练集、验证集和测试集。BUSI随机分为7:1:2用于训练、验证和测试。首先根据挑战[22]将CAMUS划分为训练集和测试集。然后,我们从训练集中随机选择10%的患者进行模型验证,其余数据作为最终训练数据。为了评估不同模型的泛化能力,US30K中的其他数据集在训练期间未见。通过在US30K的整个训练集上训练并分别在不同任务上进行评估,比较了SAMUS与其他基础模型。为了公平比较,所有基础模型在相同设置下重新实现并训练400个epoch,使用相同的单点提示。通过在每个单独数据集上进行训练,比较了SAMUS与SOTA任务特定方法。
表3:AutoSAMUS与SOTA任务特定方法的定量Dice(%)比较。AutoSAMUS-是仅更新APG时的模型。

表4:SAMUS不同组件组合的消融研究。F-Adapter和P-Adapter分别表示特征适配器和位置适配器。

比较SAMUS与SOTA基础模型。 表1和表2总结了基础模型在US30K上的通用性和泛化能力的比较。在比较的基础模型中,MSA在通用性和泛化能力方面表现最佳。与MSA相比,SAMUS在US30K的所有子数据集上均取得了显著改进。这表明CNN分支和CBA模块在补充局部信息方面对SAMUS的价值,这对于医学图像分割至关重要。
比较SAMUS与SOTA任务特定方法。 如图3所示,包括12种SOTA方法[2, 3, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33]进行比较。SAMUS在包括TN3K、BUSI、CAMUS-MYO、DDTI、UDIAT和HMC-QU的所有数据集上均优于最佳比较方法,Dice平均增加了1.97%、3.55%、0.15%、7.06%、12.22%和7.42%。这证明了通过SAMUS适应SAM到医学图像领域的必要性。
比较AutoSAMUS与SOTA任务特定方法。 为了比较AutoSAMUS与任务特定方法,我们首先将训练好的SAMUS的参数加载到AutoSAMUS中。然后,我们在三个下游任务上微调APG以及APG与SAMUS的可学习部分,结果分别由AutoSAMUS-和AutoSAMUS表示在表3中。AutoSAMUS-在DDTI和UDIAT上优于最佳比较方法,在HMC-QU上接近最佳比较方法,而AutoSAMUS在DDTI、UDIAT和HMC-QU上优于最佳比较方法,Dice平均增加了4.26%、2.5%和0.44%。
SAMUS中各组件的有效性。 如表4所示,结合SAMUS的任何组件都可以有效提高分割性能和泛化能力。结合所有四个组件,SAMUS在分割和泛化性能上表现最佳。
5 结论
在本文中,我们提出了SAMUS,一个从SAM派生的通用基础模型,及其端到端版本AutoSAMUS,用于临床友好且可泛化的超声图像分割。具体来说,我们提出了一个CNN分支图像编码器、一个特征适配器、一个位置适配器和一个跨分支注意力模块,以丰富超声对象的特征表示。此外,为了促进下游任务的临床应用,我们将SAMUS与自动提示生成器结合,实现了自动分割,取代了原始SAM的手动提示,实现了新的端到端分割范式。构建了一个包含3万+图像和6.8万+掩码的大规模超声图像数据集US30K,用于评估和潜在的临床使用。广泛的实验证明了SAMUS和AutoSAMUS的出色性能,优于SOTA的基于SAM的医学基础模型和任务特定模型。
DAMO开发者矩阵,由阿里巴巴达摩院和中国互联网协会联合发起,致力于探讨最前沿的技术趋势与应用成果,搭建高质量的交流与分享平台,推动技术创新与产业应用链接,围绕“人工智能与新型计算”构建开放共享的开发者生态。
更多推荐


所有评论(0)