A segment anything model-guided and match-based semi-supervised segmentation framework for medical imaging

Abstract

Background

Semi-supervised segmentation leverages sparse annotation information to learn rich representations from combined labeled and label-less data for segmentation tasks. The Match-based framework, by using the consistency constraint of segmentation results from different models/augmented label-less inputs, is found effective in semi-supervised learning. This approach, however, is challenged by the low quality of pseudo-labels generated as intermediate products for training the network, due to the lack of the ‘‘ground-truth’’ reference.

Purpose

This study aims to leverage the foundation model, segment anything model (SAM), to assist unsupervised learning of Match-based frameworks. Trained with an extremely large dataset, SAM-based methods generalize better than traditional models to various imaging domains, allow it to serve as an assistant to Match-based frameworks to improve the quality of intermediate pseudo-labels for semi-supervised learning.

Methods

We propose SAM-Match, a SAM-guided and Match-based framework for semi-supervised medical image segmentation. Our approach involves two main steps: First, we use pretrained Match-based models to extract high-confidence predictions for prompt generation. Second, these prompts and unlabeled images are input into a fine-tuned SAM-based method to produce high-quality masks as pseudo-labels. And the refined pseudo-labels are further fed back to train the Match-based framework. SAM-Match can be trained in an end-to-end manner, facilitating interactions between the SAM- and Match-based models.

Results

SAM-Match demonstrates robust performance across multiple medical imaging datasets, including the ACDC cardiac MRI dataset, the BUSI breast ultrasound dataset, and an in-house liver MRI dataset (MRLiver). We partitioned the datasets into training, validation, and test sets (70%, 10%, and 20% for ACDC; 60%, 9%, and 31% for BUSI; and 62%, 12%, and 25% for MRLiver). On ACDC, with only 3 labeled cases, we achieved a Dice score of 89.36% ± 0.06% on 20 test cases. For BUSI, using just 30 labeled samples for training, we attained a Dice score of 59.35% ± 0.12% on 170 test samples. On MRLiver, training with only 3 labeled cases resulted in a Dice score of 80.04% ± 0.11% on 12 test scans. Wilcoxon signed-rank tests with Bonferroni corrections between the SAM-Match framework and the other comparison methods further demonstrated the statistical significance of SAM-Match's improvement in segmentation accuracy.

Conclusions

Our SAM-Match framework shows promising results in semi-supervised semantic segmentation, effectively tackling the challenges of automatic prompt generation for SAM and high-quality pseudo-label generation for Match-based models. It can help accelerate the adoption of semi-supervised learning in segmentation tasks, particularly in data-scarce scenarios. Our data and code will be made available at https://github.com/apple1986/SAMatch.