大模型系列 | 玩转nnUNet和SAM | 如何把SAM做到即插即用，nnSAM是最好的答案！

点击下方卡片，关注「集智书童」公众号

最近计算机视觉中基础模型的发展，特别是”Segment Anything Model” (SAM)，使可伸缩且领域无关的图像分割成为通用分割工具的可能。与此同时，医学图像分割领域受益匪浅，专门的神经网络如nnUNet，它在特定领域数据集上训练，可以自动配置网络以应对特定的分割挑战。

为了结合基础模型和领域特定模型的优势，作者提出了nnSAM，它将SAM模型与nnUNet模型协同整合，以实现更准确和稳健的医学图像分割。nnSAM模型充分利用了SAM的强大和稳健的特征提取能力，同时利用nnUNet的自动配置能力，促进数据集定制化学习。

作者对nnSAM模型在不同规模的训练样本上进行了全面评估，结果显示它支持少样本学习，这对医学图像分割非常重要，因为高质量的标注数据可能稀缺且昂贵。通过融合前两者的优势，nnSAM定位自己为医学图像分割的潜在新基准，提供了一种将广泛适用性与专门效率相结合的工具。

代码:https://github.com/Kent0n-Li/Medical-Image-Segmentation

1. 简介

在现代临床工作流程中，高效准确地分割医学图像对疾病诊断和预后、治疗计划和监测以及治疗结果跟踪等方面至关重要。传统上，医学图像分割是一项非常耗时和劳动密集的任务。基于深度学习的自动分割技术的出现显著减少了放射科医生和放射肿瘤学家所需的时间和精力。在为生物医学图像分割设计的众多深度学习架构中，U-Net以其能够有效且高效地捕捉全局和局部特征以获得更好的分割结果而脱颖而出。

基于U-Net，许多研究开发了具有各种修改的架构以适应不同任务。例如，TransUNet整合了U-Net和Transformer的优势，为医学图像分割定义了一个新的基准。通过利用Transformer的全局上下文理解和U-Net的精确定位能力，TransUNet可以捕获远距离依赖性，同时保持局部结构的分割准确性。

另一个例子是UNet++，它旨在弥合编码器和解码器特征图之间的语义差距。它将深度监督编码器-解码器网络与嵌套的密集跳过路径相互连接，以增强分割准确性。另一个网络，SwinUNet引入了另一种基于Transformer的医学图像分割方法，利用U形编码器-解码器架构和跳过连接来增强局部全局语义特征学习。该模型在传统基于卷积的方法和混合Transformer-卷积技术上表现出优越性能。然而，许多分割工作需要大量的人力在架构修改和超参数调整上进行努力，以适应不同的应用或数据集。

为应对这一挑战，提出了nnUNet框架。nnUNet框架采取了独特的方法，不提出新的网络架构。相反，它重新聚焦于方法、架构搜索和数据预处理步骤，以获得最佳性能。nnUNet策略表明，通过适当的预处理和后处理组合，即使是基本的网络架构也可以在各种医学分割任务中实现最先进的性能。

然而，历史上，用于医学图像分割的深度学习模型，包括nnUNet，都是为特定数据集或应用程序量身定制的，因此很难将单一模型的有效性推广到各种分割任务。虽然nnUNet的出现标志着医学图像分割更加灵活的方法，但分割结果的质量仍然依赖于特定分割任务的大量训练数据。对于具有有限训练数据量的医学图像分割任务，”few-shot”学习解决方案非常重要，它允许基于少量样本训练新模型，并更为实际。

最近出现的SAM，这是一个在1100万张图像和超过10亿个分割Mask（SA-1B训练数据集）上进行训练的模型，显示出在各种图像类别上实现”few-shot”甚至”zero-shot”学习的巨大潜力。然而，最近的研究发现，将SAM模型直接应用于医学图像时，其准确性受到限制，需要额外的微调。

此外，SAM模型除了图像本身外，还需要输入提示（边界框、点等），这阻碍了它在完全自动化的临床工作流中的无缝集成。尽管这一方面对于多功能性是一个优势，但在需要实时或无间断程序的高通量医学场景中可能存在挑战。

最近，基于SAM框架开发了AutoSAM，它可以直接从待分割图像中学习提示，并将学习的提示用于完全自动化的分割。但是，AutoSAM需要从训练数据集中学习新的提示编码器，并可能对”few-shot”情况下的训练数据稀缺性敏感。

受到nnUNet和SAM的优点和缺点的启发，作者引入了nnSAM，这是一种新颖的即插即用解决方案，旨在增强医学图像的分割准确性。nnSAM将SAM的强大特征提取和泛化能力与nnUNet的数据中心自配置能力相结合。通过利用SAM的图像编码器并将其无缝集成到nnUNet的架构中，nnSAM生成了丰富的潜在空间表示，作为提高分割准确性的基础。SAM和nnUNet的融合特别有益于训练数据稀缺的情况，以实现高质量的医学图像分割。

本文的主要贡献总结如下：

• 作者引入了nnSAM，它是Segment Anything Model（SAM）和nnUNet的一种新颖融合。通过结合SAM的强大特征提取能力和nnUNet的自动配置设计，nnSAM确保即使在训练数据非常有限的情况下，也能实现增强的分割质量。

• 作者的全面评估揭示了nnSAM在现有最先进技术上的卓越性能，为医学图像分割提供了一个潜在的新基准。

2. 方法

2.1. 架构概述

所提出的nnSAM框架的架构如图1所示。该模型旨在结合nnUNet和SAM的优势。具体而言，nnSAM由两个并行编码器组成：nnUNet编码器和SAM编码器。SAM编码器是一个预训练的Vision Transformer（ViT）。来自两个编码器的嵌入被连接起来，然后输入到nnUNet的解码器中，以生成最终的分割地图。

此外，SAM编码器被用作即插即用的插件，在训练期间其参数被冻结。因此，只有nnUNet的编码器和解码器的权重在训练期间被更新。

2.2. 自动配置的nnUNet架构

将nnUNet集成到nnSAM中可以实现自动化的网络架构和超参数配置，使其高度适应每个医学图像数据集的独特和特定特征。这种自适应能力始于一个自配置过程，它会自动调整nnUNet编码器的架构，以适应训练数据集的特征，包括医学图像的尺寸、通道数以及分割任务中涉及的类别数。

此外，nnUNet使用自动化的预处理流程，包括对输入数据进行归一化以及应用数据增强技术，如旋转、缩放和弹性变形。这些预处理和增强步骤对提高模型的鲁棒性和准确性至关重要。

除此之外，nnUNet还可以根据数据集固有属性自动选择最有效的损失函数，并根据数据集的特性调整优化器设置。例如，对于数据集内部的检测类别不平衡情况，nnUNet可以自动配置加权损失函数以强调较小类别。这进一步辅以nnUNet的超参数调整过程，包括学习率和批量大小的网格搜索。根据每个特定的训练数据集，nnUNet的架构也会自动调整，通过动态修改参数，如层数和卷积核大小，以实现最佳性能。

这套全面的自动配置功能使nnUNet以及对应的nnSAM架构能够为每个特定的医学图像任务优化其编码器设置，从而提高其效率和准确性。由于nnSAM的层数取决于具体的数据集，图1中作者将编码器层数表示为到，解码器层数表示为到。

2.3. SAM编码器

SAM编码器是一个预训练的Vision Transformer模型，针对广泛的SA-1B分割数据集进行训练。在这个极其大的数据集上训练的SAM编码器在分割任务的领域无关特征提取方面表现出色。然而，它的分割能力高度依赖于提示，无法自动识别分割目标和底层语义。因此，nnSAM只使用SAM编码器来融合其特征提取优势，而将特定于数据集的任务（识别分割兴趣区域）留给nnUNet。

对于输入图像，其中H×W是空间维度，C是通道数，SAM编码器需要输入H×W的尺寸为1024×1024。为了满足这一要求，在nnUNet的预处理之后，作者使用线性插值将其调整为1024×1024。SAM编码器生成一个尺寸为64×64的图像嵌入S。然后，作者将此嵌入S调整为与nnSAM的解码器层D1的尺寸匹配以进行连接（图1）。

为了在速度和准确性之间取得平衡，作者使用MobileSAM，这是SAM的轻量级版本，大小不到原始SAM的1/60，但性能可媲美。MobileSAM是通过蒸馏从原始SAM中获得的，通过这种方式，从原始图像编码器中的知识被转移到轻量级版本中。

3. 实验

表1显示了在不同数量的训练数据样本（4到20）下的模型性能。所提出的nnSAM在DICE和ASD方面在所有训练样本大小上均优于所有其他分割方法。当使用20张标记图像进行训练时，nnSAM的平均DICE得分为93.75％，平均ASD为1.23毫米。被认为是最佳分割模型之一的nnUNet的分割准确性也明显高于其他方法，但略低于nnSAM。其他方法，包括SwinUNet、TransUNet和AutoSAM，精度要低得多，它们的DICE都低于90%，ASD都高于4毫米。最差的性能来自基本UNet，这是可以预期的，因为SwinUNet、TransUNet和AutoSAM都是基于预训练模型的，而UNet是从头开始训练的，最受训练样本不足的影响。

此外，随着训练样本数量逐渐减少，nnSAM相对于其他方法的优势变得更加显著。特别是当只使用4张标记图像进行训练时，对于DICE，nnSAM比第二名的nnUNet高出2.9%，并且比其他方法更高。总的来说，与其他方法相比，nnSAM在分割准确性上提供了更高的性能，特别是在训练数据量有限的情况下。

表2和表3显示了相应标签类别下DICE和ASD的性能。在这两个表中，nnSAM在大多数情况下提供了最佳结果。

作为第二名模型，nnUNet的性能最接近作者的nnSAM，然而，在图3的可视化结果中，当训练样本数量为4时，nnUNet对于Myo类别的分割效果不佳，出现了更多的假阳性。

此外，还存在一些趋势反直觉的异常值。例如，SwinUNet和TransUNet在LA标签上的结果表明，随着样本数据量的增加，ASD变得更大。

根据图2和图3，作者发现随着样本数量的增加，TransUNet和SwinUNet生成了更多的假阳性分割，并且所有这些假阳性都远离真实分割位置，导致ASD的异常结果。

总的来说，UNet和AutoSAM产生了较差的分割结果。当训练数据量有限时，UNet和AutoSAM的LV（Myo）的心肌几乎无法识别。由于AutoSAM严重依赖于自定义训练的编码器来提供自动提示，因此“few-shot”学习场景对于学习SAM分割的准确提示构成了挑战，导致准确性较低。

这些结果表明，nnSAM在仅有少量训练样本的情况下，针对具有挑战性的目标提供了卓越的准确性，这要归功于SAM编码器的强大通用性和nnUNet的自动配置框架的适应能力。

4.讨论

结果表明，nnSAM在医学图像分割方面表现出卓越的性能，尤其是在标记的训练数据有限的few-shot学习场景中。通过将预训练的SAM编码器集成到nnUNet的框架中，nnSAM可以利用SAM强大的特征提取能力，同时受益于nnUNet的自适应架构配置和超参数优化。使用MM-WHS心脏CT数据集进行的评估突出显示了nnSAM的几个关键优势。

首先，nnSAM在所有大小的训练集（4到20个样本）上始终实现了最高的准确性，超越了nnUNet、SwinUNet和TransUNet等最先进的模型。从非常少的示例中生成准确的分割能力使nnSAM在难以获取标记数据且成本高昂的医学应用中具有价值。当标记一些结构的时候，像SwinUNet和TransUNet这样的模型展示出不稳定的结果，更多的训练样本导致结果更差，这表明它们可能对训练数据分布过拟合。相反，nnSAM的分割质量随着添加更多训练数据而持续改善。与nnSAM相比，AutoSAM使用自定义编码器替换提示编码器，使其能够自动生成并提供提示给SAM。

然而，AutoSAM并未针对医学图像的语义分割进行优化，也没有像nnUNet那样强大的预处理和自动配置能力。自从nnUNet出现以来，它已成为大多数医学图像分割任务的最先进技术，代表了传统任务特定语义分割的顶级端到端模型。另一方面，SAM是一种基于提示的分割框架，是一种具有很强通用性的代表性模型。将来自两种不同分割框架的最佳模型结合起来已被证明可以有效地进一步提高医学图像分割的准确性，并设定一个潜在的新基准。

作者当前的研究存在一些限制，需要在未来的工作中解决。

• 首先，作者在一组心脏CT扫描的单个数据集上评估了nnSAM。未来的研究需要在更大更多样化的医学成像数据集上进行测试。

• 其次，目前的nnSAM框架仍需要有限数量的训练数据和标签，未来的工作需要探索只有一个样本（’one-shot’学习）或完全没有标签（’zero-shot’学习）的端到端分割的可能性。

• 此外，作者使用2D切片进行训练和测试；扩展到基于3D体积的分割可能会进一步提高分割的准确性，但目前面临将3D SAM嵌入与3D nnUNet嵌入合并的技术困难。未来的研究需要寻找潜在的解决方案或替代方案。

5.参考

[1]. nnSAM: Plug-and-play Segment Anything Model Improves nnUNet Performance

点击上方卡片，关注「集智书童」公众号