WeakMedSAM: Weakly-Supervised Medical Image Segmentation via SAM——基于 SAM 的弱监督医学图像分割,通过子类探索和提示亲和力挖掘

这篇文章的核心内容是提出了一种名为 WeakMedSAM 的弱监督医学图像分割模型,旨在减少医学图像分割任务中对大量像素级标注数据的依赖,同时利用 Segment Anything Model(SAM)的强大能力。以下是文章的主要研究内容总结:

研究背景

  • 医学图像分割的挑战:医学图像分割通常需要大量精确的像素级标注数据,这些数据的获取成本高昂,且需要专业医生的知识。因此,减少标注成本的需求迫切。

  • 弱监督学习(WSS)的优势:弱监督学习通过使用图像级标签等弱监督信息来替代像素级标注,显著降低了数据标注成本。

  • SAM 的应用:SAM 在自然图像分割中表现出色,但在医学图像分割中存在挑战,如复杂的形状、模糊的边界、小目标或低对比度等。

WeakMedSAM 模型

文章提出了 WeakMedSAM,一个基于 SAM 的弱监督医学图像分割模型,包含两个主要模块:

  1. 子类探索模块(Sub-Class Exploration, SCE)

    • 通过将每个主类(如“有肿瘤”和“无肿瘤”)细分为多个子类(如肿瘤的不同潜在变体),来缓解医学图像中类内共现问题。

    • 使用无监督聚类生成子类标签,并结合主类分类任务进行联合训练,从而学习到更准确的特征表示。

    • 实验表明,优化子类分类头可以显著提升类激活图(CAM)的质量。

  2. 提示亲和力挖掘模块(Prompt Affinity Mining, PAM)

    • 利用 SAM 的提示能力,通过在图像上应用网格点提示来获取亲和图(表示像素间的关系)。

    • 通过随机游走算法,利用亲和图对 CAM 进行细化,增强医学图像中的结构信息,而无需额外训练。

实验设计与结果

  • 数据集:在三个广泛使用的基准数据集上进行实验,包括 BraTS 2019(脑肿瘤分割)、AbdomenCT-1K(腹部器官分割)和 MSD Cardiac(心脏分割)。

  • 性能指标:使用 Dice 系数(DSC)、Jaccard 指数、平均对称表面距离(ASSD)和 Hausdorff 距离(HD95)等指标评估分割性能。

  • 结果

    • 在 BraTS 2019 数据集上,使用 SAMUS 和 EfficientSAM 时,分别达到了 79.69% 和 77.25% 的 Dice 分数,显著优于其他弱监督方法(如 TOCO 的 74.61%)。

    • 在 AbdomenCT-1K 数据集上,Dice 分数分别达到了 75.87% 和 71.15%。

    • 在 MSD Cardiac 数据集上,Dice 分数分别达到了 58.05% 和 57.62%。

    • 在小目标分割任务(如 BraTS-Core 数据集)上,Dice 分数达到了 61.85%,显示出对小目标分割的优势。

消融研究与分析

  • 模块有效性:通过消融实验验证了 SCE 和 PAM 模块的有效性。仅使用 ViT-b SAM 时,DSC 为 68.13%,引入 SCE 和 PAM 后,DSC 分别提升至 74.91% 和 76.31%,结合两者后提升至 79.69%。

  • 不同特征提取器:使用不同的预训练网络进行子类聚类,发现模型对特征提取器的性能不敏感,最终选择 ResNet18 以减少计算开销。

  • 子类数量的影响:实验表明,子类数量为 8 时性能最佳,过多的子类会降低性能。

  • 聚类算法的随机性:通过多次随机种子实验,证明了聚类算法对随机初始化的鲁棒性。

  • 与其他细化方法的比较:PAM 模块在不增加额外训练的情况下,优于其他需要额外训练或不利用结构信息的方法。

交互式 WeakMedSAM

  • 在交互式分割场景中,使用单点提示和边界框提示进行实验,发现 WeakMedSAM 在单点提示下显著优于 SAM,显示出其在利用有限交互信息方面的优势。

结论与未来工作

  • 结论:WeakMedSAM 在减少标注成本的同时,利用 SAM 的能力,显著提升了弱监督医学图像分割的性能,特别是在小目标分割任务上表现出色。

  • 未来工作:探索端到端弱监督微调 SAM 的潜力,进一步提升性能和泛化能力,重点是改善 SCE 和 PAM 模块之间的协同作用,以更高效地适应交互式医学图像分割。

这里是自己的论文阅读记录,感兴趣的话可以参考一下,如果需要阅读原文的话可以看这里,如下所示:

官方项目地址在这里,如下所示:

摘要

我们在计算机视觉任务中见证了基础模型的显著进展。目前,已有几项工作利用分割任何模型(SAM)来提升医学图像分割的性能,其中大多数工作集中在通过全监督方式对大量像素级标注的医学图像进行微调。在本文中,为了减少标注成本,我们研究了一种新颖的基于 SAM 的弱监督分割模型,即 WeakMedSAM。具体而言,我们提出的 WeakMedSAM 包含两个模块:1)为缓解医学图像中严重的共现问题,引入子类探索模块以学习准确的特征表示;2)为提升类别激活图的质量,我们的提示亲和力挖掘模块利用 SAM 的提示能力获取亲和力图,用于随机游走细化。我们的方法可以应用于任何 SAM 类的骨干网络,我们在 SAMUS 和 EfficientSAM 上进行了实验。在三个常用的基准数据集(BraTS 2019、AbdomenCT-1K 和 MSD Cardiac 数据集)上的实验结果表明,我们提出的 WeakMedSAM 具有很好的前景。

关键词—医学图像分割,弱监督分割,分割任何模型,子类探索,提示亲和力挖掘

I. 引言

分割任何模型(SAM)[1] 在计算机视觉领域取得了显著的成功,并且在采用 SAM 进行各种下游分割任务方面引起了越来越多的兴趣 [2]–[4]。其中,将 SAM 适应于医学图像分析领域 [5] ——这是一个长期且重要的方向——正受到越来越多的关注。这些基于 SAM 的模型 [5]–[8] 的共同目标是促进 SAM 有效且高效地适应医学图像分割。由于直接使用 SAM 分割医学图像的性能有时由于任务差异而不尽如人意,这些方法,例如 SAMed [7]、Med-SA [9] 和 MedSAM [8],利用下游医学数据集对 SAM 进行微调,展示了将 SAM 整合到医学图像分割中的潜力。最新工作如 Med-SA [10]、One-Prompt SAM [11] 和 SegAnyPath [12] 已经开发出更高效的适应 SAM 到医学图像的技术。然而,观察到这些方法遵循全监督范式,需要大量的医学图像数据集,且需要手动勾勒出像素级分割标签。众所周知,获取像素级标注是一个劳动密集型且耗时的过程,需要经验丰富的医生提供专业知识,这极大地阻碍了 SAM 适应分割在不同临床场景中的部署。这一观察结果突显了采用更少劳动密集型方法将 SAM 适应于医学成像领域的必要性。作为一种流行的替代方案,以摆脱大量的标注负担,弱监督分割(WSS)范式 [13] 越来越受到认可。WSS 通常采用弱监督技术,例如图像级标签 [14]–[16]、点 [17]–[19]、涂鸦 [20]–[22] 或边界框 [22]–[24],作为绕过全面像素级监督的手段。在这些弱监督方法中,我们在本文中关注图像级标签,因为它们简单且无歧义,因为图像级标签是唯一的,而其他弱监督方法(如点)可能会引入可能的歧义。通过我们的实证和实验观察,我们已经确定了与基于 SAM 的框架用于弱监督分割相关的几个挑战:a)使用弱监督的分割结果主要依赖于类间信息,即不同类别之间的区别(例如,肿瘤和健康组织、马和自行车),这使得每个类别都有其自己的任务相关激活区域。此外,类内区域是同一类别内的元素,它们经常共现(例如,肿瘤和水肿、马和骑手),导致出现虚假的任务无关激活区域 [25]–[27]。与自然图像相比,医学图像中受共现影响更为严重。对于小分割目标(如肿瘤),这种现象变得更加明显,现有的 WSS 方法往往难以有效处理这些小目标。如图 1 所示,病变的外周区域与目标区域强烈共现,导致激活区域不准确甚至完全错误。WSS 在医学成像中的这一固有挑战也在使用基于 SAM 的模型时遇到。因此,我们想知道如何减轻由共现引起的误激活的类间区域?b)CAM 主要用于识别图像中最具区分性的区域,可能会导致欠分割或过分割 [13]。因此,CAM 的细化需要实施补充策略。然而,天下没有免费的午餐。例如,基于亲和力的方法 [28] 通过训练辅助网络来细化 CAM,从而引入额外的计算开销。相反,仅依赖于像素的方法,如条件随机场(CRF)[29]、[30],无法整合结构信息。因此,是否有一种方法不仅可以利用 SAM 的现有参数,还可以利用医学结构信息来细化 CAM,而无需额外训练?为了抵抗上述挑战,本文提出了一个新颖的弱监督医学图像分割框架,即 WeakMedSAM。具体而言,为解决医学图像中的类内共现问题,一种可行的方法是将每个主要类别(例如,有肿瘤和无肿瘤)划分为几个子类(例如,肿瘤类型的潜在变化)。在训练之前,我们对属于同一主要类别的样本特征进行预聚类,以获取子类标签。随后,我们将子类分类任务与主要类别的分类任务结合起来。通过子类分类头显式学习类内表示,使得主要分类头能够实现准确的类间激活区域。在我们的实验中,我们惊讶地发现,仅优化主要类别分类头而不优化子类分类头,有时会意外地同时优化子类分类损失和 CAM 质量。这突显了类内表示在增强 CAM 质量中的重要性,并展示了子类机制显式获取和增强该表示的能力,从而提升类激活的精度。我们称此模块为子类探索(SCE)。为了充分利用 SAM 的参数和提示能力,通过网格点提示为图像样本获取亲和力图,随后从亲和力图中得出的概率矩阵决定了激活区域是否在特定像素处扩展或收缩,通过在 CAM 上执行随机游走来获得最终的伪标签。我们提出的用于细化 CAM 的方法消除了训练辅助网络的必要,从而显著减少了计算支出。我们称此模块为提示亲和力挖掘(PAM)。我们的 WeakMedSAM 被设计为一个即插即用模块,可以应用于任何基于 SAM 的模型。为了确保泛化能力,我们在 SAMUS [31] 和 EfficientSAM [32] 上进行了实验。我们有意避免使用在大型医学数据集上经过广泛训练的 SAM 模型,如 MedSAM [8],以防止由于包含我们的实验数据集在其训练数据中而产生的潜在偏差。此外,我们通过使用不同的预训练特征提取器(包括在 ImageNet 和 MIMIC-CXR 数据集上预训练的模型)来评估我们的方法的鲁棒性。通过在三个基准数据集——BraTS 2019、AbdomenCT-1K 和 MSD Cardiac 数据集——上进行广泛评估,我们的方法展示了很有希望的结果。我们的方法在 BraTS 2019 数据集上分别使用 SAMUS 和 EfficientSAM 实现了 79.69%/77.25% 的 Dice 和 5.57/10.35 体素的 ASSD,与最新的 WSS 方法相比,Dice 为 74.61,ASSD 为 11.91 体素。此外,我们通过修改 BraTS 数据集(称为 BraTS-Core)来验证我们方法在小目标分割上的有效性,并与全监督方法进行比较,以突出我们方法的实际潜力。总之,我们的主要贡献总结如下:

  • 我们首次尝试研究基于 SAM 的弱监督医学图像分割模型,以同时减少标注成本并借鉴 SAM 的出色能力。

  • 我们引入了一个子类探索模块,有效地缓解了医学图像中具有挑战性的共现问题,并因此产生了更精确的类激活区域。

  • 我们提出了一个提示亲和力挖掘模块,利用 SAM 的现有提示能力整合结构信息来细化 CAM,而无需额外训练。

II. 相关工作

A. 标签高效图像分割

在标签高效图像分割领域,半监督和少样本学习方法因其能够在有限的标记数据上实现鲁棒分割而受到广泛关注。半监督分割方法 [33]–[36] 利用标记和未标记数据来提升模型性能。另一方面,少样本分割 [37]–[39] 通过采用元学习框架来解决从最少标记示例中学习的挑战,这些框架从支持图像中提取类别特定的原型,以指导对查询图像的分割。尽管半监督和少样本方法在减少注释需求方面取得了显著进展,但我们的工作将重点转移到一个更高效的标签范式上。我们的方法采用弱监督分割框架,仅依赖于图像级类别标签,进一步减轻了数据注释的负担。

B. 弱监督图像分割

弱监督分割(WSS)[14]–[16] 是一种仅使用图像级标签进行图像分割的范式。基本上,初始伪标签通常是使用 CAM 生成的。然而,CAM 的一个常见缺点是它倾向于仅激活最具区分性的区域。为了克服这一限制,近期研究提出了各种训练策略。例如,一些方法整合了擦除 [40]、在线注意力积累 [41] 和跨图像语义挖掘 [42] 等技术来增强分割过程。其他方法利用辅助任务来正则化训练目标,包括视觉词汇学习 [43] 和尺度不变性正则化 [44]。此外,某些技术 [45]、[46] 利用额外的显著性图作为监督,以有效抑制背景区域并识别非显著性对象。此外,某些方法 [47]–[49] 对比像素和原型表示,以鼓励更全面地激活目标区域。与这些方法不同,我们提出了一个基于 SAM 提示能力的弱监督分割方法,整合医学图像中的结构信息而无需额外训练。此外,我们的方法在分割小目标(如小肿瘤)方面显著优于上述 WSS 方法,证明了在处理医学图像中的细粒度结构方面的优越性能。

C. SAM 调整用于医学图像

SAM 在应用于自然图像时表现出色。然而,当处理具有复杂形状、模糊边界、小尺寸或低对比度的对象时,它在某些医学图像分割任务中面临挑战 [5]、[6]。为了使 SAM 有效地适应医学图像领域,已提出各种方法 [7]、[8]、[31] 使用下游医学数据集对 SAM 进行微调。Ma 等人利用超过一百万张图像开发了 MedSAM [8]。此外,Zhang 等人引入了 LoRA 到 SAM 中,形成了 SAMed [7]。除了上述微调方法外,最近还出现了更多创新的方法,利用 SAM 进行医学图像分割。例如,医学 SAM 适配器(Med-SA)[10] 采用一种轻量级但有效的适应技术,将特定领域的医学知识整合到分割模型中。此外,One-Prompt SAM [11] 巧妙地处理了在推理阶段未见的任务,仅用一个提示即可在一次前向传递中完成处理。此外,SegAnyPath [12] 引入了一个多尺度代理任务和一个创新的任务引导的混合专家架构,展示了在病理图像分割中的卓越性能。然而,值得注意的是,这些现有方法依赖于全监督标签进行 SAM 的微调,需要大量医学图像进行精确标注。相比之下,我们的方法仅使用图像级类别标签,显著减少了数据标注成本,这对于在临床场景中准确且易于部署的医学图像分割至关重要。

III. 方法论

如图 2 所示,我们提出的 WeakMedSAM 包括两个模块:a)SCE(子类探索模块):使用图像级标签对 SAM 的 ViT 编码器进行弱监督微调,同时利用子类特征获取 CAM。b)PAM(提示亲和力挖掘):使用提示亲和力图对 CAM 进行随机游走细化,以增强医学图像中的结构表示。WeakMedSAM 的整体架构如图 2 所示,具体细节如下。

图 3:子类和主类分类的激活区域。同一主类中的不同子类会触发不同的类内判别区域。通过利用子类分类头提取额外的类内信息,主类分类头能够获得更具鲁棒性的类间激活表示。 

A. SAM 微调结构

我们提出的 WeakMedSAM 方法与任何基于 SAM 的模型兼容。在本研究中,我们使用了 SAMUS [31] 和 EfficientSAM [32],其中 SAMUS 代表传统的 SAM 微调框架,而 EfficientSAM 代表经过知识蒸馏后的参数高效 SAM 微调框架。1)SAMUS:SAMUS 的实现是在 SAM 的 ViT 编码器中加入一个并行的 CNN 分支,利用跨分支注意力来增强医学图像的分割。随后,开发了一个位置适配器和一个特征适配器,以促进 SAM 从自然领域到医学领域的适应。CNN 分支由一系列相互连接的卷积池化块组成,按顺序排列。跨分支注意力模块作为 CNN 和 ViT 分支之间的连接。所有特征适配器共享一致的结构,包括三个主要组成部分:1)向下的线性投影,2)激活函数,3)向上的线性投影。SAMUS 有效地减少了与微调相关的计算开销,同时不改变 SAM 本身的参数,从而代表了传统的微调框架。

图 4:未引入 SCE 的训练过程。仅专注于优化主类分类损失,可能会导致改进的子类特征表示和更准确的类激活区域。

2)EfficientSAM:EfficientSAM 提出了一种新颖的 SAM 预训练方法,称为 SAM 驱动的掩码图像预训练。这种方法通过将已建立的 MAE [50] 预训练技术与 SAM 模型相结合,为分割任务生成轻量级的 ViT 主干,从而促进高质量预训练 ViT 编码器的开发。具体而言,EfficientSAM 利用 SAM 编码器生成特征嵌入,并使用轻量级编码器训练掩码图像模型,以重建来自 SAM 的特征,而不是传统的图像块。EfficientSAM 重新设计了 SAM 的 ViT 编码器,显著减少了参数数量,为通过知识蒸馏对 SAM 进行微调提供了策略。WeakMedSAM 被设计为一个即插即用模块,可以应用于任何基于 SAM 类架构的分割网络,其性能和复杂性受到骨干 SAM 网络的影响。为了验证 WeakMedSAM 的多功能性,我们在上述两种 SAM 类骨干网络上进行了测试,证明了它能够在广泛的 SAM 类骨干网络上提供卓越的性能。后续的方法描述与 SAM 骨干网络无关。

B. 子类探索

我们不使用像素级全监督标签,而是仅使用图像级分类标签来微调 SAM。为了缓解医学图像中 WSS 的共现现象并强调类间知识,我们引入了一个子类分类任务,该任务由通过聚类获得的子类标签进行监督。这个任务显式地学习不期望的类内表示,同时允许主要分类任务激活任务相关区域。这种方法在医学成像中尤其重要,因为类内共现现象很常见,并且它进一步增强了模型准确激活小目标(如肿瘤)的能力。SCE 模块如图 3 所示。1)获取子类标签:对于每个主要类别,记为 pc,其中 c ∈ {1, . . . , C},我们指定 K 个子类,表示为 sk c,其中 k ∈ {1, . . . , K}。对于每个图像样本 x,其主要类别标签为 yc p ∈ {0, 1}C,相应的主要类别 c 的子类标签是一个长度为 K 的 0-1 向量,表示为 yc s ∈ {0, 1}K。由于没有子类的真值标签可供直接优化,我们通过聚类生成伪子类标签。具体而言,我们对每个主要类别进行聚类,基于从特征提取器中提取的图像特征。聚类算法可以是任何无监督聚类方法,例如 K-means,特征提取器可以是任何预训练的图像编码器。在第 IV-C 节中,我们进一步分析了这一过程,特别是针对聚类算法的随机性和特征提取器的结构及预训练数据集的影响,表明我们的方法对这些因素具有很强的鲁棒性。2)联合训练:我们的最终目标是开发一个子类分类头 Hs,参数化为 θs,同时共享 SAM 的图像编码器 E 与主要分类头 Hp,参数化为 θp。对于 Hp 和 Hs,我们使用多标签二进制交叉熵损失作为分类损失 L。一旦我们通过上述聚类过程获得了子类的伪标签 ys,我们就会联合优化这两个分类器 Hp 和 Hs:

其中 E 表示 SAM 的图像编码器,N 表示图像总数,λ 表示用于平衡两个损失函数的权重,我们将其设置为 0.5。通过这种方法,子分类头显式地探索类内子空间,从而使主要分类头专注于学习干净的类间表示,以便获得更准确的 CAM。3)分析:为了证明 SCE 模块的有效性,我们设计了以下实验:在不优化子类分类头的参数的情况下,仅优化 SAM 图像编码器和主要类别分类头。我们在训练过程中跟踪主要类别分类损失 Lp、子类分类损失 Ls 和 CAM 的 Dice 系数 Sdsc。我们得到了一个有趣的实验结果,展示了经过细化的类内特征表示与更准确的类激活区域之间的相关性。如图 4 所示,在训练初期(图 4 的部分 (a)),Lp 迅速收敛,但 Sdsc 达到一个平台期,而 Ls 保持相对不变。这种优化过程是直观的,并且自然地符合没有 SCE 模块的模型的预期行为。

接下来(图 4 的部分 (b)),随着我们持续仅优化 Lp,Sdsc 显著提高,伴随着 Ls 的降低。我们将此归因于经过细化的类内子空间特征表示,这得到了 Ls 降低的支持。这表明,在减轻了任务无关的类内表示对图像编码器的影响后,依赖于类间信息的 CAM 的准确性得到了提高。最后(图 4 的部分 (c)),在没有显式学习类内表示的情况下,Ls 的收敛停止了。持续优化 Lp 导致过拟合,表现为虚假激活区域的出现,导致 Sdsc 下降。这进一步强调了消除类内干扰以增强 CAM 的重要性。值得注意的是,并非所有训练过程都遵循这一范式,仅优化 Lp 并不一定会导致进一步探索类内空间。因此,引入 SCE 模块以显式优化 Ls。SCE 有助于模型获得类内空间的细化特征表示,从而增强激活区域的准确性。

C. 提示亲和力挖掘

借助像 SAM 这样的大型模型的提示能力,我们的目标是在不需要额外训练的情况下实现图像上相邻坐标之间的类别无关亲和力。这些亲和力随后被用作随机游走过程中的转移概率,允许 CAM 在相同结构实体的邻近区域内传播。这种传播显著提升了 CAM 的质量。1)从 SAM 获取亲和力图:两个坐标之间的亲和力代表它们的类别无关结构特征之间的相似性,而特定显著区域内的亲和力对于指导细化更有意义。这两个属性——结构相似性和空间局部性——与 SAM 的提示机制一致,其中点提示会导致具有相似结构特征的局部激活区域。如图 5 所示,为了利用 SAM 的提示能力获取亲和力图,首先,图像被一个大小为 8×8 的均匀网格 G 划分。然后,对于每个网格区域 i,在网格区域的中心点 pi 应用一个点提示,这会产生一个掩码预测图 maff i ∈ RH×W,即掩码解码器的输出。因此,图像的整体亲和力图可以定义如下:

其中 norm(·) 对图 maff i 的亲和力值进行归一化。然后,全局亲和力图 Maff 将用于生成后续随机游走的转移概率矩阵 T。2)使用亲和力图修订 CAM:从 SAM 中获得的局部结构亲和力被转换为转移概率矩阵 T。该矩阵促进了对图像内结构区域敏感的随机游走,促进了这些区域内的激活分数的传播,如图 6 所示。一对特征向量之间的结构亲和力由它们的 L1 距离表征。具体而言,我们用 aij 表示特征 i 和 j 之间的结构亲和力,计算如下:

请注意,亲和力仅在半径为 γ 的局部圆内的特征之间计算。这里我们取 γ = 5。计算出的亲和力共同形成一个亲和力矩阵 A,其中对角线元素设置为 1。从这个亲和力矩阵中,转移概率矩阵 T 用于随机游走过程,按照以下程序导出:

其中超参数 β 被赋予一个大于 1 的值。通过将原始亲和力矩阵 A 乘方为 β,记作 A◦β,我们有效地抑制了 A 中不重要的亲和力。这一调整使得随机游走过程更加谨慎地传播。此外,对角矩阵 D 的计算有助于对 A◦β 进行逐行归一化。通过使用转移概率矩阵 T,亲和力传播过程通过一系列随机游走操作实现。具体而言,CAMs 乘以 T 来执行这种传播。这种迭代传播过程持续进行,直到达到预定义的迭代次数 t。因此,类别 c 的修订 CAM,记作 M ∗ c,使用以下表达式获得:

其中操作 vec(·) 表示矩阵的向量化,t 表示迭代次数。

D. 复杂性讨论

WeakMedSAM 是任何基于 SAM 模型的扩展。与原始 SAM 模型相比,WeakMedSAM 额外集成了两个线性层,专门用于子类探索(SCE)模块——主要类别分类头和子类别分类头,增加了极小的计算开销。值得注意的是,提示亲和力挖掘(PAM)模块没有增加任何额外的模型参数。该模块的计算开销仅包括提示编码器和掩码解码器,这些是轻量级的,甚至可以在网络浏览器中执行 [1]。这种设计确保了我们的方法适用于训练和推理过程,能够在各种 GPU 配置上执行,而无需大量的计算资源。在第 IV-C7 节中,我们进一步对所提出方法的复杂性进行了数值讨论,详细说明了其在不同计算环境中的效率和可扩展性。

IV. 实验

A. 实验设置

1)数据集:我们使用了三个广泛使用的数据集进行研究:BraTS 2019 [55]、AbdomenCT-1K [56] 和 MSD Cardiac 数据集 [57]。BraTS 2019 数据集包含 335 个多模态扫描,每个扫描都配有专家分割掩膜。这些扫描包括四种模态:T1、T1c、T2 和 FLAIR,而在这里我们仅使用 FLAIR 模态。该数据集特别用于二元分割任务,我们的重点是区分健康和不健康的目标。AbdomenCT-1K 数据集是一个全面且多样化的腹部 CT 器官分割数据集,包含来自 12 个医疗中心的 1000 多个 CT 扫描,涵盖了多阶段、多供应商和多疾病病例。我们的目标是分割其中的四个腹部器官:肝脏、肾脏、脾脏和胰腺。MSD Cardiac 数据集包含用于左心室分割任务的 MRI 图像。它包含 30 个患者病例,每个病例包含多个 2D MRI 切片。我们按照 8:1:1 的比例随机将其划分为训练、验证和测试集,基于患者进行划分。在我们的分析中,我们将 3D-MRI 扫描的各个切片视为 2D 图像,专注于切片级别的分割任务。为了进一步验证我们的方法在小目标(如小肿瘤)分割上的性能,我们在 BraTS 数据集中去除了肿瘤周围的水肿,仅保留肿瘤区域。我们称这个数据集为 BraTS-Core,它遵循与原始 BraTS 数据集相同的划分方案和标签定义。关于 BraTS-Core 数据集的详细描述在第 IV-C6 节中提供。

2)网络架构:我们保持 SAM 的提示编码器和掩码解码器处于冻结状态,因此仅对 SAM 的图像编码器进行训练,使用 SAMUS 和 EfficientSAM。对于 EfficientSAM,我们使用参数最少的 ViT-Ti 编码器作为骨干网络。值得注意的是,我们提出的 WeakMedSAM 可以即插即用于任何基于 SAM 的模型。为了不失一般性,后续的分析实验在 SAMUS 中进行。与使用编码器最终颈部层的输出不同,我们选择使用来自 Transformer 块的输出作为图像嵌入进行分类,以获得更好的类激活表示。此外,主要类别和子类别的分类头均采用 1×1 卷积层。对于使用伪标签进行训练的分割网络,我们使用广泛应用于医学图像分割任务的 U-Net [58] 网络。

3)实现细节:我们使用 Python 和 PyTorch 框架进行实现。代码在四块 NVIDIA GTX 2080Ti GPU 上运行。学习率遵循一个周期策略,最高达到 10^-4,然后在后续迭代中衰减。批量大小设置为 24,总训练周期数为 10。图像被裁剪为 256×256 的大小。此外,我们为子类别损失使用了常量权重 λ = 0.5。

4)评估指标:为了评估这些伪标签的质量,我们使用了以下四个评估指标:Dice 系数(DSC)、Jaccard 指数、平均对称表面距离(ASSD)和 Hausdorff 距离的 95% 分位数(HD95)。DSC 和 Jaccard 指数是广泛使用的像素级相似性度量,用于评估分割模型的性能。另一方面,ASSD 和 HD95 通过测量预测结果与真实标签之间的表面距离来量化分割结果的准确性。

B. 与其他方法的比较

  1. 与弱监督分割方法的比较:为了评估我们提出的方法 WeakMedSAM 的有效性,我们对类别激活图(CAMs)应用了阈值,以生成伪标签,随后利用这些伪标签训练分割网络并获得最终的分割结果。我们将该方法与近期的弱监督分割方法进行了比较,结果如表 I 和图 7 所示。我们的方法在多个数据集上展现出了卓越的准确性,并在两种不同的基于 SAM 的网络上表现出色,从而验证了我们方法的通用性。

  2. 与交互式 SAM 类方法的比较:尽管 SAM 并未专门在医学图像数据集上进行训练,但它展现出了零样本泛化的潜力 [59]。特别是,SAM 的交互式分割功能,通过点和边界框等提示引导,可以极大地提升终端用户的体验。利用视觉基础模型进行交互式分割,在减少繁重的像素级标注成本方面具有巨大的潜力。因此,我们提出了一个假设性的交互式分割场景,其中专家使用视觉提示来指定感兴趣的区域。为了使这种交互式分割的成本与我们的图像级标签获取成本相当,我们采用了一个单点提示。这个提示的坐标对应于 SAM 交互式分割中使用的分割标签的中心点。

在最近的研究中,许多工作都集中在对大规模医学图像数据集进行 SAM 的微调上。在实际临床应用中,使用这些微调后的模型进行交互式分割可能更具重要意义。因此,我们在与 SAM 的比较中也加入了 MedSAM。研究结果表明,SAM 在某些样本中表现出强大的零样本泛化能力,如图 8 所示。然而,在大多数情况下,由于缺乏与任务相关的上下文信息,SAM 的分割结果主要依赖于像素信息,而医学图像中这种信息严重不足。因此,SAM 往往会将空白区域解释为背景,而不是非目标区域,从而导致严重的过分割或欠分割。因此,直接将 SAM 应用于医学图像分割并非最佳选择。相比之下,MedSAM 相比 SAM 在整体分割性能上有了显著提升。尽管它略逊于我们的 WeakMedSAM,但通过增加少量额外的交互式提示,它可以取得出色的结果。

重要的是,使用弱监督框架中的类别级标签,消除了在推理阶段需要临床医生进行交互的需求。WeakMedSAM 的整个推理过程是完全自动化的,进一步减少了分割任务所需的时间。

C. 消融研究与分析

消融分析的定量结果如表 III 所示。该表表明,基于 ViT-b SAM 的骨干网络在 BraTS 2019 数据集上实现了 68.13% 的 DSC(Dice 系数)。引入所提出的子类分类模块和提示亲和力挖掘模块后,DSC 分别显著提升至 74.91% 和 76.31%。当这两个模块结合时,模型的性能提升至 79.69%。

  1. SCE 特征提取的不同骨干网络:由于缺乏子类的真值标签,我们通过无监督聚类生成伪标签。具体来说,基于特征提取器提取的图像特征对每个主类进行聚类。一个潜在的问题是,执行特征提取的骨干网络的性能可能最终会影响子类分类的最终性能。我们使用不同的预训练网络进行了预聚类。如表 IV 所示,我们的网络对骨干网络的性能并不敏感。因此,为了减少计算开销,我们在预聚类阶段使用了 ResNet18 作为特征提取器。

  2. SCE 特征提取器的不同预训练数据集:尽管我们的实验使用了 ImageNet 预训练模型用于 SCE 模块,但我们认识到考虑医学图像预训练特征提取器的重要性。然而,为各种医学图像模态识别或训练模型仍然是一个挑战。现有的文献表明,即使在存在显著领域偏移的情况下,ImageNet 预训练的特征提取器在下游任务中仍然有效 [60]–[63],在医学成像中也是如此 [64]–[66]。鉴于我们在第 IV-C1 节中证明的 SCE 模块对特征提取器容量不敏感,使用 ImageNet 预训练的特征提取器是一个合理且实用的选择。为了进一步分析这一点,我们使用在 ImageNet 和 MIMIC-CXR [67] 数据集上预训练的 ViT 模型评估了 SCE。图 9 展示了使用 ImageNet 和 MIMIC-CXR 预训练的 ViT 对 BraTS 数据集进行聚类的结果,表 V 展示了使用这些提取器的分割性能。这些结果表明,SCE 对特征提取器的预训练数据集具有鲁棒性。

  3. 不同子类数量:引入子类分类任务的目的是明确地定义类内信息,这可能会扭曲激活区域。通过子类分类头进行学习,从而使得主类分类头能够专注于获取类间信息。然而,每个主类关联的子类数量是在预聚类过程中手动确定的。这意味着不同数量的子类可能会影响模型的性能。通过实验不同数量的子类,如图 10 所示,我们观察到当子类数量为 8 时,模型性能达到最佳。作为一个未来的研究方向,开发一种自适应的方法来确定子类的数量将是有益的。这可以减少多余子类的出现,从而提高方法的效率。

  4. 聚类算法随机性的分析:为了应对聚类算法固有的随机性对 SCE 模块性能可能产生的影响,我们使用多个随机种子进行了聚类实验。结果如图 11 所示,对于给定的一组特征,聚类算法对随机种子初始化的变化表现出鲁棒性。并且分割结果也保持了高度稳定,所有指标的变化都小于 2%。因此,我们的 SCE 模块对聚类过程引入的随机性不敏感。

  5. 与其他细化方法的比较:为了增强 CAM,人们已经探索了各种方法。这些方法可以大致分为两类:一类是需要额外训练以纳入结构信息的方法,例如像素级语义亲和力(PSA)[28];另一类是不需要额外训练的方法,如条件随机场(CRF)[29]、[30],但它们未能利用结构信息。相比之下,我们提出的 PAM 模块旨在提取结构信息,同时不需要任何额外的训练。为了证明 PAM 模块的有效性,我们对其进行了与其他方法的比较分析。如表 VI 所示,我们的结果表明,我们提出的方法优于其他方法。

  6. 小目标分割的分析:我们在多个数据集上评估了我们的方法的性能。如图 12 所示,BraTS 数据集的目标平均相对大小较小。我们在 BraTS 数据集上进行了详细的实验,包括消融研究和超参数敏感性分析,证明了我们的方法在小目标(如小肿瘤)上的有效性。为了进一步验证小目标上的性能,如图 14 所示,我们修改了 BraTS 注释,仅保留非增强肿瘤核心(NET)和 GD 增强肿瘤(ET),排除了肿瘤周围水肿(ED),创建了 BraTS-Core 数据集。如图 13 所示,BraTS-Core 的目标平均相对大小比 BraTS 更小。我们在表 VII 中将 BraTS-Core 上的分割性能与近期的弱监督分割方法和全监督的上限进行了比较。该上限定义为 BraTS 2019 排行榜 [1] 中肿瘤核心(Dice_TC)的 DSC 指标的平均值。我们的方法在小目标上表现出显著的优势。然而,与全监督分割相比,仍存在性能差距,这表明未来还有改进的空间。

  7. 进一步的计算复杂性数值分析:为了进一步分析我们方法的计算复杂性,我们重申,我们提出的模块引入的开销是最小的。SCE 模块的分类头仅由几个 MLP 组成,而 PAM 模块利用了 SAM 的提示编码器和掩码解码器,这些组件被设计为轻量级操作,甚至可以在网页浏览器中运行。主要的计算成本来自于 SAM 的图像编码器,在我们的训练过程中,该编码器大部分是冻结的。重要的是,我们的方法可以作为即插即用的增强模块应用于任何基于 SAM 的架构。为了进行定量评估,表 VIII 提供了常用 SAM 类模型的参数计数。

  8. 与全监督上限的比较:我们在 BraTS 数据集上比较了我们的弱监督方法与全监督上限的性能。与表 VII 一致,我们使用了该上限的 DSC 和 HD95 指标的平均值。如表 IX 所示,尽管全监督方法天生具有更优越的分割性能,但我们的方法显著减少了对大量数据标注的需求,同时并未导致性能完全失败,展示了其在实际临床应用中的潜力。

D. 参数敏感性

为了专注于核心方法而不是广泛的训练优化技巧,我们在 BraTS 数据集上进行了超参数消融研究和敏感性分析。从这些分析中得出的超参数集随后被应用于 AbdomenCT-1K 和 MSD Cardiac 数据集,以评估所选设置的泛化能力。

利用亲和图来获取随机游走的概率矩阵 T 的方法涉及超参数。在公式(4)中,超参数 β 大于 1,用于抑制矩阵 A 中的不显著亲和力。在公式(5)中,超参数 t 决定了随机游走的迭代次数。我们系统地进行了大量的实验来评估这些超参数的影响。我们的实证研究表明,当随机游走相关性的阈值 β 设置为 4,且 t 也设置为 4 时,亲和力的细化效果最佳。

作为一种启发式研究,以确定其他数据集的最优参数,我们对 AbdomenCT-1K 数据集进行了额外的超参数分析,重点关注 SCE 模块中的子类数量 K 和 PAM 模块中的随机游走迭代次数 t。作为一个多器官分割数据集,AbdomenCT-1K 拥有更多的主类,这降低了类内共现的重要性,与只有单一主类的 BraTS 相比。因此,如图 15 所示,较小的 K 是最优的,K = 4 时性能最佳,而 BraTS 的最优 K = 8。t 的最优值在各个数据集中保持一致,t = 4 时表现最佳。

E. 交互式 WeakMedSAM

与利用大量像素级标签进行密集训练的全监督方法(如 MedSAM)不同,WeakMedSAM 依赖于图像级监督,这本质上限制了其交互式分割性能。我们的目标是证明 WeakMedSAM 的弱监督微调并不会破坏 SAM 架构中图像编码器与其他组件之间的一致性。为了验证 WeakMedSAM 在交互式分割场景中的适用性,我们设计了两种类型的提示:一种是从分割标签的质心派生的单点提示,另一种是覆盖整个分割标签的边界框提示。与原始 SAM 的对比实验如图 16 所示,结果表明,SAM 和 WeakMedSAM 在使用边界框提示时表现良好。然而,在使用单点提示时,WeakMedSAM 显著优于 SAM,显示出其在利用有限交互信息方面的卓越能力。这些发现表明,仅使用图像级分类标签对 SAM 进行微调以开发适用于交互式分割的模型是可行的;然而,我们承认 WeakMedSAM 并非专门为这一应用而设计,这表明需要进一步的研究来提升其在此类场景中的性能。

V. 结论

在本研究中,我们探讨了在 Segment Anything Model(SAM)的指导下,弱监督医学图像分割的范式。我们提出的模型 WeakMedSAM 包含两个模块:1)第一个模块,即 SCE,涉及对 SAM 的 ViT 编码器进行弱监督微调,以利用子类特征,有效消除潜在的共现问题,以获得可靠的 CAM 结果;2)第二个模块,即 PAM,执行 CAM 的随机游走细化,利用新出现的亲和图来增强类内表示。从三个广泛使用的基准数据集中得出的实验结果,证实了本文提出的 WeakMedSAM 模型的前景广阔。在更广泛的背景下,我们的方法还为适应 SAM 提供了一种新的视角,通过减少对大量精确标注数据的依赖,将其应用于各种下游领域,这也可以被采用在其他场景中。未来的研究将探索使用弱监督方法对 SAM 进行端到端微调的潜力,这可能会进一步提升性能和泛化能力,这将专注于改善 SCE 模块中的 ViT 编码器微调与 PAM 模块中的基于亲和力的细化之间的协同作用,确保 SAM 更加协调和高效地适应交互式医学图像分割。

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

Together_CZ

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值