SAM-Med | 英伟达基于SAM提出医学图像标注效率神器

集智书童公众号

发布于 2023-09-04 11:06:15

8500

发布于 2023-09-04 11:06:15

文章被收录于专栏：集智书童

近年来，大型视觉模型——SAM已经彻底改变了计算机视觉领域，尤其是图像分割。SAM提出了一种新的快速分割范式，展示了其显著的零样本泛化能力。广泛的研究探索了SAM在各种下游任务中的潜力和局限性。在这项研究中介绍了

SAM^{Med}

，这是一个利用SAM功能的医学图像标注增强框架。

SAM^{Med}

框架由2个子模块组成，即

SAM^{assist}

和

SAM^{auto}

SAM^{assist}

使用即时学习方法展示了SAM对下游医学分割任务的泛化能力。结果显示，仅使用大约5个输入点就显著提高了分割精度。

SAM^{auto}

模型旨在通过自动生成输入Prompt来加快标注过程。

所提出的SAP-Net模型仅用5个带标注的slices就实现了卓越的分割性能，肾脏和肝脏分割的平均Dice系数分别为0.80和0.82。总体而言，

SAM^{Med}

在医学图像标注方面显示出了有希望的结果。这些发现突出了在医学图像标注任务中利用大规模视觉模型的潜力。

1、简介

近年来，大语言模型（LLM）取得了显著的成功，并引起了学术界和工业界的广泛关注。值得注意的是，ChatGPT是LLM最突出的模型之一，显示出其理解和生成类人语言的非凡能力。它充分展示了大型模型的潜力和非凡的可扩展性。它首次展示了通用人工智能以前所未有的熟练程度理解和处理人类语言的非凡潜力。这一突破性的成果不仅彻底改变了自然语言处理，而且影响了计算机视觉领域。

最近，Kirillov等人提出了“Segment-Anything模型”（SAM），这是一个在迄今为止最大的分割数据集上构建的大型视觉模型，在1100万张许可的自然图像上有超过10亿个Mask。值得注意的是，SAM展现了卓越的零样本迁移能力，使其能够熟练地概括以前未见过的数据。

通过定义一个新的可Prompt分割任务，SAM根据特定的分割Prompt（如指定目标的点和边界框）有效地分割图像。这种方法为SAM提供了一个独特的优势，可以适应不同的输入线索，促进精确和不同的目标分割，有助于其在具有挑战性的场景中的卓越性能。基于Prompt的分割丰富了其功能，使其成为各种图像分析应用中的宝贵资产。

1.1、Medical Annotations

医学图像标注是医疗保健和医学研究领域的一项重要任务。医学图像的准确有效标注对于疾病诊断、肿瘤检测等各种下游应用至关重要。传统的标注工具，如ITK-SNAP、3D Slicer，在医学成像界已经很常见。这些工具专注于为用户提供稳定的标注平台。随着深度学习技术在医学分割中的发展，以神经网络为分割Backbone开发了几种先进的图像标注工具，如RITM、SimpleClick和MIDeepSeg。

1.2、About Segment-Anything Model

SAM由3个核心组件组成：图像编码器、Mask解码器和Prompt编码器。图像编码器使用基于视觉Transformer的架构来有效地提取图像特征。在这方面，SAM提供了3个预训练的图像编码器，每个编码器被定制为不同的尺度，表示为ViT-B、ViT-L和ViT-H，分别包含91M、308M和636M个参数。预训练的图像编码器在将高分辨率图像（尺寸为1024×1024）嵌入到低分辨率图像嵌入（尺寸为64×64）中起着关键作用。

另一方面，Mask解码器采用了轻量级的Transformer网络，能够根据提取的图像特征和给定的Prompt信息生成精确的图像分割结果。Mask解码器包含两个Transformer层，一个动态Mask预测头和一个IoU分数回归头。

为了便于用户输入Prompt细节的全面表示，Prompt编码器被战略性地设计。SAM支持4种不同的Prompt类型：点、框、文本和Mask，每种类型都使用特定的编码方法来确保最佳表示。例如，点和框通过傅立叶位置编码与前景和背景的2个可学习标记相结合进行编码。MaskPrompt保留与输入图像相同的空间分辨率，并通过卷积运算进行编码。文本Prompt由CLIP中经过预训练的文本编码器进行编码。

1.3、SAM in medical imaging

SAM彻底改变了图像分割领域，促进了零样本处理，并支持对各种下游任务的进一步开发。在医学图像处理领域，SAM已成为研究的焦点，大量预印本论文证明了这一点。

在这些文献中，有几项研究专门关注利用SAM在特定医学领域进行分割任务。其中包括病理学图像分割、CT扫描中的肝肿瘤分割、腹部CT图像分割、MRI中的脑肿瘤分割、息肉检测和眼科图像分割。同时，一些工作进行了更全面的评估，比较了多个器官部位和各种医学成像模式下基于SAM的分割。

值得注意的是，迄今为止，Mazurowski等人已经进行了全面的评估实验。他们发现，在医学成像中，每个对象上的绘图框Prompt产生了最佳的分割性能。这一发现与其他研究结果非常一致。除了研究分割性能外，Ma等人和Li等人对SAM进行了微调，以解决特定的医学分割应用。他们的工作证明了在医学图像上微调SAM的显著效果，从而提高了医学成像目的的分割精度。

此外，SAM的影响延伸到了各种应用，正如一项成功将SAM集成到3D slices器中的著名研究所强调的那样。这种集成大大增强了SAM在医学成像界的实用性和可访问性。

1.4、Contributions

在本研究中提出了由2个子模块组成的

SAM^{Med}

框架，即

SAM^{assist}

和

SAM^{auto}

，以利用大型视觉模型SAM的显著泛化能力来加速医学图像标注过程。通过利用SAM模型的强大功能，显著降低了标注医学图像的复杂性和耗时，从而提高了整个标注工作流程的效率。

所提出的

SAM^{assist}

模型采用即时学习技术，有效地将下游医学分割任务适应SAM。即时学习的使用使SAM能够在最小的额外计算资源下实现显著的性能改进，使其成为医学图像标注任务的高效解决方案。

提出了所提出的

SAM^{auto}

模型，以进一步促进标注过程。通过引入Prompt生成器模块，可以有效地减少手动交互过程。研究了三种不同的自动Prompt生成方法。还提出了一个空间感知的原型网络，即SAP-Net，用于准确的Prompt生成。通过利用few-shot学习的能力，SAP-Net能够自动生成SAM所需的Prompt。

2、相关工作

由于医学图像标注被广泛认为是一项艰巨的任务，这对该领域的人类专家提出了重大挑战。因此，人们致力于解决这一问题，以期找到有效的解决方案。

已经探索的一种方法是利用最小的用户交互，特别是通过极端点击，来减轻人类专家的负担。Roth等人采用极值点作为注意力机制来增强随机walker算法生成的初始粗分割的分割结果。类似地，Luo等人利用极值点生成距离图，作为神经网络的外部线索，并通过额外的用户点击来细化初始分割。

另一种解决方案是使用最小的标注数据来实现与完全监督分割相比可比较的分割精度。Roy等人采用了一种少快照学习方法，使用有限数量的带标注的支持示例训练体积分割模型，证明了零样本传输能力。蔡等人提出了一种半监督技术，将正交标注slices作为一致性项，解决了医学图像分割中的弱标注和半监督挑战。

其他标注框架，如MONAI Label和AbdomenAtlas，利用预训练的医学分割模型的能力来促进高效的人在环标注过程。这些全监督方法的有效性在很大程度上取决于预训练模型，从而限制了它们进行零样本分割的能力。在需要标注一个看不见的类的情况下，有必要重新训练或微调底层模型。

如第1.3节所述，越来越多的研究将SAM纳入其特定的医学图像分割任务中。在这些工作中，马等人和李等人对SAM进行了微调，以使其适应其医学分割应用。通过微调SAM的Mask解码器模型，获得了显著的改进。值得注意的是，通过微调SAM的Mask解码器模型实现了显著的改进。然而，与上述工作类似，对下游任务的SAM进行微调会损害大型模型的泛化能力。

3、本文方法

在本节中，首先介绍了问题设置，然后详细介绍了

SAM^{Med}

标注工具的架构设计和训练策略，最后描述了评估策略。

SAM^{Med}

主要由2个模块组成：Assist和Auto模块，即

SAM^{assist}

和

SAM^{auto}

。

SAM^{assist}

旨在帮助用户更有效地对医学图像进行标注。

SAM^{auto}

旨在自动生成标注，而无需用户手动操作。

3.1、Problem setting

SAM^{Med}

被设计为医学图像的标注工具。首先，

SAM^{Med}

需要初始标注

D_{Train}=\{(I^i_T,L^i_T)\}^N_{i=1}

来训练模型。初始手动标注的

D_{Train}

包括N对原始图像slices

I_T

及其相应的二进制标签

L_T

。

SAM^{assist}

模型的目的是从

D_{Train}

中学习模型

F_{assist}(·)

，以便在给定新的查询图像slices

I_Q

及其相应的手动输入Prompt

P_Q

的情况下，推导出二进制分割

L_Q

。

SAM^{auto}

的目标是从

D_{Train}

中学习一个模型

F_{auto}(·)

，以便只给定一个新的查询图像slices

I_Q

，就可以推断出二值分割

L_Q

。

3.2、Architectures

如前所述，建议的

SAM^{Med}

由2个模块组成：SAMasiset和

SAM^{auto}

。这两个模块协同工作，最大限度地提高了医疗标注程序的效率。

1、

SAM^{assist}

正如许多先前的工作所研究的那样，原始SAM模型在各种医学图像分割任务上的性能有限。分割性能在很大程度上取决于用户的输入Prompt。尽管SAM在许多视觉任务中表现出令人印象深刻的零样本性能，但其有效性仍然受到自然图像和医学图像之间的领域差距的影响。

为了利用大视觉模型显著的泛化能力，作者结合了即时学习方法来重新定义大视觉模型的下游任务。与专注于调整原始模型以适应下游任务的传统基于微调的方法相比，基于即时学习的方法旨在调整下游任务以适应原始大模型。与基于微调的方法相比，即时学习的关键优势是能够在显著减少训练数据需求和训练耗时的情况下取得有效结果。

SAM^{assist}

的概述如图2所示。训练数据集

D_{Train}=\{(I^i_T,L^i_T)\}^N_{i=1}

由N对带标注的图像slices

I_T

和相应的标注

L_T

组成，

I^i_T∈I_T

，

L^i_T∈L_T

。

SAM^{assist}

的目标可以描述为：

其中

P_T

表示由从标注

L_T

提取的随机生成的Prompt组成的输入Prompt集。

是SAM模型

的可训练参数。具体而言，

对应于Prompt编码器的参数。

表示损失函数，如交叉熵损失、Dice损失等。

通过学习一种新的Prompt编码器，重新定义了下游分割任务的语义，从一般的前景/背景视觉分割转移到特定的医学分割，例如肾脏/背景、肝脏/背景分割。

2、

SAM^{auto}

与

SAM^{assist}

模型不同，

SAM^{auto}

模型被设计为在没有用户交互的情况下自动生成标注。

SAM^{auto}

旨在进一步推进医学标注程序。

SAM^{auto}

的体系结构框架如图3所示。引入了一个“Prompt生成器”模块，表示为G，它可以自动生成查询图像的潜在Prompt，而不是依赖用户的手动输入Prompt。该模块在有限数量的手动标注slices上进行训练。

SAM^{auto}

中使用的图像分割模块是预训练的

SAM^{assist}

模型。因此，

SAM^{auto}

模块的目标可以描述如下：

因此，主要的挑战在于为SAM模型开发一个应用程序)的Prompt生成器。可以探索各种方法来生成查询图像slices的Prompt。在本研究中，主要研究了三种方法：

Prompt propagation

从众多标签传播方法中获得灵感，Prompt传播是一种为新图像生成潜在Prompt的简单直观的方法。考虑到医学图像slices之间的一致性，基于带标注的种子Prompt，Prompt传播可以有效地推断出可能的Prompt。然而，采用即时传播方法的关键挑战之一是如何建立适当的传播标准。

此外，即时传播的几个固有缺陷仍有待解决。首先，它无法处理离散对象，如细胞核和肺结节分割任务。其次，由于其组织组成的多样性，传播背景Prompt是困难的。

在这项工作中，通过动态阈值方法实现了一种简单的即时传播策略。具体的算法细节可以在算法1中找到。前景点将基于强度一致性原则递归地传播到相邻slices。

Prompt classification

生成Prompt的另一种方法是Prompt分类，它直接将Prompt（最好是点Prompt）分类为前景类和背景类。这种方法利用带标注slices的固有特性来训练直接的分类器。通过分析SAM训练有素的图像编码器提取的图像特征，分类器学习区分前景Prompt和背景Prompt。

与Prompt传播方法相比，基于分类的方法提供了一种更系统、更灵活的Prompt生成方法。通过利用标注数据中存在的丰富信息，分类器可以很容易地应用于其他未标记的数据，而无需种子Prompt。尽管有其优点，但迅速分类也带来了一些挑战。设计一个有效的分类器来抑制有限标注数据上的过拟合问题是快速分类方法的关键。

Coarse segmentation

除了前面提到的方法外，作者还研究了作为Prompt生成过程的一部分，直接从查询图像slices中分割对象的可行性。粗略分割通过提供前景的粗略空间定位来提供Prompt生成的另一种替代方法。通过直接分割目标，作者的目标是捕捉它们的边界和空间范围，这在生成与感兴趣的目标对象更紧密对齐的Prompt时非常有用。

与专注于生成Prompt本身的Prompt传播和Prompt分类不同，粗分割旨在通过明确地描绘前景区域来定位和提取感兴趣对象特有的相关特征。请注意，粗分割的目的不是取代SAM的Mask解码器，它试图利用有限的标注来定位目标对象，并为训练有素的SAM模型生成准确的Prompt。

在这项工作中，使用few-shot学习技术设计了一个轻量级的ROI分割器。与传统的全卷积网络（FCN）相比，如U-Net和Deeplab，它们是为全监督精细分割而设计的，few-shot分割方法旨在解决有限标记数据的挑战。作者提出的空间感知原型网络，即SAP-net，试图利用相邻医学slices之间的位置一致性和原型学习的能力来准确定位目标区域。

SAP-Net的体系结构如图4所示。使用PANet作为Backbone。PANet是原型网络的一个变体，它引入了Assist对齐损失来正则化支持集和查询集上的原型。

在训练阶段，所有带标注的slices

D_{Train}=\{(I^i_T,L^i_T)\}^N_{i=1}

将被划分为支持集

S=\{(I^j_s,L^j_s)\}^J_{j=1}

和查询集

Q=\{(I_q^k,L^k_q)\}^K_{k=1},J+K=N

。没有直接微调SAM的预训练图像编码器，而是添加了一个小型可训练FCN网络来调整预训练图像编码的特征。

遵循PANet的相同架构，使用分割损失

L_{seg}

和

L_{alig}

来优化网络。

L_{seg}

鼓励使用来自支持集的原型对查询图像进行高分割质量。相反，

L_{alig}

鼓励从查询生成的原型与支持集的原型很好地对齐。具体来说，

L_{seg}

可以公式化如下：

其中d表示诸如Dice距离、交叉熵距离之类的任何距离测量。

\hat L_q

表示查询图像的预测Mask，其计算如下：

cos表示余弦距离，该余弦距离用于测量原型

p_c

和图像特征

F_q

之间的距离。

P=\{p_c|c∈C_i\}

是包含所有类原型的原型集。

是一个权重因子。原型通过掩蔽全局平均池化计算。

为了充分利用相邻slices之间的位置一致性，通过引入位置编码层来提供空间感知功能[ref59]。因此，图像特征

F_q

可以通过以下方式获得：

其中，

\varepsilon

是SAM的预训练图像编码器，

\tau

是从

\varepsilon

调整图像特征的附加FCN网络。

[·,·]

表示连接操作。

γ(·)

是位置编码器层，其计算如下：

其中

B∈R^{2×d}

是从

N(0,σ^2)

分布中采样的。

是位置嵌入特征维度。反之亦然，关于支撑集，

L_{align}

之前的所有计算都是相同的。

4、实验

4.1、

SAM^{assist}

鉴于

SAM^{assist}

是专门设计的半自动分割工具，主要目标是评估其在最大限度地减少手动交互方面的功效。

SAM^{assist}

的评估包括4个关键的分割任务：肾脏分割、前列腺分割、肝脏分割和脾脏分割。

SAM^{assist}

模型将首先在几个随机选择的带标注的slices上进行训练，然后在其余slices上进行测试。为了模拟真实的标注过程，避免假设用户输入的任何先验分布。相反，从前景和背景区域随机采样点来训练新Prompt。每个案例都将单独进行训练和测试。

图5展示了

SAM^{assist}

模型在4种不同固体组织中的性能。如图所示，使用

SAM^{assist}

，即使只有2个带标注的slices用于即时训练，目标上的单点也显示出显著的分割性能。

详细的分割结果如图6所示。一个测试slices的分割结果证明了

SAM^{assist}

的显著改进。

如前所述，与微调方法相比，快速学习方法需要更少的资源。表1不仅提供了数据需求，还提供了时间成本，概述了与4个细分任务相关的时间成本。该表表明，2个和5个标注slices的成本分别约为9秒和30秒。考虑到标注过程涉及人在环，以秒为单位的训练时间成本不太可能对用户体验产生实质性影响。

4.2、

SAM^{auto}

在这项工作中，主要研究了显式Prompt生成的“prompt propagation”和“coarse segmentation”方法。至于算法1中描述的“prompt propagation”策略，从标注slices中随机选择的点将利用直接的阈值规则传播到相邻slices。当没有满足指定标准的点时，此迭代过程将终止。

为了提高鲁棒性，将传播单独应用于每个带标注的slices。因此，在两个带标注的slices的情况下，传播被执行两次，然后所得到的分割被集成以获得最终的分割。点传播过程的一个示例如图7所示。