[Paper Review] 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 Education14 citations

TL;DR

nnSAM combines SAM encoder with nnUNet to improve medical image segmentation, achieving higher Dice and lower surface distance, especially in few-shot learning scenarios.

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.

Motivation & Objective

• Motivate improved medical image segmentation when labeled data is scarce.
• Leverage SAM's robust feature extraction with nnUNet's auto-configurable architecture.
• Enable few-shot learning by fusing pretrained Vision Transformer features with dataset-tailored learning.
• Demonstrate superior performance across varying training sample sizes.

Proposed method

• Integrate a frozen SAM encoder with the nnUNet decoder in a two-branch architecture.
• Concatenate embeddings from nnUNet encoder and SAM encoder before decoding segmentation.
• Use MobileSAM as a lightweight, fast encoder to balance speed and accuracy.
• Employ nnUNet's auto-configuration to adapt architecture, preprocessing, and hyperparameters to each dataset.

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

Experimental results

Research questions

• RQ1 Does nnSAM improve segmentation accuracy over state-of-the-art models across different training sample sizes?
• RQ2 Can the SAM encoder provide robust feature extraction that benefits nnUNet in medical image segmentation?
• RQ3 Is few-shot performance of nnSAM superior to existing methods on cardiac CT segmentation?

Key findings

• nnSAM consistently achieves the highest Dice scores across training sizes (4–20) compared to baselines.
• With 20 labeled images, nnSAM attains Dice of 93.75% and ASD of 1.23 mm on average.
• nnUNet is strong but slightly outperformed by nnSAM, while AutoSAM, SwinUNet, TransUNet, and UNet lag behind, especially with fewer samples.
• nnSAM's advantage grows as training data decreases, showing robustness in few-shot settings.
• The combination leverages SAM's general feature extraction with nnUNet's auto-configured learning to boost performance.

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