Semi-Supervised Medical Image Segmentation via Knowledge Mining from Large Models

TL;DR

This work tackles data scarcity in medical image segmentation by enabling a lightweight model (U-Net++) to mine domain-specific knowledge from a large, generalist model (Segment Anything Model, SAM). The method trains U-Net++ on limited labels, then uses its predictions on unlabeled data to create prompts (points and bounding boxes) that guide SAM to generate pseudo labels, which in turn refine the U-Net++ through either one-time or continuous scheduling. Across Kvasir-SEG and COVID-QU-Ex, this knowledge mining approach yields consistent Dice and IoU improvements, often surpassing fully supervised baselines even with reduced labeled data, and maintains computational efficiency suitable for clinical use. The results highlight the feasibility of leveraging large foundation models to boost small, domain-specific models in data-constrained medical imaging tasks, with potential gains from domain-tuned SAM variants and self-supervised enhancements.

Abstract

Large-scale vision models like SAM have extensive visual knowledge, yet their general nature and computational demands limit their use in specialized tasks like medical image segmentation. In contrast, task-specific models such as U-Net++ often underperform due to sparse labeled data. This study introduces a strategic knowledge mining method that leverages SAM's broad understanding to boost the performance of small, locally hosted deep learning models. In our approach, we trained a U-Net++ model on a limited labeled dataset and extend its capabilities by converting SAM's output infered on unlabeled images into prompts. This process not only harnesses SAM's generalized visual knowledge but also iteratively improves SAM's prediction to cater specialized medical segmentation tasks via U-Net++. The mined knowledge, serving as "pseudo labels", enriches the training dataset, enabling the fine-tuning of the local network. Applied to the Kvasir SEG and COVID-QU-Ex datasets which consist of gastrointestinal polyp and lung X-ray images respectively, our proposed method consistently enhanced the segmentation performance on Dice by 3% and 1% respectively over the baseline U-Net++ model, when the same amount of labelled data were used during training (75% and 50% of labelled data). Remarkably, our proposed method surpassed the baseline U-Net++ model even when the latter was trained exclusively on labeled data (100% of labelled data). These results underscore the potential of knowledge mining to overcome data limitations in specialized models by leveraging the broad, albeit general, knowledge of large-scale models like SAM, all while maintaining operational efficiency essential for clinical applications.

Figures (2)

• Figure 1: Knowledge mining procedure for SAM. The U-Net++ is first trained in a supervised setting and directly adopted for SAM knowledge mining. The red square implies unlabeled images. The fire and snowflake icons indicate trainable and frozen modules, respectively. The respective supervised loss and pseudo loss are illustrated in Section \ref{['loss']}.
• Figure 2: Sample results on the Kvasir-SEG testing set for qualitative analysis.