0-AutoProSAM Automated Prompting SAM for 3D Multi-Organ Segmentation

AutoProSAM : Automated Prompting SAM for 3D Multi-Organ Segmentation

AutoProSAM ： 用于 3D 多器官分割的自动提示 SAM

arXiv:2308.14936v3 [cs.CV] 26 Jun 2024

背景

SAM在临床环境中，创建有效的提示非常具有挑战性和耗时，需要医生等领域专家的专业知识。最初为 2D 自然图像设计的 SAM 在 3D 医学图像分割任务上执行次优任务。这种性能归因于自然图像和医学图像之间的领域差距以及 2D 和 3D 图像之间的空间特性不同，特别是在多器官分割中。本文作者提出AutoProSAM 利用 SAM 的基础模型功能，自动执行基于 3D 多器官 CT 的分割，而无需依赖领域专家的提示。该方法利用参数高效的适应技术使 SAM 适应 3D 医学图像，并结合了特定于该领域的有效自动提示学习范式。 通过消除对手动提示的需求，它增强了 SAM 的 3D 医学图像分割能力，并在基于 CT 的多器官分割任务中实现了SOTA性能。

自然图像与医学图像数据集之间的差异（a）手动提示流程（b)

为了实现 SAM 的提示生成过程的自动化，主要采取了两个方向。

1. 涉及为 2D 图像分割或 3D 图像中的每个 2D 切片生成 2D 边界框提示。虽然这些方法是一个很好的起点，但它们并没有解决 3D 图像中固有的 3D 空间分辨率问题。
2. 训练辅助网络来为 SAM 从输入图像生成提示。但是，这些模型未能充分利用 SAM 广泛的预训练功能。

本文的方法通过利用 SAM 编码器的功能来生成自动提示，从而克服了这些限制。因此，解决了以前工作的缺点，并解锁了 SAM 在基于 CT 的 3D 多器官分割任务中的全部潜力。

实验方法

SAM架构

具体架构详见SAM。

尽管 SAM 在 2D 自然图像分割方面取得了成功，但在应用于 3D 体积医学图像时，它面临着重大挑战。一个关键问题是该模型依赖于切片预测，而切片预测没有考虑切片间空间上下文，从而影响了其对复杂医疗任务的适用性。此外，医学图像和自然图像之间固有的域差异导致它们在医疗应用中的性能受到限制。为了有效应对这些挑战并优化医学成像任务的 SAM，量身定制的适应和模型的微调变得至关重要。

处理 3D 医学输入

为了增强 SAM 处理 3D 医学图像的能力，提出了一种名为 AutoProSAM 的适应策略，如图所示。该策略有两个主要目标：首先，使模型能够直接学习 3D 空间模式，其次，通过从预训练模型继承大多数参数，同时引入易于调整的增量参数来确保连续性。详细设计阐述如下。

位置编码增强功能

在预训练的 ViT 模型中，有一个大小为 C×H×W 的查找表，其中 C 代表通道，H 代表高度，W 代表宽度。此外，用零初始化一个大小为 C × D 的可调查找表（其中 D 表示卷补丁的深度）。为了获得 3D 点 （d，h，w） 的位置编码，将带有 （h，w） 的冻结查找表中的嵌入添加到带有 （d） 的可调查找表中的嵌入中。

Patch 嵌入调整

利用 1×k×k 和 k×1×1 3D 卷积的组合来近似 k ×k ×k 卷积的效果（例如，内核大小 k = 14）。1×k ×k 卷积是使用预训练 2D 卷积的权重进行初始化的，并且在微调阶段保持不变。对于新引入的 k × 1 × 1 3D 卷积，应用深度卷积来减少需要调整的参数数量。这种方法有助于管理模型的复杂性。

调整注意力块

注意力块可以直接调整以适应3D功能。对于 2D 输入，查询的大小为 [B，HW，C]，对于 3D 输入，可以毫不费力地将其修改为 [B，DHW，C]，同时保留所有预训练的权重。实现了类似于SwinUNETR中的滑动窗口机械，以减轻尺寸增加对内存的影响。这种方法有助于优化模型的性能，同时管理内存需求。

瓶颈修改

用对应的3D对应物替换瓶颈中的 2D 卷积，并从头开始训练它们以提高性能（将 2D ViT 平滑地过渡到 3D ViT，重复使用大多数参数）,并建立轻量适配器微调。

在整个训练阶段，只调整卷积、深度适配器和归一化层的参数，同时将所有其他参数保持在冻结状态。这种冻结方法可以提高训练期间的记忆效果。微调适配器和归一化层有助于弥合自然图像和医学图像之间的差距，使模型能够更有效地适应医学图像域。

自动提示生成器

整个过程如图所示。不是使用手动生成的点或边界框，而是在最后一个注意力和深度适配器作块之后直接获取输出特征图。该 APG 遵循完全基于卷积神经 （FCN） 的编码器-解码器设计。该生成器拥有轻量级结构，利用基于 3D 的卷积运算，可以毫不费力地从头开始学习。这能够为不同的医疗分割任务量身定制精确的提示生成。它消除了对额外手动生成提示的需要，从而简化和加快了多类医学图像分割任务。

MASK解码器

在多个级别上具有跳跃连接的 U 形网络对于成功至关重要，同时为了满足轻量化的需求，在解码器中加入了多层聚合机制 （MLAM） ，利用图像编码器第 1-4 阶段的中间特征图，以及来自 APG 的提示嵌入，在不影响效率的情况下丰富掩码特征图。为了提高分辨率细节，对蒙版特征图进行上采样以匹配原始分辨率，然后将其与原始图像连接起来。此串联映射使用另一个 3D 凸视图进行融合，以生成最终蒙版。这种方法在蒙版生成过程中有效地将高分辨率细节与原始图像数据相结合。简化了原来的 SAM，只专注于目标下游任务，省略了多任务生成和模糊感知等功能。掩码解码器的主干主要由轻量级 3D 卷积层组成。

实验结果

实验结果如图。

可视化结果。

作者同时也对基于SAM的其他模型进行比较。

总结

本文介绍了 AutoProSAM，这是一种旨在增强 3D 多器官医学图像分割的 SAM 的新方法。通过将 SAM 从 2D 自然图像扩展到 3D 医学图像，通过参数高效的微调和自动提示生成器 （APG） 来解决域差距和空间差异，该生成器可自动创建提示，无需手动输入。在基于公共和私人 CT 的数据集上进行的广泛实验表明，AutoProSAM 的性能优于最先进的模型，更高的 Dice 和 NSD 分数证明了这一点。