SAM Carries the Burden: A Semi-Supervised Approach Refining Pseudo Labels for Medical Segmentation

TL;DR

This work leverage SAM's abstract object understanding for medical image segmentation to provide pseudo labels for semi-supervised learning, thereby mitigating the need for extensive annotated training data and outperforms intensity-based post-processing methods, state-of-the-art supervised learning for segmentation (nnU-Net), and the semi-supervised mean teacher approach.

Abstract

Semantic segmentation is a crucial task in medical imaging. Although supervised learning techniques have proven to be effective in performing this task, they heavily depend on large amounts of annotated training data. The recently introduced Segment Anything Model (SAM) enables prompt-based segmentation and offers zero-shot generalization to unfamiliar objects. In our work, we leverage SAM's abstract object understanding for medical image segmentation to provide pseudo labels for semi-supervised learning, thereby mitigating the need for extensive annotated training data. Our approach refines initial segmentations that are derived from a limited amount of annotated data (comprising up to 43 cases) by extracting bounding boxes and seed points as prompts forwarded to SAM. Thus, it enables the generation of dense segmentation masks as pseudo labels for unlabelled data. The results show that training with our pseudo labels yields an improvement in Dice score from $74.29\,\%$ to $84.17\,\%$ and from $66.63\,\%$ to $74.87\,\%$ for the segmentation of bones of the paediatric wrist and teeth in dental radiographs, respectively. As a result, our method outperforms intensity-based post-processing methods, state-of-the-art supervised learning for segmentation (nnU-Net), and the semi-supervised mean teacher approach. Our Code is available on GitHub.

Figures (5)

• Figure 1: Proposed pipeline: We use SAM to refine predictions $\mathcal{P}$ of the segmentation model $f_\theta$ trained on few labelled data. After mask cleaning, we extract a bounding box ( ), one positive ( ) and multiple negative ( ) seed points as sparse prompts for SAM, obtaining refined pseudo labels $\mathcal{R}$. Finally, we train a new segmentation model $f_\varphi$ with $\mathbf{R}$ on unlabelled data. See Sec. \ref{['sec:sam_pseudo_label']} for details.
• Figure 2: Dice scores reached by a selection of methods trained with an increasing number of labelled data $\mathcal{X}_\text{train}^n$ (see Sec. \ref{['sec:datasets']} for details). Line plots show the Dice means. We omit the box plots for some $n$ to increase readability.
• Figure 3: The median result on wrist test split, including DSC ($\mu\pm\sigma$) for a selection of methods. a) U-Net $f_\theta$, b) Mean Teacher (MT) and c) U-Net $f_\varphi$ trained with SAM refined pseudo labels $\mathcal{R}$ on unlabelled data $\mathcal{Y}$ (ours).
• Figure 4: First/second row shows the best/worst case of test split including DSC ($\mu\pm\sigma$) for a selection of methods. Second column: U-Net, third column: Mean Teacher (MT) and fourth column: U-Net trained with SAM refined pseudo labels on unlabelled data $\mathcal{Y}$ (ours).
• Figure 5: First/second group shows the best/median result of test split including DSC ($\mu\pm\sigma$) for a selection of methods. "MT" describes the Mean Teacher and "ours" the U-Net $f_\varphi$ trained with SAM refined pseudo labels $\mathcal{R}$ on unlabelled data $\mathcal{Y}$.