Leveraging SAM for Single-Source Domain Generalization in Medical Image Segmentation

The Segment Anything Model (SAM) proposed by Kirillov et al. (2023) stands out as a groundbreaking foundational model that introduces the concept of promptability, enabling robust generalization across numerous domains. Trained on the large SA-1B dataset, SAM has demonstrated exceptional performance in zero-shot segmentation of natural images, showcasing its versatility and adaptability. However, when it comes to medical image segmentation, SAM’s original model encounters limitations due to its lack of domain-specific knowledge in the realm of medical imaging as demonstrated in Zhang & Jiao (2023).

To address this challenge, previous works have explored various strategies to enhance SAM’s performance in medical image segmentation tasks. One such approach involves fine-tuning SAM using medical images, as demonstrated in projects like MedSAM(Ma & Wang, 2023) and MSA(Wu et al., 2023). Some of the work like Skinsam(Hu et al., 2023) and SAMUS(Lin et al., 2023) applied SAM to more specific segmentation tasks including ultrasound, 3d images, or skin cancer. They aim to impart SAM with a better understanding of medical images’ unique characteristics and intricacies, enabling it to produce more accurate and reliable segmentation results. Another promising solution is the utilization of auto-prompting techniques, which have shown promising improvements in SAM’s performance across different domains in projects like AutoSAM(Shaharabany et al., 2023) and AutoSAM Adapter(Li et al., 2023).

Despite these advancements, it is important to note that the current performance achieved by these approaches may not be sufficient for clinical applications that demand higher levels of accuracy and reliability. Additionally, it is worth highlighting that there is no noticeable research focused on the broader aspect of generalization across domains in the context of SAM. While SAM has demonstrated impressive generalization capabilities in various scenarios, there remains a need to explore and develop techniques that enable SAM to effectively generalize across diverse domains, including medical imaging. By addressing this research gap, we can unlock SAM’s full potential in tackling the challenges associated with domain generalization and further elevate its applicability in real-world medical image segmentation tasks.

Image segmentation is critical for computer-aided diagnosis and treatment planning. However, the effectiveness of deep network models trained for specific domains can be hindered when confronted with new, unseen target domains due to unpredictable domain shifts. This challenge has spurred significant interest in the field of domain generalization (DG), which aims to develop models that can generalize well across diverse domains, even those that have not been encountered during training. Several recent works have achieved good performance in medical image segmentation, by applying data augmentation and learning domain-invariant features. So far, some works have applied domain generalization to medical image segmentation. For instance, Liu et al. (2021) proposed a method based on meta-learning by encouraging the shape compactness and shape smoothness for prostate segmentation. Wang et al. (2020) developed a Domain Knowledge Pool to learn and memorize the prior information extracted from multi-source domains, which helps to augment the original image. Hu et al. (2022) designed a domain adaptive convolution module and content adaptive convolution module and incorporated both into an encoder-decoder backbone. Zhou et al. (2022) introduced a domain-specific restoration module for regularization and a Random Amplitude Mixup module with low-level frequency knowledge.

Among the various scenarios in DG, single-source domain generalization presents a distinct challenge. Unlike multi-source domain generalization, which benefits from multiple source domains, single-source DG aims to train a predictive model that is both robust and generalizable using only a single source domain. This poses an extreme lack of data, making it more difficult to achieve reliable and accurate segmentation results. Recent research efforts have primarily focused on addressing this challenge through innovative techniques such as data augmentation and the integration of multi-source DG techniques like adversarial training and contrastive learning.Ouyang et al. (2022) proposed a simple causality-inspired data augmentation approach to expose a segmentation model to synthesized domain-shifted training examples. Su et al. (2023) proposed a location-scale augmentation strategy, that engages the inherent class-level information and enjoys a general form of augmentation. Xu et al. (2022) proposed a mutual information regularizer to enforce the semantic consistency between images from the synthetic domains and estimated the consistency by patch-level contrastive learning. However, the outcomes obtained thus far have been considerably unsatisfying.

The overview of our methods architecture is shown in Figure 1. Our approach compromises two primary stages. In the first stage, we utilize a traditional segmentation model to generate coarse prediction masks. We then apply a mask-filtering module to produce refined bounding boxes, which are used in the second stage to fine-tune SAM. Each stage and its main components are described in detail in the following sections.

In the first stage, we train a Resnet backbone as a segmentation network to produce coarse prediction masks. The segmentation network is given input of $(I,M_{GT})$ and outputs $M_{C}$, where $I,M_{GT}$ and $M_{C}$ denote the source image, ground truth masks, and coarse predictions respectively. Resnet was selected as the backbone network due to its widespread adoption and established reputation in the domain of medical image segmentation.

During the training phase, the network is exclusively trained using data from a single domain. In the subsequent testing phase, the trained network is directly employed on unfamiliar target domains without any additional modifications.

To ensure the production of high-quality bounding boxes, we acknowledge the presence of considerable noise within the coarse masks initially predicted by Resnet. This noise has the potential to significantly degrade the quality of the resulting bounding boxes, which would lead to unpromising mask prediction results from SAM.

To address this issue, we introduce a mask-filtering module that effectively eliminates noise and retains the largest continuous mask, thereby improving the overall quality of the generated bounding boxes. The underlying concept of this mask-filtering module is straightforward yet impactful. Despite the presence of considerable noise in the output of the segmentation network, it is highly probable that the network accurately identifies the largest region corresponding to the target object. Leveraging this insight, we employ a breadth-first search (BFS) algorithm to identify the largest connected component of the mask, discarding all other components. This refinement process yields improved predictions and facilitates the generation of more precise bounding boxes.

While implementing this module, we observed that the computational speed of the algorithm is significantly influenced by the size of the input image. Considering that the main objective of this step is to eliminate smaller components in the coarse prediction masks, preserving the original size is unnecessary. We devised a strategy to expedite this process by resizing the initial coarse prediction masks $M_{C}$ to a smaller dimension $M_{C}^{\prime}$, (e.g. from $512\times 512$ to $128\times 128$), before executing the algorithm. Subsequently, to restore the original size of the masks, we upscale the corresponding bounding box $B_{R}$’s parameters by the corresponding factor (e.g. 4 for $512\times 512$ to $128\times 128$). By adopting this approach, we achieve a balance between computational efficiency and the desired accuracy of the bounding boxes.

Following the acquisition of refined bounding boxes, we employed a fine-tuning strategy on SAM, similar to the approach used in MedSAM. During this process, we exclusively fine-tuned SAM’s lightweight mask decoder, while keeping the parameters of the image encoder and prompt encoder fixed throughout training.

However, it is worth noting that MedSAM resizes all images to a standardized dimension of $1024\times 1024$ to accommodate SAM’s embedding module. Unfortunately, this resizing operation can incur unnecessary computational overhead and impede training speed in certain situations.

In contrast, to fully exploit the capabilities of SAM’s robust image encoder and its multi-box prediction functionality, we adopted a straightforward concatenation method. Specifically, we merged four individual $512\times 512$ images to create a larger composite image with dimensions of $1024\times 1024$. Consequently, the associated bounding boxes $B_{R}$ and ground truth masks $M_{GT}$ were adjusted accordingly. This merging approach enables us to harness the benefits of SAM’s pre-trained embedding modules within its image encoder, while simultaneously expediting the training and testing processes.

Similar to section 3.2, during the training phase, SAM is only fine-tuned using only the data from the same single source domain, while tested on all other target domains.

Dataset    In this study, we assess the performance of our methodology using the publicly available Prostate dataset, which is a diverse collection of T2-weighted MRI prostate images and corresponding masks, spanning six distinct domains.

Evaluation Metric    We employ a widely adopted evaluation metric in medical image segmentation domains - the Dice coefficient (Dice) metric. The Dice coefficient provides a quantitative measure of the quality of segmentation results, with higher values indicating better segmentation performance.

Implementation Details    To mitigate the impact of noise in the coarse predictions, we established a higher confidence threshold of $\theta_{1}=0.75$ for the mask predictions generated by the Resnet backbone. Additionally, considering the superior segmentation performance and increased robustness of the SAM model, we set a lower confidence threshold of $\theta_{2}=0.5$ for SAM’s predictions. In the training process, we use an Adam solver with a base learning rate of 0.0001 and fine-tune SAM for 200 epochs on each domain of the prostate dataset with a batch size of 16.

We report the performance of our approach on the Prostate dataset in cross-site evaluation, as listed in Table 1. Our approach accomplishes the highest Dice score of $79.54\%$, outperforming all previous methods by a large margin. Compared with the former best ACSDG results, our result is $8.12\%$ higher (absolute) and $11.37\%$ better (relative).

It is also worth noting that although the previous state-of-the-art methods demonstrated impressive performance in certain domains, they exhibited limitations and shortcomings when applied to other domains. In contrast, our method achieves the best performance in all 6 domains, demonstrating a stronger domain generalization ability.

We conducted a series of ablation experiments on the Prostate dataset to analyze the design and parameter selection in our method.

Ablation on the confidence threshold of SAM    We used different threshold parameters $\theta_{2}$ to filter the final prediction of SAM. The performance varies as shown in Table 2. SAM showcases exceptional precision in mask prediction, achieving highly accurate overall prediction results even when the confidence score exceeds only 0.5. However, when a higher confidence threshold is employed, the model might mistakenly eliminate parts that are, in fact, correct predictions. If not specified, the confidence threshold $\theta_{2}$ is set to 0.5 as the default setting.

Ablation on different modules    The quantitative findings of this section are presented in Table 3. Initially, we fine-tuned SAM using ground truth bounding boxes directly obtained from the ground truth masks, representing the upper limit of our method. This approach yielded outstanding results, with an average Dice score surpassing $90\%$. However, when we employed the most straightforward method of using the full image size as the bounding box input, SAM’s performance noticeably decreased to approximately $74.41\%$. This decline highlights the significant influence of bounding box accuracy and quality on the overall segmentation performance of SAM.

Figure 2 shows the qualitative outcomes of fine-tuning SAM using our approach on Domain E and Domain F of the Prostate dataset. We randomly selected four instances from each domain and applied SAM, fine-tuned on Domain F, to segment instances in Domain E, and vice versa. The results demonstrate that the process of predicting coarse masks and applying our mask-filtering module to generate refined bounding boxes can effectively enhance the accuracy of SAM during the fine-tuning process.