SAM3D: Segment Anything Model in Volumetric Medical Images

Volumetric segmentation is crucial in medical image analysis, finding applications in pathology diagnosis, surgical planning, and computer-aided diagnosis. Volumetric medical images like CT, MRI, OCT, and DBT offer a 3D view of anatomical structures. Segmentation identifies regions of interest for better interpretation.

Deep learning, particularly UNet [1] and variants [2, 3, 4], made strides in 3D medical segmentation but faced limitations. Transformer-based models like Vision Transformer (ViT) [5] and Swin-UNet [4] showed promise in capturing long-range relationships. Combining CNNs and Transformers in models like TransUNet [3], UNETR [6], and HiFormer [7], yielded promising results. However, these models prioritize precision, leading to increased complexity and training time. Leveraging pretrained models offers an alternative. SAM, a transformer-based model pretrained on large-scale datasets, has shown generalizability in segmentation tasks. SAM-based models for medical images piqued interest.

This work introduces SAM3D, an architecture for volumetric medical segmentation, combining the SAM encoder and a lightweight 3D CNN decoder. Unlike traditional slice-by-slice processing, SAM3D extracts features across the entire volume, improving segmentation while maintaining simplicity and computational efficiency. Contributions include applying the SAM encoder to process 3D volumes, designing SAM3D for effective 3D medical segmentation, and validating its performance on various datasets, such as ACDC [8], Synapse [9], MSD BraTS [10], and MSD Lung [10]. SAM3D demonstrates competitive results, marking a novel approach to 3D volumetric imaging.

Segmentation Methods using CNNs and Transformer: Various methods leverage a combination of CNNs and Transformer architectures for segmentation tasks. TransUNet [3] integrates CNNs and Transformer within a U-shaped architecture to capture both local and global information. Swin-Unet [4] replaces U-Net’s convolutional blocks with Swin Transformer blocks. nnUNet [2] introduces a self-adapting framework for 2D and 3D medical segmentation. MISSFormer [11] enhances hierarchical feature representation using Enhanced Transformer Blocks. TransDeepLab [12] combines Swin Transformer blocks with ASPP module and cross-contextual attention. HiFormer [7] introduces the Double-Level Fusion (DLF) module. UNETR [6] encodes 3D input patches with Transformers and combines feature extraction with a CNNs-based decoder. Swin UNETR [13] enhances UNETR with Swin Transformer blocks. nnFormer [14] interleaves CNNs and Transformer blocks with feature pyramids. UNETR++ [15] introduces the efficient paired-attention (EPA) module.

In contrast to existing approaches that involve fine-tuning SAM and handling 3D images as sets of 2D slices, and unlike the conventional CNNs/Transformer-based methods that typically require large model designs, our proposed SAM3D effectively and efficiently harnesses SAM’s capabilities for 3D medical segmentation. It does so without the need for large model architectures or depending on slice-by-slice predictions. This approach enhances the model’s ability to perceive anatomical structures and capture global information.

In this section, we introduce our model, SAM3D, and explain the rationale behind its simple design. Our goal is to leverage SAM without the need for extensive parameter retraining or complex task-specific modules.

In contrast to SAMed [17] and MedSAM [18], where all three components of SAM are fine-tuned, our approach involves freezing SAM’s image encoder and training a new lightweight 3D decoder. SAM3D leverages SAM by initially processing images slice by slice and then incorporating a lightweight 3D decoder to capture depth-wise relationships between slices. The overall architecture of SAM3D is depicted in Figure 1 and can be summarized as follows: a volumetric input $I\in\mathbb{R}^{H\times W\times D}$ is divided into $D$ 2D slices, each of dimension $H\times W$. We duplicate each channel three times to generate the slices that have dimension of $H\times W\times 3$. The pretrained SAM encoder processes these slices, generating 3D slice embeddings denoted as $F$. The depth-wise relationships among these slice embeddings are effectively captured by our proposed 3D decoder. Additionally, we remove the prompt encoder from SAM to ensure that feature extraction remains uninhibited across different modalities.

Additionally, we employ the deep supervision technique for multiple decoding stages. Specifically, the output features of each decoding stage pass through a segmentation block, consisting of one 3 x 3 x 3 and one 1 x 1 x 1 convolution layer, to generate predictions for one typical stage. To calculate the loss value for one typical stage, we down-sample the ground truth to match the prediction resolution. Consequently, the final loss can be defined as follows:

here, L is set to 3, representing the number of decoder layers. $\alpha_{l}$ signifies the hyperparameter controlling the contribution of different resolutions to the final loss function. In practice, we set $\alpha_{2}=\frac{\alpha_{1}}{2}$ and $\alpha_{3}=\frac{\alpha_{1}}{4}$ with all $\alpha$ hyperparameters normalized to 1.

A. Datasets and Evaluation Metrics.

Synapse comprises eight abdominal organs in a large dataset and the performance comparison is shown in Table 1. Among the models evaluated, UNETR++ (a Transformer-based model) achieved the best results with 42.9M parameters, while nnFormer ranked second with 150.5M parameters. Notably, SAMed_s distinguishes itself by achieving impressive results with a modest 6.32M parameters and a DSC of 77.78%. SAMed_s shares a similar architecture with our SAM3D, fine-tuned from SAM, but differs in processing methods. SAMed_s employs a straightforward slice-by-slice approach, while SAM3D considers depth-wise information. Despite this difference, both models are efficient in parameter usage. SAMed_s requires 6.32M parameters, whereas SAM3D excels with just 1.88M parameters. Furthermore, SAM3D achieves a DSC score exceeding 1.78%, demonstrating superior performance compared to SAM-based methods with lightweight models.

While SAMed is exclusive to the Synapse dataset, our SAM3D can be evaluated on a variety of other datasets, including Cardiac, Brain Tumor, and Lung. In Table 2 and 3, it is evident that SAM3D competes favorably with SOTA CNNs/Transformer-based networks on the Cardiac ACDC and Lung datasets. For instance, SAM3D surpasses TransUnet’s performance on the ACDC dataset with a 0.41% increase in DSC while utilizing less than 50$\times$ the number of parameters. Table 4 further illustrates SAM3D’s competitiveness with other leading models on the Brain Tumor Brats dataset, despite its significantly lower parameter count. For example, SAM3D achieves a 1.8% DSC improvement compared to UNETR, while requiring less than 20$\times$ the number of params. It is worth noting that the MRI scans in Brats contain four modalities, which explains SAM3D’s parameter count being four times that of other single-modality models.

Fig. 3 visually presents samples from the Synapse dataset. In this illustration, we compare our approach (in the third column) with the outcomes obtained from SAMed and SAMed_s, which represent SOTA in SAM-based methods for volumetric medical image segmentation. Despite the reduced trainable params, SAM3D exhibits superior segmentation performance compared to the other two methods.

In this study, we introduce SAM3D, an efficient and simple SAM-based model tailored for volumetric medical image segmentation. Our approach harnesses the capabilities of a SAM pre-trained encoder coupled with a lightweight 3D decoder. Through extensive experimentation, we have established that SAM3D competes effectively with current SOTA 3D neural networks and Transformer-based models while demanding significantly fewer parameters (50$\times$ fewer). Furthermore, SAM3D outperforms other lightweight networks in the context of volumetric segmentation. As SAM has already made a substantial impact on natural image segmentation, our research extends its potential to the domain of medical image segmentation. We anticipate that this work will serve as an inspiration for future researchers, fostering advancements in the field of medical segmentation

Additionally, our simple decoder leaves room for developing a more complex architecture, which could potentially enhance the model’s performance. This presents a promising avenue for further research and development.

Nhat-Tan Bui and Ngan Le are supported by the National Science Foundation (NSF) under Award No OIA-1946391 RII Track-1, NSF 1920920 RII Track 2 FEC, NSF 2223793 EFRI BRAID, NSF 2119691 AI SUSTEIN, NSF 2236302. Minh-Triet Tran is sponsored by Vietnam National University Ho Chi Minh City (VNU-HCM) under grant number DS2020-42-01. Dinh-Hieu Hoang is funded by Vingroup Joint Stock Company and supported by the Domestic Master/ PhD Scholarship Programme of Vingroup Innovation Foundation (VINIF), Vingroup Big Data Institute (VINBIGDATA), code VINIF.2022.ThS.JVN.04.