[论文解读] nnSAM: Plug-and-play Segment Anything Model Improves nnUNet Performance

Yunxiang Li, Bowen Jing+3 more|arXiv (Cornell University)|Sep 29, 2023

Artificial Intelligence in Healthcare and Education被引用 14

一句话总结

nnSAM 将 SAM 编码器与 nnUNet 结合，以提高医学图像分割，在 Dice 更高、表面距离更低，特别是在小样本学习情景。

ABSTRACT

Automatic segmentation of medical images is crucial in modern clinical workflows. The Segment Anything Model (SAM) has emerged as a versatile tool for image segmentation without specific domain training, but it requires human prompts and may have limitations in specific domains. Traditional models like nnUNet perform automatic segmentation during inference and are effective in specific domains but need extensive domain-specific training. To combine the strengths of foundational and domain-specific models, we propose nnSAM, integrating SAM's robust feature extraction with nnUNet's automatic configuration to enhance segmentation accuracy on small datasets. Our nnSAM model optimizes two main approaches: leveraging SAM's feature extraction and nnUNet's domain-specific adaptation, and incorporating a boundary shape supervision loss function based on level set functions and curvature calculations to learn anatomical shape priors from limited data. We evaluated nnSAM on four segmentation tasks: brain white matter, liver, lung, and heart segmentation. Our method outperformed others, achieving the highest DICE score of 82.77% and the lowest ASD of 1.14 mm in brain white matter segmentation with 20 training samples, compared to nnUNet's DICE score of 79.25% and ASD of 1.36 mm. A sample size study highlighted nnSAM's advantage with fewer training samples. Our results demonstrate significant improvements in segmentation performance with nnSAM, showcasing its potential for small-sample learning in medical image segmentation.

研究动机与目标

• 在标注数据稀缺时推动改进的医学图像分割。
• 利用 SAM 的鲁棒特征提取与 nnUNet 的自动配置架构。
• 通过将预训练的 Vision Transformer 特征与数据集定制学习融合，实现少样本学习。
• 在不同训练样本规模下展示更优性能。

提出的方法

• 在两分支架构中将冻结的 SAM 编码器与 nnUNet 解码器整合。
• 在解码分割前，将 nnUNet 编码器和 SAM 编码器的嵌入拼接。
• 使用 MobileSAM 作为轻量级、快速的编码器，以在速度和准确性之间取得平衡。
• 利用 nnUNet 的自配置来使架构、预处理和超参数适应每个数据集。

Figure 1: The architecture of nnSAM, which integrates nnUNet’s encoder with the pretrained SAM encoder. The correspondingly concatenated embeddings are input into nnUNet’s decoder to output the final segmentation. A cardiac sub-structure segmentation example is presented. (LV: left ventricle; RV: ri

实验结果

研究问题

• RQ1 nnSAM 是否在不同训练样本规模下相较于最先进模型提升分割准确性？
• RQ2 SAM 编码器是否可提供有益于 nnUNet 的鲁棒特征提取以用于医学图像分割？
• RQ3 在心脏 CT 分割任务中，nnSAM 的少样本表现是否优于现有方法？

主要发现

• 与基线相比，nnSAM 在不同训练规模（4–20）下始终取得最高的 Dice 分数。
• 在有 20 张标注图像时，nnSAM 的 Dice 达到 93.75%，ASD 为 1.23 mm。
• nnUNet 较强但略落后于 nnSAM，而 AutoSAM、SwinUNet、TransUNet 与 UNet 则落后，尤其是在样本较少时。
• 随着训练数据减少，nnSAM 的优势进一步扩大，显示出少样本设置中的鲁棒性。
• 该组合通过将 SAM 的通用特征提取与 nnUNet 的自配置学习结合，提升了性能。

Figure 2: Example 1 of segmentation visualization results for different methods on different numbers of training samples.