2D医学图像分割大模型：SAM-Med2D

pub: 30 August, 2023 ArXiv
[ paper] [ code]

Introduction:

1 现有医学图像分割方法仅适用于特定的模态、器官或病变。
2 通用AI模型如ChatGPT2、ERNIE Bot 3、DINO、SegGPT、SAM等使得一个模型可以处理多个任务。
3 因为SAM只在自然图像上进行训练，缺乏与医学影像相关的特定知识，导致SAM在多模态和多目标医学数据集上很难泛化。如图1所示，自然图像和医学图像在像素强度、颜色、纹理和其他分布特征方面存在显著差异。

将SAM从自然图像转移到医学图像上，提出SAM-Med2D以解决以下问题：

● 如何将SAM在医学图像领域进行微调？

在将SAM迁移到医学领域时，不同的提示模式在不同的分割任务中起着重要作用。如图2（c）所示，使用边界框提示可以得到相对准确的息肉掩膜。相反，当点击前景点时，掩膜质量较差（列d）。随着点击次数的增加，分割结果逐渐改善，甚至超过了边界框提示的性能（例如肝脏分割）。在心肌分割时，使用边界框提示可能会在结果中包含不感兴趣的区域，而点提示允许我们逐渐获取所需的掩膜。因此，本文旨在对三种提示模式（点、边界框和掩膜）进行微调，以满足医学图像分割任务中不同场景的需求。根据表1，与其他方法相比，我们的方法进行了更全面的微调，这意味着我们的方法在医学图像领域具有更广泛的提示分割能力。

表 1 说明了在特定的有限规模医疗数据集上微调 SAM 的方法。 虽然这些方法有经验证的有效性，它们仅增强 SAM 在特定场景下的分割能力类似于训练数据集。 因此，它们对于更多样化的医疗的适用性受到限制。

● 如何评估SAM-Med2D的性能？

为了全面评估SAM-Med2D在医学影像领域的性能，我们从多个角度进行了评估。主要关注以下几个方面：1）数据多样性：我们在不同的成像模态（如CT、MR、X射线）、器官和多种病理条件（如肿瘤、炎症等）上评估SAM-Med2D，以确保其能够对不同类型的医学图像进行分割。2）微调策略：我们将SAM-Med2D的默认微调策略与其他替代策略进行比较，例如不同的模型配置，以展示SAM-Med2D在医学影像领域的成功应用。3）泛化能力：我们使用9个MICCAI2023数据集评估SAM-Med2D的鲁棒性，以确保其在未见过的医学图像环境中能够进行准确的分割。
通过综合评估，我们发现SAM-Med2D具有以下能力：1）医学专业知识：相比于SAM，SAM-Med2D在处理复杂的器官结构、病变以及边界不清晰的情况下表现出更好的性能。这意味着SAM-Med2D能够准确识别和分割医学图像中具有挑战性的区域，从而提供更精确的诊断和治疗支持。2）广泛的分割能力：SAM-Med2D展现了在不同提示模式下广泛的分割能力，使其能够在不同场景下完成分割任务。这意味着医生和医学影像专业人员可以利用SAM-Med2D进行更精确和准确的分割操作，从而提高医学图像分析结果的效率和可靠性。3）泛化能力：SAM-Med2D展现出强大的泛化能力，能够直接应用于未见过的医学图像数据，并产生良好的分割结果。

Methodology：

在介绍SAM-Med2D之前，让我们简要回顾一下SAM的架构。SAM由三个主要组件组成：大规模图像编码器、提示编码器和轻量级掩膜解码器。这个框架允许基于不同的提示为同一图像生成不同的掩膜。图像编码器利用预训练的Visual Transformer (ViT) [22]来处理高分辨率入，并输出与原始图像尺寸相比为1/16比例的特征图。提示编码器包括稀疏提示和密集提示，将点、边界框或文本映射到256维向量，并在每个级别上对密集提示进行带有GELU激活函数的卷积下采样。掩膜解码器通过交叉注意机制接收来自两个编码器的嵌入信息，并更新图像嵌入和提示嵌入。在这项工作中，我们对SAM进行微调，创建了SAM-Med2D，有效地将该框架扩展到医学图像领域。现在，我们将详细讨论SAM-Med2D的每个组件和微调策略：

●Adapting图像编码器

作为SAM中参数最多的部分，在微调过程中全局更新图像编码器会产生巨大的计算成本。为了以较低的成本将医学领域的知识纳入图像编码器中，我们引入了适配器（ Adapter）技术。具体而言，在微调过程中冻结原始图像编码器的所有参数，并为每个Transformer块部署一个适配器，如图4所示。我们在通道和空间维度上对图像编码器进行适应。对于通道维度，我们首先使用全局平均池化将输入特征图的分辨率压缩到C×1×1。然后，我们使用一个线性层对通道嵌入进行压缩，再使用另一个线性层进行恢复，压缩比为0.25。最后，我们通过sigmoid函数获取通道维度的权重，并将其与输入特征图相乘，作为下一级的输入。对于空间维度，我们使用卷积层将特征图的空间分辨率下采样两倍，然后使用转置卷积进行空间分辨率的恢复，保持与输入相同的通道数。在每个适配器层之后添加跳跃连接。

●提示编码器和掩膜解码器

SAM中的提示编码器支持四种类型的提示：点、边界框、掩膜和文本提示。由于缺乏针对医学图像-文本对齐的大规模预训练模型，文本提示的使用受到限制。因此，在微调过程中，我们只考虑剩下的三种提示模式。与之前只微调单个提示的方法相比[20, 21]，我们保留了提示的全部功能，并增强了它们在医学影像领域的适用性。具体而言，SAM-Med2D同时使用稀疏提示（点和边界框）和密集提示（掩膜）。对于稀疏提示，每个点被表示为其位置编码的向量嵌入，以及指示其前景或背景位置的两个学习嵌入的和。每个边界框使用其左上角和右下角的位置编码，以及表示“左上角”和“右下角”的学习嵌入作为向量嵌入。对于密集提示，我们使用模型第一次迭代后生成的低分辨率特征图作为掩膜提示，应用两个卷积嵌入来将输入掩膜的规模缩小4倍，输出通道数为原始输入的1/4和1/16。最后，使用1×1卷积将通道维度映射为256。
我们没有对掩膜解码器的结构进行任何更改，并在训练过程中持续更新其参数。为了使模型能够识别模棱两可的情况，每个提示同时预测多个掩膜（默认为三个）。在反向传播过程中，我们选择与基准标注相对于地面实况具有最高交并比（IoU）分数的预测掩膜计算损失，并相应地传播梯度。我们将上一轮生成的低分辨率特征图映射到[0, 1]的范围内，作为当前迭代的密集提示。在实际的训练过程中，我们观察到即使只有稀疏提示，模型也可以迅速收敛，这削弱了密集提示的影响。因此，我们采用了SAM的训练策略，在最后一次迭代和一个随机的中间迭代中，只提供密集提示，以鼓励模型从提供的掩膜中获益。

●微调策略

与SAM和其他交互式分割方法类似，我们通过模拟交互式分割来训练SAM-Med2D。对于每个数据批次，我们训练模型进行9次迭代。在第一次迭代中，我们以相等的概率随机选择前景点或边界框作为稀疏提示。前景点从真值中采样，边界框是真值的最大外接矩形，每个坐标偏移不超过五个像素。值得注意的是，除了第一次迭代，适配器层、提示编码器和掩膜解码器的参数同时更新，后续迭代只更新掩膜解码器的参数。从第二次迭代开始，我们从先前掩膜预测与真值之间的误差区域中随机选择1、3、5或9个点作为后续的稀疏提示，允许用户执行单点或多点的交互式分割。SAM-Med2D旨在通过针对医学影像领域专门调整SAM框架，提高医学图像的分割效果。它采用了适配器技术，扩展了提示的功能，并使用基于模拟交互式分割的微调策略。

●Evaluate SAM-Med2D

针对之前评估匮乏数据和缺乏一般医学图像分割方法基准的限制，我们将进行SAM-Med2D的全面多维度评估，为未来研究人员提供交互式分割方法的基准。在模型方面，我们以SAM作为基线模型，并从其交互方法中选择了两种直观的提示模式，即Bbox和Points，来评估SAM-Med2D。这个选择是因为Bbox和Points是常用的交互模式，可以作为简单高效的注释方法。通过在这两种交互模式下评估SAM-Med2D的性能，我们可以深入探讨它们在医学图像分割任务中的优势和局限性。Bbox交互模式通过边界框引导算法进行分割，直观且容易理解。它提供了目标的粗略位置和形状信息，从而指导算法进行更准确的分割。然而，在某些情况下，准确地使用Bbox包围目标可能具有挑战性，特别是对于形状复杂、边缘模糊或重叠对象的目标。这可能导致分割结果不准确和遗漏分割的情况。另一方面，点交互模式通过标记目标的关键点或区域来引导进行分割。这种模式可以提供更精确的分割指导，特别适用于形状复杂或局部细节丰富的目标。这些评估将帮助我们更深入地了解不同交互模式在医学图像分割中的优势和局限性，从而激发研究人员设计和开发更灵活高效的交互模式。
在数据方面，我们将在包括MRI、CT、超声等10种不同模态的医学图像上评估SAM-Med2D的性能。这个全面的评估将使我们能够了解SAM-Med2D在特定模态下的优势和挑战，并揭示其在多模态图像中的潜在应用能力。不同的医学图像模态具有独特的特征和噪声来源，因此在这些模态上评估SAM-Med2D的性能将提供更全面的了解其适用性和鲁棒性。此外，考虑到不同的解剖结构和器官具有独特的形态、特征和变异模式，我们对四个解剖结构和31个主要器官进行了SAM-Med2D的评估。这些评估帮助我们深入了解SAM-Med2D在不同场景中的性能差异，并针对特定结构和器官的挑战进行有针对性的改进。最后，我们非常重视SAM-Med2D的泛化能力，因此在9个MICCAI 2023数据集上进行了测试。这些数据集代表了来自不同来源、机构或设备的医学图像，提供了多样性。通过在这些数据集上评估SAM-Med2D，我们可以验证其对新数据的泛化能力。这在验证我们的方法在广泛的临床场景和数据源中的适用性方面是至关重要的。

通过以上全面评估，我们将能够深入了解SAM-Med2D的性能和适用性，为未来的研究人员和开发者提供有价值的参考和基准。这些评估结果将对医学图像分割的发展产生积极影响，推动设计和应用更准确高效的方法。

Dataset：

最近的研究再次证实了训练数据量在大型模型的学习能力中的关键作用[7, 8, 23]。通过从更大规模的数据中进行学习，模型可以获得更丰富的领域特定知识，并更好地适应各种应用场景。尽管SAM是在超过10亿个掩膜上进行训练的，但在医学图像分析领域，由于自然图像和医学数据之间存在显著的领域差距，其性能并不理想。为了弥补这一差距，我们收集和整理了迄今为止最大的医学图像分割数据集。该数据集由多个公共和私有数据集组成，确保了全面的覆盖范围和多样性。图3（b）展示了数据集的10种不同成像模态及其相应的数据比例。为了提高视觉呈现效果，我们使用对数尺度来可视化数量上的差异。基于解剖结构和病变的存在，我们将数据集分为头颈部、胸部、腹部、盆腔和病变等类别（图3（c））。此外，我们从这些数据集的271个标签中整理和合并了31个主要器官，如图3（a）所示。这几乎涵盖了当前公开数据集中所有可用对象类型，弥补了SAM在医学领域知识方面的不足。

Experiment:

●实验细节

我们的方法使用PyTorch实现，并在8个NVIDIA Tesla A100 GPU上进行训练，每个GPU具有80GB内存。考虑到内存限制，我们在这项工作中只对SAM（SAM-B）的基础模型进行微调。我们使用Adam优化器，初始学习率为1e-4，总共训练12个epochs，在第7和第10个epochs时将学习率除以2。在训练过程中，所有图像都被调整为256x256的分辨率。我们的调整策略是对宽度和高度均小于256的图像边缘进行零填充，而对其他情况下的图像进行双线性插值调整。对于每个图像，我们随机选择5个对应的掩膜。如果可用的掩膜较少，我们会随机复制样本。为了充分利用GPU内存，每个GPU同时处理50个图像及其对应的250个掩膜。监督掩膜预测的损失函数是focal loss [31] 和dice loss [32]的线性组合，比例为20:1。此外，还使用IoU预测与预测掩膜与真值掩膜之间的均方差损失进行训练。我们使用Dice系数来评估分割结果。

●总体性能

表2呈现了SAM、FT-SAM（仅微调掩膜解码器）以及我们的SAM-Med2D在测试集上的总体性能结果。我们发现，在边界框提示（Bbox prompt）模式下，FT-SAM的Dice分数比SAM提高了11.93%，而我们的SAM-Med2D实现了更显著的性能提升，Dice分数为79.30%（即提升了17.67%）。这表明在大规模数据集上进行微调可以导致更好的目标领域可迁移性。我们还在点提示模式下模拟了交互式分割。在该模式下，我们从前景中随机采样一个点作为第一个提示点，并且后续的提示点随机选择在分割结果与真值之间的误差区域内。另外，前一次迭代生成的低分辨率掩膜作为掩膜提示，结合前面的提示点一起作为模型的输入。实验结果表明，使用单个点提示时，SAM的表现较差，即使分辨率达到1024×1024，Dice分数比FT-SAM低23.23%。随着提示点的增加，不同模型的性能显著提高，微调方法甚至超过了边界框提示模式。这证明在医学图像中使用基于点的交互式分割是可行且有效的。此外，对比1024×1024分辨率下的SAM整体分割性能不如经过微调的方法。这表明经过微调的模型学习到了医学领域的特定知识，而通过低成本的微调是减少领域差异的一种有效和可行的方法。

●在不同的解剖结构上

如图5（a）所示我们评估了不同模型和分辨率在头颈部（H&N）、胸部（Tx）、腹部（Abd）、盆腔（Pl）和其他区域的分割性能，其中“其他区域”包括病变和四个解剖结构之外的区域。我们主要关注使用边界框提示和1个点提示（1 pt prompt）的模型的分割性能。观察到，在使用边界框提示时，1024×1024分辨率下的SAM在胸部、腹部和其他区域的性能优于FT-SAM。然而，在头颈部区域，其表现较差。这可能归因于头颈部区域中病变或器官相对较小，边界不太清晰，使得模型在没有微调的情况下难以适应这种类型的分割任务。与其他方法相比，我们的SAM-Med2D在所有解剖结构的Dice分数上都具有优势。由于1个点提示提供的信息有限，不同类别之间存在性能差异。有趣的是，我们发现经过微调的SAM显著优于原始的SAM。这是因为微调方法从大规模医学影像数据集中学到了目标区域内点的位置关系，从而实现更准确的决策。基于上述结果，我们得出结论，SAM-Med2D在不同解剖结构的分割任务中表现出色，在盆腔和胸部区域的Dice指标方面取得了满意的结果。然而，值得注意的是，头颈部区域在不同模型和分辨率下的性能似乎相对较差，这表明需要采取额外的改进措施。

表3呈现了点提示模式下的性能结果。为了保持公平，我们在SAM（256）、FT-SAM和SAM-Med2D的预测中使用了相同的初始点。观察到，随着迭代次数的增加，分割性能显著提高，SAM在所有模态上实现了超过10%的改进。更令人印象深刻的是，我们的SAM-Med2D仅通过单点交互就超越了其他方法使用五个点进行交互的性能。这凸显了在医学影像领域中，大规模预训练支持的点提示可以有效应用，相比边界框提示实现更高效的交互。此外，我们注意到，在组织病理学和显微镜模态中，1024×1024分辨率下的SAM的多点交互实际上表现不如单点交互。这可能是因为模型已经根据单个点提示做出了最优决策，而进一步的点提示旨在纠正分割结果，可能会干扰模型的判断。这一观察结果也表明，微调可以增强SAM的性能上限。

●在不同模态数据上的性能评估中

图5（b）总结了四种方法在边界框提示模式下的表现。这四种方法在皮肤镜、内窥镜、眼底摄影、组织病理学和显微镜等****不同模态的数据中都达到了超过70%的Dice分数。当以1024×1024分辨率部署预测图像时，SAM在内窥镜、组织病理学和显微镜模态中优于其他方法。我们将这个结果归因于以下几个因素：1）这三种模态来自2D数据集，并由RGB图像组成，与自然图像具有相似之处。2）微调方法受到可用数据量的限制，导致性能受到一定程度的约束（如图中红线所示，表示经过对数转换的掩膜数量）。3）更大的图像分辨率提供了更多的细节，从而导致更高的预测性能。在相同的分辨率设置下，我们的SAM-Med2D明显优于SAM，可以有效处理所有成像模态的数据。直接比较不同模态之间的性能可能是不公平的，因为涉及到不同类型的对象和数据规模的变化。

●主要器官的表现评价

图5（c）呈现了FT-SAM和SAM-Med2D在30多个器官上的Dice分数差异。条形图表示使用5个点提示获得的结果，而五角星表示相应的边界框提示结果。SAM-Med2D在24个器官中获得了更高的结果，最大差异达到了6.95%。此外，观察不同提示方法下相同器官的实验结果，可以发现5个点提示和边界框提示之间的性能差距相对较小。这一发现表明，在将SAM应用于医学领域时，边界框交互可能不是最有效的方法（随着点交互数量的增加，边界框提示的绝对优势逐渐减弱）。对于某些骨骼区域，如肋骨、肩胛骨和锁骨，点提示策略可能更有效。总体而言，这些结果表明，使用适配器方法进行微调的模型可以实现出色的分割性能，并且点交互在器官分割中具有重要潜力。

●评估SAM-Med2D的泛化能力

如表4所示，我们在9个公开可用的数据集上进行了测试，其中大部分来自MICCAI2023竞赛。为了公平比较，我们统一使用256×256分辨率的图像进行测试。表4显示，SAM在边界框提示下表现出色，加权平均Dice指数为85.35%。然而，在单点提示下的性能并不令人满意（48.08%）。由于适配器层参数是即插即用的，我们测试了两种情况：保留和移除适配器层参数。当保留适配器层参数时，我们的SAM-Med2D在边界框提示下实现了81.93%的Dice分数，并且在移除适配器层参数时性能提高了8.19%。此外，我们观察到SAM未能有效适应点提示，在9个数据集中最好的分割性能仅达到51.05%。相比之下，我们的SAM-Med2D在点提示下获得了合理的分割结果。值得注意的是，在推理过程中移除适配器层参数时，SAM-Med2D在1个点提示下的性能与SAM在边界框提示下非常接近（83.41% vs. 85.35%），这为数据注释和分析节省了大量的时间和成本。总而言之，SAM只在边界框提示下显示出良好的泛化性能，而我们的SAM-Med2D在两种提示模式下实现了更好的泛化效果。

●定性比较

我们定性地比较了SAM-Med2D和SAM的分割掩膜。SAM的视觉结果来自于256×256或1024×1024分辨率中更好的那个。图6的前三行展示了两种方法在三种模态上的分割性能。在大多数情况下，由边界框提示得到的分割结果可以定位目标区域，但是我们SAM-Med2D的视觉结果边界更清晰、更接近真值。而对于1个点提示，SAM难以定位目标区域，导致分割结果与预期结果之间存在显著差异。最后三行展示了两种模型在肝脏、肺部和前列腺器官上的分割结果。在边界框提示模式下，两种方法生成了类似质量的掩膜。通过观察3个点和5个点提示的结果，我们可以看到更多的点提示会产生更好的分割效果。在有相同数量点提示的模型中，SAM-Med2D比SAM更好地描述了目标区域，这意味着SAM-Med2D需要较少的交互操作和更短的时间来达到所需的结果。这对于数据注释或伪标签生成具有优势。我们将这一现象归因于SAM-Med2D通过从大规模数据集中学习获得了与医学影像领域相关的特定领域知识。这与本文的动机相一致，即为SAM在稳健可靠的医学图像分割方面奠定基础。

图7展示了在同一图像内合并多个目标区域的结果。当目标边界清晰时，SAM和我们的SAM-Med2D之间存在微妙的视觉差异。在其他情况下，SAM-Med2D能够准确地分割出人眼难以识别的部分。另一方面，在1个点提示的场景中，SAM在许多器官上常常失败，并且难以定位目标区域。这再次证明，对SAM进行大规模数据微调可以实现更好的领域可迁移性。

Conclusion:

本研究通过在大规模医学图像数据集上对SAM进行微调，获得了SAM-Med2D，并能显著改善各种医学图像分割任务的性能。我们采用了两种明确的提示策略来生成掩膜，进行定量和定性比较。在相同的分辨率下，只有仅微调掩膜解码器（FT-SAM）在边界框提示模式下实现了11.93%的改进，而完全微调的SAM-Med2D实现了17.67%的改进。令人惊讶的是，我们的方法在1个点提示模式下表现出了压倒性的优势（18.94% vs. 70.01%）。此外，SAM-Med2D在两种提示模式下都展现出了优秀的泛化能力，表明其在医学领域具有实际价值。

从解剖角度来看，在1024×1024的分辨率下，SAM在胸部、腹部和其他区域的性能优于FT-SAM，而SAM-Med2D在整体分割性能上优于所有其他方法。关于不同模态的比较，SAM在目标模态数据类似于自然图像属性时表现出良好的泛化能力。我们对30多个主要器官进行了两种微调方法的比较，SAM-Med2D在24个器官上获得了更好的结果，最大改进幅度达到了6.95%。此外，在9个公开数据集上进行的泛化实验显示，基于大规模数据集预训练的模型具有很强的领域可迁移性。虽然边界框提示始终优于1个点提示，但增加更多的点明显改善了分割结果，甚至超过了边界框提示模式。在使用点提示模式时，SAM-Med2D能够更快地生成所需的掩膜，甚至优于使用边界框提示模式的其他方法。

然而，未来还需要进一步优化以实现更强大的医学图像分割。从定性的分割结果来看，对于复杂的形状/边界、小尺寸或低对比度的对象，不同提示模式产生的分割结果仍有改进的空间。在未来的工作中，建立相关的优化策略可能会提升分割结果，例如为不同器官设置窗口宽度，并为交互式分割设计边界损失。除了本文中使用的提示策略，自然语言也可以作为医学图像分割中常见的用户交互形式，但目前缺乏相关的数据集。这是我们积极追求的方向，旨在为SAM-Med2D赋予医学领域中的自然语言理解能力，以满足不同用户需求。尽管我们已经在超过1970万个掩膜上对SAM-Med2D进行了训练，但与SAM的训练数据量和分辨率相比仍存在差距。这导致SAM-Med2D能够有效处理常见器官或病变，但在医学领域的“万能性”上还有一定的不足之处。我们计划通过数据引擎生成更大量、更多样化的高质量掩膜，使SAM-Med2D能够真正地对各种类型的医学图像进行分割。

总而言之，本研究通过对大规模医学图像数据集上的SAM进行微调，获得了SAM-Med2D，并取得了令人满意的性能改进和泛化能力。我们将提供代码和预训练模型供研究人员使用，并希望该工作对医学计算机视觉领域的研究人员提供见解和机会，促进未来的研究和改进。

有任何问题欢迎联系：438285719@qq.com 共同学习交流