[論文レビュー] 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 の特徴をデータセットに合わせた学習と融合して few-shot 学習を可能にする。
• 異なる学習サンプルサイズにわたって優れた性能を示す。

提案手法

• 凍結された SAM エンコーダを nnUNet デコーダと二系統アーキテクチャで統合する。
• デコード前に nnUNet エンコーダと SAM エンコーダからの埋め込みを連結する。
• モバイル SAM を軽量で高速なエンコーダとして使用し、速度と精度のバランスを取る。
• 各データセットに対してアーキテクチャ、前処理、ハイパーパラメータを自動的に適応させる 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 エンコーダが医用画像セグメンテーションに有益な堅牢な特徴抽出を提供するか。
• RQ3 心臓 CT セグメンテーションにおいて nnSAM の few-shot パフォーマンスは既存手法を凌駕するか。

主な発見

方法	DICE 4	DICE 8	DICE 12	DICE 16	DICE 20	ASD 4	ASD 8	ASD 12	ASD 16	ASD 20
UNet	59.31	62.46	66.86	71.86	76.52	19.54	20.82	17.86	19.79	18.15
SwinUNet	81.24	80.82	83.15	84.37	86.88	4.79	5.29	3.61	4.12	4.17
TransUNet	81.23	82.34	84.82	87.05	87.11	4.18	8.79	9.77	11.03	11.99
AutoSAM	65.10	65.89	67.63	77.55	78.53	16.55	19.60	16.98	16.73	15.92
nnUNet	81.77	84.45	88.36	92.35	93.15	6.97	4.90	3.15	1.56	1.40
nnSAM	84.67	86.36	90.74	93.20	93.75	3.87	3.29	2.18	1.43	1.23

• nnSAM は訓練サイズが 4–20 の間で一貫して最高の Dice スコアを達成する。
• 20 枚のラベル付き画像で、nnSAM は Dice が 93.75%、ASD が 1.23 mm の平均を達成する。
• nnUNet は強力だが nnSAM にわずかに及ばず、AutoSAM、SwinUNet、TransUNet、UNet は、特にサンプル数が少ない場合に後れを取る。
• nnSAM の利点は訓練データが減少するにつれて拡大し、few-shot 設定での堅牢性を示す。
• この組み合わせは SAM の一般的な特徴抽出と nnUNet の自動構成学習を活用して性能を向上させる。

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