【SAM医学分割】半监督三维医学图像分割的 SAM 框架
Segment Anything Model（SAM）微调技术在完全监督的医学影像分割中表现出卓越的性能，但需要精确的注释。为了降低注释成本并保持令人满意的性能，这项工作利用 SAM 的功能建立了半监督医学图像分割模型。重新思考了有效性、效率和兼容性的要求，提出了一个三阶段框架，即缝合、微调和再训练（SFR）。目前的微调方法大多涉及二维切片微调，忽略了相邻切片之间的上下文信息。本文的拼接策略可以减少自然图像和三维医学图像之间的不匹配。拼接后的图像用于微调 SAM，为伪标签提供稳健的初始化。之后，训练一个三维半监督分割模型，同时保持与 V-Net 等传统分割器相同的参数大小。SFR 框架即插即用，可轻松兼容各种流行的半监督方法。同时还开发了一个扩展框架 SFR+，通过置信度估计进行选择性微调和再训练。广泛的实验验证了SFR 和 SFR+ 在五个数据集的中度注释和稀缺注释中都取得了显著的改进。特别是，SFR 框架仅用 LA 数据集的一个标注数据就将平均教师的 Dice 分数从 29.68% 提高到 74.40%。
INTRODUCTION
最近，视觉分割的一般基础模型因其出色的分割和泛化能力在医学图像领域引起了广泛关注。尽管这些基础模型在医学图像分析领域取得了显著的进步，但由于特定的模式、复杂的成像技术和多变的组织等不可避免的因素，利用统一的模型分割所有医学图像有时仍具有挑战性。为了解决这个问题，最近有几项研究提出，要么专注于提示工程，要么设计适配器进行微调]，以借用基础模型的能力，例如从 SAM 衍生出的 MSA和 SAMed ，来完成特定的任务。
这些大多是采用微调或适应技术的全监督方法。然而，全监督医学影像分割依赖于经验丰富的专家划定的大量精确注释，这使得标注过程变得繁琐、耗时，甚至主观。
最近的趋势表明，在某些情况下，半监督方法的性能几乎可以与完全监督方法相媲美。例如，在注释率为 40% 的 BTCV 数据集上，[16] 的性能优于完全监督方法。而 [17] 在 LA 数据集上的性能在有 20% 标记数据的情况下，仅比完全监督结果低 0.8%。因此，不禁要问，目前基础模型的成功是否能促使我们开发出一种有效的半监督医学图像分割模型？
基于上述目标，本文希望在设计半监督医学图像分割框架之前，重新审视几个重要因素。
-如何有效初始化？ 根据以往的研究，在半监督情况下，初始化阶段伪标签的质量对接下来的分割起着重要作用。与自然图像不同，三维医学图像的切片间连续性对于准确的目标分割至关重要。此外，医学图像的分辨率通常相对较低。现有的医学影像拟合策略，包括直接放大二维切片和调整位置嵌入的大小，都是采用逐片微调的方式，而忽略了三维图像中存在的固有切片间相关性。因此，是否有更好的方法利用基础模型提高医学图像初始伪标签的质量 ？
-如何提高效率？ 基础模型是在参数相对较大的大规模数据集上预训练的。现有的微调方法在推理过程中仍然保留了基础模型的原始参数大小，甚至引入了额外的参数。在分割医学图像时，真的需要这么大的参数尺寸吗？首先，现有的只涉及小尺寸参数的模型在分割器官时确实表现良好。其次，注意到最近的研究 揭示了基础模型中的冗余问题，揭示了大规模预训练模型被过度参数化 。第三，与自然图像相比，医学图像的外观往往具有标准化的视图和相对有限的纹理变化 。针对这些问题，是否有可能在保持良好结果的同时摆脱过度参数化的基础模型 ？
-如何保持兼容性？ 一方面，近年来半监督学习作为一种有吸引力的策略出现，并被广泛应用于医学影像分割任务，同时也提出了很多半监督医学影像分割方法。能否让基础模型更好地服务于现有的半监督方法呢 ？另一方面，计算机视觉和机器学习领域关于半监督学习的研究进展仍然有助于发展新的半监督医学影像分割模型。框架将来还能兼容这些新方法吗 ？
有鉴于此，本文认为，在基础模型时代，一个有前途的半监督医学影像分割框架应该是性能有效、参数高效和兼容性出色的。因此，本文提出了一个简单明了的缝合、微调和再训练（SFR）框架来实现上述目标。首先开发了一种拼接策略，对切片进行拼接操作，生成与高分辨率输入相匹配的图像，从而更好地利用切片间的关系和维度信息。然后将拼接后的图像输入 SAM 进行微调。 之后，在 SAM 的指导下训练一个小规模的三维分割模型，同时保持参数大小不变。经过微调的 SAM 提供了有利的初始化，并与各种三维模型兼容。 此外，还开发了扩展框架 SFR+，引入了置信度估计和选择性训练策略，以提高对无标记数据的利用率。
**** -提出了一种新颖的框架，既能利用基础模型的能力，又能确保半监督分割的性能，并进一步降低标注成本，该框架包括三个阶段，即拼接、微调和再训练。
-拼接策略在伪标记初始化方面简单而有效，与当前的调整大小/直接微调策略有很大区别。
-在推理过程中，参数大小与主流分割器（如 V-Net ）保持一致，大大小于基础模型的参数大小。
-框架即插即用，可以很容易地与现有的大多数流行的半监督分割方法相结合。
RELATED WORK
Foundation Models in Medical Images
- 视觉基础模型：视觉基础模型已受到广泛关注，并在包括分割在内的各种计算机视觉任务中表现出令人印象深刻的性能。这些模型的突出例子包括 SAM 、SegGPT 、SEEM 、SLiMe 和 SAM 2 ，以及它们的扩展应用。这些模型利用大规模图像数据集来学习通用的视觉表征，并展示了卓越的泛化能力。
在医学图像领域，UniverSeg通过提供一组图像标签对示例，实现了二维医学图像的通用分割。STU-Net 是专门针对 CT 模式的基础模型，其最大的变体包含 14 亿个参数。此外，SAM 已成为最流行的图像分割模型之一，许多研究已将其扩展到医学图像领域。SAM-Med2D 是一种二维模型，在 460 万张医学图像上对 SAM 进行了微调。SAM-Med3D采用了类似于 SAM 的架构，但它是从头开始训练的，没有使用 SAM 的预训练权重。由于 SAM 的出色性能和广泛适用性，所提出的框架将其作为默认的基础模型。
- 使 SAM 适应 3D 医学影像 ：SAM 的零镜头能力不足以确保其在医学图像中的直接应用。为了将其强大的分割能力扩展到医学图像中，许多研究都致力于通过不同的图像处理和微调策略进行微调。
对于图像处理而言，三维医学图像与预训练的自然图像之间在图像分辨率上的差异是一个挑战，目前已提出两种策略来解决这一问题 。第一种是上采样微调 ，即对每个切片进行上采样以直接匹配输入分辨率。第二种是小尺寸微调 ，即通过双线性插值来缩小输入尺寸。然而，这两种策略都是基于二维输入，而对于三维医学图像，需要通过分割每个切片来生成预测结果。3DSAM-Adapter和 SAMMed3D将 SAM 扩展到了三维架构，但由于模型尺寸较大，它们增加了额外的训练开销。与此相反，本文的拼接策略旨在适应图像尺寸和分辨率的变化，创建大尺寸的拼接图像，有效捕捉相邻切片的空间信息。
微调方法包括只微调子部件参数和加入适配器。MedSAM 对 SAM 的掩码解码器进行了微调，并冻结了编码器，但其性能显示落后于特定医疗模型，特别是在边界区域方面。MSA 和 SAMed 采用了参数高效微调技术，使用适配器和基于低秩的策略（LoRA） 进行微调。
METHOD
Notations and Framework Overview
给定 m 幅已标注图像和 n 幅未标注图像，第 i 幅（1 ≤ i ≤ m）已标注图像及其地面实况分别记为 Xl i 和 Yl i，第 j 幅（1 ≤ j ≤ n）未标注图像记为 Xu j。这里，Xl i、Xu j∈RH×W×D，Yl i∈{0，1，...，K - 1}H×W×D，其中 H、W、D 表示三维医学图像的相应维度，训练的输入补丁大小为正方形，即 H = W。K 是要分割的不同类别的数量。
如图 2 所示，构建的 SFR 框架由以下三个模块组成：拼接模块、微调模块和再训练模块。拼接模块可减少自然图像与三维医学图像之间的不匹配 。拼接后的图像被输入到 SAM 中进行微调，并为半监督模块提供初始伪标签。然后，训练三维半监督分割模型。此外，还提出了 SFR+，通过置信度估计有选择地进行微调和再训练，从而增强了所提框架。
步骤 1：拼接模块。如前所述，通过减少预训练样本（即自然图像）和微调样本（即医学图像）之间的分辨率差异 ，拼接模块可以利用切片拼接函数 FC(-) 将三维标注体 Xl i 转换为大尺寸二维图像 Ml i∈ RHd×W d。此外，拼接后的地面实况 Nl i 也是类似得到的。Ml i 和 Nl i 按 d×d 网格排列，其中 d = ⌈√D⌉，如果 d × d > D，将在所有切片后拼接零点。
步骤 2：微调模块。首先利用拼接后的标注图像 Ml i 及其地面真实图像 Nl i，通过常用策略（如 LoRA ）对以 θ 为参数的 SAM 进行微调。这样做的目的是缩小自然图像和医学图像之间可能存在的分布偏移 。以 LoRA 为例，将微调函数表示为 FLoRA(-)，其参数为θ，输入为 Ml i 和 Nl i：
然后，利用微调 SAM 的预测函数 FFT(-)为未标记图像生成高质量的伪标签，并通过拼接反变换 F-1 C (-) 生成三维伪标签 ，如下所示：
步骤 3：重新训练 SSL 模块。最近 SSL 网络上流行自训练、平均教师等、本模块可采用最先进的 ACMT、MagicNet 方法。具体来说，SSL 网络从微调的 SAM 中学习伪标签 。将 SSL 网络记为 FS(-)，参数为 ω，最优 ω∗ 的再训练模块如下：
SFR Framework
1) 缝合模块 ：为了从二维自然图像适应三维医学图像，本文认识到输入分辨率和图像维度是关键因素 。三维体积的切片间空间信息与目标识别息息相关，而在高分辨率图像上训练的大型模型很难推广到低分辨率的医学图像切片上。受此启发，所提的拼接策略（如图 3 所示）将医学图像与自然图像分辨率相匹配，并补充了三维医学图像特有的空间排列。预训练 SAM 模型的输入空间分辨率为 1024 × 1024。拼接策略将三维体积（原始三维图像或三维补丁）逐片排列成 d × d 网格，生成大小为 1024 × 1024 的二维图像。由于不同医疗数据集的切片尺寸各不相同，因此切片分辨率和拼接切片数量之间存在性能权衡。所提方法可以有效管理不同切片数量的图像。对于小比例切片，如 LA 数据集，使用原始三维图像作为输入卷。对于大规模切片，如 BTCV数据集，采用常见的三维处理方式，将体积划分为多个补丁作为输入体积，然后拼接所有切片以实现 1024×1024 的最终尺寸，避免直接缩小大切片。与图 3 中的小尺寸输入微调方法和直接升采样微调方法相比，发现拼接策略有效地解决了图像尺寸和分辨率差异带来的挑战。
从切片的连续性和上下文的完整性出发，对拼接策略进行了深入研究。
切片连续性。 由于三维医学影像固有的空间连续性，本文探索了三维医学影像与切片连续性之间的关系。对于拼接的二维图像，分割模型通过自我关注机制学习各切片之间的特征相关性，因此它很难在没有切片顺序的情况下捕捉上下文信息和形状的连贯性。为了研究切片连续性的重要性，采用了两种方法来破坏连续性：1) 随机改变切片顺序；2) 随机旋转和翻转每个切片。同样的切片操作也适用于地面实况掩码，以保持输入数据和标签之间的一致性。如图 4 所示，观察到，打乱切片顺序会导致形状一致性的丧失和性能的下降，结果强调了切片连续性在三维器官分割中的重要性。
上下文完整性 。拼接模块将一个体素（三维原始图像或三维补丁）重组为 1024 × 1024 图像，从而在单张图像中完整呈现该体素。为了探索上下文完整性对拼接的影响，将医学切片与自然图像进行拼接，同时保持每个切片的分辨率不变。如图 5 所示，通过逐步增加医学切片的数量，从自然图像过渡到医学图像。具体来说，在训练过程的开始阶段，从 PASCAL VOC 2012 数据集中加入自然图像，该数据集共包含 2,913 幅图像。每次迭代，都会从该数据集中随机选择一批自然图像。随着训练的进行，逐步用医学图像切片替换部分自然图像，到最后阶段，只使用医学图像切片。测试时，始终使用完全拼接的医学图像切片。虽然这种方法看似逐渐从自然图像特征调整为医学图像特征，但它实际上破坏了医学图像中解剖结构的完整性，导致性能下降。
这些观察结果表明，拼接策略必须保持切片的连续性和上下文的完整性，从而有效弥合自然图像和医学图像之间在领域和空间维度上的差距。
2) 微调模块 ：微调模块对视觉基础模型进行微调。作为最流行的通用图像分割模型之一，SAM 是微调模块的默认设置，表示为 FFT(-)。剔除了所有提示，并在推理过程中执行自动分割。框架并不局限于特定的微调策略，它可以用于不同的策略中。将微调损失统一表示为：
其中，Pl i = FFT(Ml i) 是微调模块的预测值。以往的研究主要涉及仅对子部件参数进行微调以及加入适配器。
子部件微调 。子部件微调方法直接修改模型参数。MedSAM-v1冻结了图像编码器和提示编码器，只对掩码解码器进行了微调，而 MedSAM-v2 则对图像编码器和掩码解码器都进行了微调。然而，其整体性能仍然落后于医学图像分割的专家模型，尤其是在边界共识方面。
Adapter Tuning 。适配器调整是指在原始基本模型中插入适配器，只对适配器进行调整，而冻结所有预训练参数。适配器由向下投影层、ReLU 激活层和向上投影层组成。基于低秩的微调策略（LoRA）将可训练的低秩分解矩阵注入预训练模型的各层。SAMed冻结了 SAM 的图像编码器，通过添加旁路采用 LoRA，并对掩码解码器进行微调。
由于 LoRA 可以与原始预训练权重合并进行推理，因此本文采用它作为微调模块方法 FLoRA(-)。根据文献[10]，对于 SAM 的分类头，模糊性预测由每个语义类别的确定预测输出来代替。
LoRA 为跨域微调提供了一种实用而有效的方法，LoRA 与本文为医学图像量身定制的特定拼接策略相结合，确保模型能够很好地处理自然图像数据与医学图像数据之间的巨大差异。
本文注意到，拼接二维切片可能会导致器官出现在周围和相似的区域，因此需要引导 SAM 捕捉这些相似性 。拼接策略保留了切片之间的空间关系，可以有效利用相邻切片中同一目标的信息 。例如，如图 6 所示，同一语义类别（如脾脏）出现在三个相邻切片的相似区域（左上角）。逐片预测的上采样微调方法将整个脾脏（红色）混淆为中间切片中的左肾脏（蓝色）。相反，缝合微调方法可以在所有三张切片中正确识别脾脏。通过结合切片间信息，所提的方法能有效指导 SAM 在不同切片上识别同一器官。相比之下，拼接策略保留了切片之间的空间关系，可以充分利用相邻切片上同一目标的信息。
为了验证伪标签在不同数据集上的有效性，以 LA 和 BTCV 数据集为例，在图 7 中展示了微调模块对单个和多个目标数据集的预测。掩码解码器输出一个大尺寸的二维掩码，随后将其还原为三维体积，作为 SSL 模块的伪标签。
- 再训练 SSL 模块 ：如上所述，训练数据包括有标签数据集 L = (Xl i , Yi) m i=1 和无标签数据集 U = Xu j n j=1。重新训练 SSL 模块 FS(-) 的训练目标可表述为 ：
其中，Lsup 和 Lunsup 分别是监督和非监督项，λ 则是它们之间的权衡。伪标签指导 Lpl 是一种无监督损失。
本文为再训练模块研究了四种三维医学影像半监督方法，包括两种经典方法（即自训练和平均教师）和两种先进方法（即 ACMT 和 MagicNet）。
C. SFR+ Framework
为了进一步增强所提的框架，开发了 SFR 的扩展版本，命名为 SFR+，如图 8 所示。在 SFR 框架中，使用标注样本训练的微调模块向再训练模块单向提供未标注样本的伪标签 。为了更有效地从两个模块中提取和利用未标记数据的信息， SFR+ 引入了置信度估计策略，以区分置信样本和不确定样本 。这种方法可以对微调模块和再训练模块进行选择性优化 。
- 置信度估计 ：SFR+ 引入了置信度估计，以确定如何处理每个未标记样本。计算每个未标记样本的体素级平均置信度，并根据阈值对它们进行分类 。对于每个未标记样本 Xu j，体素级平均置信度为：
其中，FS(Xu j,v) 是模型对体素 v 的预测值，c 是类别索引。样本根据阈值 τ 进行分类：
-高置信度样本：如果 Cu j ≥ τ，则认为该样本可信，并更新微调模块。
-低置信度样本：如果 Cu j < τ，样本将被归类为不确定样本，并发送至再训练模块。
- 选择性训练策略 ：在 SFR+ 中，微调和再训练模块中的选择性学习可以更有效地处理未标记样本。一方面，高置信度样本被用于更新微调模块，确保只有来自未标记数据的可靠信息才有助于进一步完善。另一方面，低置信度样本可从伪标签中获益，从而进一步改进再训练模块。
通过有选择性地在两个模块之间交替更新，SFR+ 降低了由于伪标签单向传输而可能产生的不准确预测所带来的错误传播风险。