Source-linked AI summary
SAM-Med3D: Towards General-purpose Segmentation Models for Volumetric Medical Images
Haoyu Wang, Sizheng Guo, Jin Ye, Zhongying Deng, Junlong Cheng, Tianbin Li, Jianpin Chen, Yanzhou Su, Ziyan Huang, Yiqing Shen, Bin Fu, Shaoting Zhang, Junjun He, Yu Qiao
TL;DR
Task-specific volumetric segmentation models have limited generalization across structures and modalities, motivating a general-purpose alternative. The paper introduces SAM-Med3D, a fully learnable 3D promptable model trained on SA-Med3D-140K and evaluated across diverse tasks, reporting broad performance and transfer potential. A key scope boundary is that 58% of validation cases come from sources unseen during training, while some official benchmark training data are included in training.
Problem
Existing volumetric medical segmentation models specialize in particular organs or lesions and therefore have limited generalization across anatomical structures and modalities.
Method
SAM-Med3D is a fully learnable 3D promptable segmentation model trained in two stages on the SA-Med3D-140K volumetric dataset.
Results
Across 16 public datasets, SAM-Med3D demonstrates competitive interaction-efficient segmentation, broad performance on seen and unseen targets, and promising downstream transfer, including a 60.12% improvement over SAM with one point per volume.
Takeaways & Limitations
SAM-Med3D shows potential as a general-purpose pre-trained model for promptable and semantic segmentation of diverse volumetric medical targets.
Takeaways & Limitations
The validation design includes benchmark training data in the model’s training set, although 58% of validation cases come from sources unseen during training.
Abstract
from arXiv · showhide
Existing volumetric medical image segmentation models are typically task-specific, excelling at specific target but struggling to generalize across anatomical structures or modalities. This limitation restricts their broader clinical use. In this paper, we introduce SAM-Med3D for general-purpose segmentation on volumetric medical images. Given only a few 3D prompt points, SAM-Med3D can accurately segment diverse anatomical structures and lesions across various modalities. To achieve this, we gather and process a large-scale 3D medical image dataset, SA-Med3D-140K, from a blend of public sources and licensed private datasets. This dataset includes 22K 3D images and 143K corresponding 3D masks. Then SAM-Med3D, a promptable segmentation model characterized by the fully learnable 3D structure, is trained on this dataset using a two-stage procedure and exhibits impressive performance on both seen and unseen segmentation targets. We comprehensively evaluate SAM-Med3D on 16 datasets covering diverse medical scenarios, including different anatomical structures, modalities, targets, and zero-shot transferability to new/unseen tasks. The evaluation shows the efficiency and efficacy of SAM-Med3D, as well as its promising application to diverse downstream tasks as a pre-trained model. Our approach demonstrates that substantial medical resources can be utilized to develop a general-purpose medical AI for various potential applications. Our dataset, code, and models are available at https://github.com/uni-medical/SAM-Med3D.
1 Introduction
Volumetric medical segmentation remains dominated by task-specific models, while direct or slice-wise adaptations of 2D SAM inadequately capture 3D information. SAM-Med3D addresses this gap with a fully 3D, promptable model and broad evaluation.
- Motivation: Task-specific volumetric models generalize poorly across anatomical structures and modalities, requiring separate models for different clinical applications.This specialization increases development costs and wastes resources.
- Motivation: Directly applying 2D SAM to volumetric medical images is infeasible because of the medical-natural image domain gap and missing 3D spatial modeling.Slice-by-slice aggregation additionally fails to capture inter-slice correlations and consistent 3D predictions.
- Approach: SAM-Med3D is a fully learnable 3D promptable segmentation model trained with a two-stage procedure on large-scale volumetric medical data.It uses a few 3D prompt points to segment diverse structures and lesions across modalities.
- Resources: SA-Med3D-140K contains 22K medical images, 143K masks, and 245 categories assembled from public and private datasets.The dataset is described as significantly larger than existing public medical image segmentation datasets.
- Results: SAM-Med3D shows competitive performance with fewer prompts, broad segmentation across targets and modalities, and promising transfer to downstream tasks.The evaluation covers 16 public datasets and includes seen and unseen tasks.
2 Related Works
Prior work has mainly explored adapting SAM to volumetric medical imaging or evaluating its medical segmentation performance, with reported limitations in general effectiveness.
- SAM for Volumetric Medical Images: Prior volumetric-medical efforts primarily fine-tune SAM or incorporate it into medical image-analysis pipelines.MedLSAM, for example, uses a two-stage model with a localization model and SAM.
- Evaluation of SAM in Medical Imaging: Evaluations across medical datasets generally report SAM performance lagging behind state-of-the-art methods or remaining subpar.Reported behavior varies with prompt type, target connectedness, and anatomical task.
3 SA-Med3D-140K Dataset
SA-Med3D-140K is a large, heterogeneous volumetric segmentation dataset assembled from public and private sources, standardized through quality-control processing and split to test generalization.
- Data Pre-processing: The preprocessing pipeline removes targets below 1 cm3 or 1.5 cm in any dimension and discards volumes whose background exceeds 99%.It also converts multiclass masks to one-hot formats and reduces noise using connected-domain processing.
- Data Splitting: The dataset is partitioned into 131K training masks and 12K validation masks, with training data from 58 public datasets plus all private data.The validation set contains 16 public benchmark datasets.
- Data Splitting: Among validation cases, 58% come from data sources unseen during training, while all ultrasound data are assigned to validation for cross-modality evaluation.This split is designed to assess generalization across data sources and modalities.
4 Method
SAM-Med3D replaces frozen 2D adaptations with a fully 3D architecture trained from scratch, using 3D prompts and a two-stage training strategy to improve volumetric modeling and generalization.
- Architecture Selection: The study compares 3D adapters, full fine-tuning, and training from scratch using seen AMOS and unseen TotalSegmentator targets.Dice scores are averaged across one- and ten-point prompt settings.
- Architecture: SAM-Med3D uses a 3D image encoder, 3D prompt encoder, and 3D mask decoder with 3D positional encoding and 3D layers.The architecture directly processes 3D images and 3D prompts.
- Architecture Selection: Training from scratch provides the best trade-off because adapters perform suboptimally on seen targets while full fine-tuning struggles on unseen targets.This motivates selecting a fully 3D architecture trained on 3D medical data.
- Training: The two-stage procedure first pre-trains on all 131K training masks for 800 epochs, then fine-tunes on approximately 75K high-quality masks.The major difference between stages is the training data.
- Inference: During evaluation, SAM and SAM-Med2D process slices individually, whereas SAM-Med3D uses patch-based inference for 3D volumes over 128.The patch procedure crops around the initial point and extends inference with 50% overlapping sliding windows when needed.
5 Experiments
SAM-Med3D achieves strong overall segmentation performance while substantially reducing inference time and interaction requirements across a 16-dataset validation set.
- 60.12% improvement in overall Dice scores over the comparison methods, with SAM-Med3D consistently exceeding SAM and SAM-Med2D across prompt counts.
- 1% to 26% of SAM’s inference time is required by SAM-Med3D for targets of different sizes.The fully 3D architecture enables greater throughput and reduces interaction burden.
- The validation set compares methods across 16 datasets, with N ranging from 10 to 200 slices containing the target.The table also represents interaction time as τ, usually exceeding 1 second.
- When N surpasses 20, SAM costs more inference time than SAM-Med3D even when human interaction time is excluded.N denotes the number of slices containing the target.
- The comparison includes task-specific and general-purpose methods on six public benchmarks, with marked items denoting unseen data sources or modalities.
Comparison with Task-specific Models
SAM-Med3D remains competitive with task-specific models across six public benchmarks and can outperform all other competitors on most benchmarks with 10 prompt points.
- SAM-Med3D achieved competitive Dice scores across all six public benchmarks compared with task-specific models.The task-specific comparisons used nnU-Net and UNETR trained separately on each benchmark’s training set.
- With 10 points, SAM-Med3D outperformed all other competitors on four benchmarks, including unseen data sources and modalities.
- On unseen Ultrasound, increasing prompt points produced more significant performance improvements for SAM-Med3D.
- SAM-Med2D surpassed state-of-the-art 3D task-specific models partly because its training data included 2D Ultrasound slices.
- Table 5 organizes validation results by anatomical structure and lesion, using Abd&Tho for Abdominal and Thorax targets.N denotes the count of slices containing the target, ranging from 10 to 200.
Evaluation on Different Anatomical Structures
Across anatomical structures and lesions, SAM-Med3D is robust with sparse prompts, while brain targets remain challenging and SAM performs poorly with one point per slice.
- With one point, SAM-Med3D identifies anatomical structures and lesions more effectively than SAM, which often performs poorly even with 3N or 5N points.The difficulty is especially pronounced for cardiac and muscle regions with unclear boundaries.
- Brain structures and lesions were more challenging for all evaluated models than other seen anatomical targets.
- Figure 4 compares Dice scores across modalities and prompt counts, including unseen organs and lesions marked with symbols.
- With sufficient prompts, SAM-Med2D becomes competitive with SAM-Med3D on brain structures, but extra prompts do not similarly close the lesion gap.The authors associate this difference with lesions’ greater dependence on prompts.
Evaluation on Different Modalities
SAM-Med3D performs strongly on CT and MR and remains comparable on unseen Ultrasound with sufficient prompts; transferability is also evaluated through semantic and promptable segmentation.
- SAM-Med3D performs better on the mainstream CT and MR modalities and achieves comparable performance on unseen Ultrasound with sufficient prompts.
- SAM-Med3D using one point outperforms SAM-Med2D with N points on 45 of 49 targets, with improvements up to +68.2%.The targets comprise 44 major organs and five kinds of lesions.
- Transferability for semantic segmentation is assessed by comparing UNETR with and without the SAM-Med3D pre-trained ViT encoder.
- Transferability for promptable segmentation is assessed across first-stage pre-training and multiple second-stage fine-tuning datasets.The final comparison uses high-quality data selected from 44 datasets in SA-Med3D-140K.
Evaluation on Transferability
SAM-Med3D is evaluated as a transferable model for semantic and promptable segmentation across seen and unseen datasets. The results support its potential as a pre-trained feature extractor while showing that fine-tuning choices trade target-specific expertise against general-purpose performance.
- Semantic segmentation: UNETR with SAM-Med3D’s pre-trained ViT encoder is evaluated on AMOS, Totalsegmentator, CAS2023, and SEG.A.2023.The latter two are unseen datasets from the MICCAI 2023 Challenge.
- SAM-Med3D demonstrates potential as a powerful pre-trained model for multiple downstream segmentation tasks.
- Promptable segmentation: Fine-tuning on specific datasets improves expertise for targeted entities but decreases performance on non-target entities.Carefully selected, high-quality data further strengthens general-purpose segmentation.
5.2 Qualitative Evaluation
Qualitative visualizations compare SAM, SAM-Med2D, and SAM-Med3D across point counts, anatomical structures, modalities, and views. SAM-Med3D requires fewer prompts and produces more consistent predictions between slices.
- The visual comparison spans different point numbers, anatomical structures, and modalities.
- SAM-Med3D requires fewer prompts than the other methods.
- SAM-Med3D shows better inter-slice consistency than all other methods.
6 Conclusion
SAM-Med3D is a general-purpose promptable model trained as a fully learnable 3D network on large-scale 3D medical data. Evaluations across 16 public datasets verify its efficiency and efficacy, including a 60.12% improvement over SAM with one point per volume, while indicating downstream pre-training potential.
- SAM-Med3D segments seen and unseen anatomical structures and modalities using only a few prompt points.
- The model is trained as a fully learnable 3D network on a large-scale 3D medical dataset in two stages.
- 60.12% improvement over SAM with one point per volume is reported across evaluations on 16 public datasets.
- SAM-Med3D shows promising potential as a pre-trained transformer model for multiple downstream segmentation tasks.