[논문 리뷰] nnSAM: Plug-and-play Segment Anything Model Improves nnUNet Performance

Yunxiang Li, Bowen Jing외 3명|arXiv (Cornell University)|2023. 09. 29.

Artificial Intelligence in Healthcare and Education인용 수 14

한 줄 요약

nnSAM은 SAM 인코더를 nnUNet와 결합하여 의료 영상 분할 성능을 향상시키고, Dice를 더 높이고 표면 거리 ASD를 더 낮추며, 특히 소샷 학습 시나리오에서 두드러진다.

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의 자동 구성 가능한 아키텍처를 활용한다.
• 사전 학습된 비전 트랜스포머 특징을 데이터셋에 맞춘 학습과 융합하여 소샷 학습을 가능하게 한다.
• 다양한 학습 샘플 크기에서 우수한 성능을 입증한다.

제안 방법

• 동결된 SAM 인코더를 nnUNet 디코더와 결합한 이중 분기 아키텍처를 구현한다.
• Segmentation 디코딩하기 전에 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.