Auto-Generating Weak Labels for Real & Synthetic Data to Improve Label-Scarce Medical Image Segmentation

Authors: tanvi  deshpande
eva  prakash
elsie  gyang
curtis  langlotz
andrew  ng
jeya  maria
jose  valanarasu
Venue: MIDL_io_2024_Conference
Published: 2024-01-29 23:14:26

Abstract: The high cost of creating pixel-by-pixel gold-standard labels, limited expert availability, and presence of diverse tasks make it challenging to generate segmentation labels to train deep learning models for medical imaging tasks. In this work, we present a new approach to overcome the hurdle of costly medical image labeling by leveraging foundation models like Segment Anything Model (SAM) and its medical alternate MedSAM. Our pipeline has the ability to generate *weak labels* for any unlabeled medical image and subsequently use it to augment label-scarce datasets. We perform this by leveraging a model trained on a few gold-standard labels and using it to intelligently prompt MedSAM for weak label generation. This automation eliminates the manual prompting step in MedSAM, creating a streamlined process for generating labels for both real and synthetic images, regardless of quantity. We conduct experiments on label-scarce settings for multiple tasks pertaining to modalities ranging from ultrasound, dermatology, and X-rays to demonstrate the usefulness of our pipeline. The code will be made public after review.

Main paper: 
 introduction: 
 The process of automatically identifying and delineating specific structures within medical images, i.e. segmentation, holds immense use-cases in various tasks, including diagnosis, treatment planning, and surgical procedures. Deep learning-based methods such as (Ronneberger et al., 2015; Zhou et al., 2018; Milletari et al., 2016; Valanarasu et al., 2020 Valanarasu et al., , 2021;; Chen et al., 2021; Tang et al., 2022; Ma et al., 2024) have recently constituted advancements in medical image segmentation. Most of these methods are fully supervised networks and need a good amount of data and labels for them to perform well.While it is typical in computer vision tasks to obtain many segmentation labels through crowd workers, medical imaging tasks require experts to annotate the images. Due to the high cost of obtaining gold-standard labels from qualified medical professionals, practitioners often encounter the issue of label scarcity in many medical imaging tasks. Annotating medical images for segmentation is also particularly difficult due to the fine-grained nature of the label. In contrast, however, unlabeled data is more freely available. In addition, generative models like GANs (Goodfellow et al., 2020) and diffusion models (Kazerouni et al., 2023) can help generate synthetic medical images, further boosting the data size. However, this does not help augment datasets in a fully supervised pipeline, as we also need paired labels with the images to train the model.Recently, Kirillov et al. (2023) released Segment Anything Model (SAM), a vision foundation model trained on over 11 million images and 1 billion masks. It has the unique ability to segment any image out-of-the-box in a zero-shot setting. Segmentation using SAM can be done automatically without any inputs. However, this typically does not lead to optimal performance, and so segmentation is usually done with the help of prompting techniques on the image like points or boxes (Kirillov et al., 2023) . As SAM was trained on natural images, it is often found to be suboptimal when applied to medical imaging tasks (de Oliveira et al., 2023) . To address this gap, Ma et al. (2023) introduced MedSAM, a fine-tuned version of SAM created specifically for the medical imaging domain, trained on over 1 million images across 15 modalities. MedSAM was found to be superior in terms of performance on a variety of medical image segmentation tasks when compared to SAM.Figure 1 : Labels obtained from different configurations of SAM and MedSAM. Our method auto-generates effective input prompts (bounding boxes) using only very few annotations to generate high quality weak labels while SAM and MedSAM fail in auto options and are sensitive to input prompts in manual option.Intuitively, MedSAM should help us annotate the unlabeled images out-of-the-box and obtain labels to augment our training dataset. However, one major limitation of SAM and MedSAM is their sensitivity to input prompts. Both models are only as good as the point or box inputs they receive on a given image. In specific, for medical images, the area of interest to be segmented may be particularly subtle within the image with thin fine-grained boundaries. Also, inputs that include too much of the target or background result in over-or undersegmentation, requiring precise input prompts to get a good segmentation. This can be seen in Fig. 1 , where we show labels generated from SAM and MedSAM. MedSAM clearly gives us better labels than SAM, but it is very sensitive to the bounding box prompt we use and is still constrained by the manual prompting part. We also note that auto-prompting MedSAM performs poorly and often leads to blank segmentation predictions. Therefore, selecting an appropriate input prompt is key to ensuring success. Also, eliminating the need for manual prompting could allow practitioners to auto-generate labels for any number of unlabeled real or synthetic data, which would be very useful in label-scarce scenarios.To this end, we present a new pipeline that tackles the challenge of limited labeled data by harnessing the power of foundation models like MedSAM. Our approach leverages coarse labels, generated by training segmentation models on few labels (ranging from 25 to 50), to guide the selection of inputs for unlabeled data fed into the SAM model. This process effectively creates a richer dataset, enabling the training of a significantly more accurate model while using only a few gold-standard labels and eliminating the need for time-consuming, expensive manual labeling. This enriched dataset fuels the training of deep learning models with improvements in the dice accuracy ranging from 6.6% to 72.3% for medical image segmentation datasets like BUSI, ISIC, and CANDID-PTX.In summary, our contributions are as follows:• We introduce a pipeline that automatically generates efficient weak labels for any unlabeled data using MedSAM, eliminating the need for manual prompts.• We show that when using these weak labels to augment training datasets, we observe notable performance gains in deep learning models. This is especially impactful in scenarios with limited labeled data.• We conduct thorough validation of the method and ablation studies across ultrasound (BUSI), dermoscopy (ISIC), and X-ray (CANDID-PTX) datasets. 
related works: 
 After the introduction of SAM (Kirillov et al., 2023) and MedSAM (Ma et al., 2023) , these vision foundation models have been directly applied to a myriad of medical imaging tasks, such as brain tumor segmentation (Peivandi et al., 2023) , eye feature segmentation (Maquiling et al., 2023) , liver tumor segmentation, and lung nodule segmentation (He et al., 2023) , in both fine-tuning and zero-shot settings. Several approaches have also been introduced tailoring SAM for specific medical imaging tasks. For example, Lin et al. (2023) tailors SAM specifically for ultrasound segmentation by injecting features into SAM's encoder through a parallel CNN branch. Shin et al. (2023) jointly trains a network on heterogeneous ultrasound datasets using condition embedding blocks along with SAM, to allow SAM to adapt to each dataset separately. Chen et al. (2023) adapts SAM for 3D medical image segmentation, by incorporating 3D adapters into SAM's encoder.There has also been work involving self-prompting SAM for medical image segmentation. One branch of such work involves separately learning prompts for SAM and combining them with SAM's existing architecture. One such approach includes learning a pixel-wise classifier from SAM's own embedding space and encoder to prompt SAM in few-shot settings (Wu et al., 2024) . In addition, Lei et al. (2023) use extreme points from 3D medical images to generate 2D bounding-box prompts for SAM. Pandey et al. (2023) first train YOLOv8 on several image-mask pairs to detect bounding-boxes for the regions to segment, which are then fed to SAM. Lastly, Anand et al. (2023) use localization to a template image to prompt SAM in one-shot settings. The other branch of such work involves altering SAM's architecture to incorporate learning prompts. Shaharabany et al. (2023) learn prompts for SAM by training a separate prompt encoder to automatically generate prompts, rather than conditioning on manual prompts, as SAM does. Cui et al. (2023) use SAM to generate weak segmentation labels, which are then used to fine-tune SAM.Annotations generated from SAM can also be used to augment existing segmentation pipelines, such as U-Nets (Zhang et al., 2023) . In addition, there has been work regarding sampling prompts from SAM (Qi et al., 2023) , and SAM has been used for a variety of non-medical tasks, such as inpainting (Yu et al., 2023) and captioning (Wang et al., 2023) . 
method: 
 We propose a pipeline that automatically generates weak labels for both unlabeled real and synthetic data using MedSAM, eliminating the need for manual prompting. Our pipeline includes training a small model on the available gold-standard labels, using predictions from this model to generate prompts for MedSAM, and retraining using a larger dataset consisting of the gold-standard and weak-labeled data. Our method is illustrated in Fig. 2 .Figure 2 : An illustration of our pipeline to auto-generate weak labels for unlabeled data.We use the limited annotations to train an initial model that generates low-quality coarse labels on unlabeled data. We then select inputs from these coarse labels as prompts to MedSAM to create higher-quality weak labels. These weak labels are used to train a stronger segmentation model.Preliminaries A label-scarce scenario in a medical imaging setting can be defined as a task for which there exists a small amount of labeled (gold-standard) data, and a larger pool of unlabeled data. Let the small labeled dataset be denoted by D labeled , where we assume less than 50 gold-standard labels from medical practitioners. We chose this number after considering some real clinical scenarios, in which we encountered situations where it was difficult to obtain more than 50 gold standard labels. We acknowledge that this number may vary depending on the specific task, and therefore, we conduct ablation studies using different label quantities in Section 5. Notably, acquiring 50 high-quality labels is typically more feasible than the hundreds to thousands often required for traditional segmentation models. We define the larger pool of unlabeled data as D unlabeled , which may comprise real medical scans or synthetic images generated by models like diffusion models. The evaluation set, denoted as D test , remains separate for performance assessment. We also define weak labels as labels that are not gold-standard (as they are not created by experts), but are still useful in improving the model performance.Coarse Label Generation To make use of the full potential of limited annotations, we first train an initial model θ on the labeled dataset D labeled . This model is used to generate coarse labels for the unlabeled dataset D unlabeled . These coarse labels, despite potentially limited accuracy due to the model's training on a smaller number of labels, serve as valuable inputs for automatically prompting MedSAM to produce higher-quality weak labels.The prompts to MedSAM can be either points or bounding-boxes. Both "positive points" (corresponding to portions of the image to segment) and "negative points" (corresponding to portions of the image not to segment) can be given as input point prompts. For bounding-boxes, one needs to draw a box around their region of interest and feed that as the prompt input to MedSAM. Although these prompts are an easier alternative to pixel-by-pixel annotations, they still require manual effort and a good knowledge about the segmentation task.To auto-generate labels using MedSAM without any manual intervention, we first pick up the coarse label prediction from the initial model θ. If the coarse label is empty, we filter out the sample, since there is no basis for a prompt to MedSAM. If not, we pick the largest contiguous blocks from the coarse label and filter out small blocks, as those are likely to be noise. For point inputs, we use the middle of the largest contiguous region(s) of the coarse label as the prompt to MedSAM; in addition, low-probability points from the original prediction mask are used as "negative points" for the MedSAM model. For bounding-box inputs, we compute the minimum and maximum indices of the largest contiguous region(s) of the coarse label as the prompt to MedSAM. An example of the input-selection process is shown in Fig. 3 . We fixed these techniques after thorough experimentation on various input selection methods (App. C). While the input selection technique we describe works well for multiple tasks, we do acknowledge that for some tasks, there could be more intelligent prompting techniques that could perform better. However, we would like to point out that our input selection technique is much more generic and can be helpful in obtaining useful weak labels. It should also be noted that these weak labels are far better than the ones generated by SAM and MedSAM through automatic option (App B).To ensure the quality of weak labels generated from unlabeled images, we implement a filtering process that eliminates extreme cases. Specifically, masks containing a disproportionately high percentage of pixels (> 97%) belonging to either the background or the segmented class are excluded. This filtering step safeguards against misleading or uninformative masks that could potentially hinder model performance.Augmenting the training data Finally, the weak labels generated for unlabeled images are combined with D labeled to form the final augmented dataset D augmented , consisting of images, where a small number of labels (25-50) are gold-standard, and the rest are weak labels. This augmented dataset will be used in the final training of the model.In cases where data scarcity might further exacerbate label scarcity, we explore the potential of synthetic image generation as a means to potentially increase dataset size. Thus, we additionally generate synthetic images from each dataset in order to evaluate the performance of our pipeline on such data. We train a denoising diffusion probabilistic model (DDPM) with a UNet backbone following (Ho et al., 2020) . We set the hyperparameters following (Wolleb et al., 2021) , notably with a learning rate of 1e -4 , image size of 256, and batch size of 4. 
experiments and results: 
 In this section, we give details about the datasets we use, experimental setup, and the results that validate the usefulness of our method.We utilize four datasets featuring 2D binary medical image segmentation tasks. Within each dataset, we randomly select N random samples from the training set to be the "goldstandard" labels. Note that we experiment on different values of N to mimic different label-scarce settings. The rest of the training set is considered unlabeled data (i.e. we do not use the ground-truth labels in our experiments, but generate weak labels for these images). This approach enables us to train and evaluate our model without relying on the full set of ground-truth annotations. BUSI: The Breast Ultrasound Images Dataset (BUSI) consists of ultrasound images in three classes: benign, malignant, and normal (Al-Dhabyani et al., 2020) . We focus on segmentation of benign and malignant lesions from the ultrasound images. The number of samples with segmentation masks is 650. We use a randomized train-test split of 80-20, resulting in 130 samples in the test set. ISIC: The International Skin Imaging Collaboration (ISIC) dataset consists of dermoscopic lesion segmentation tasks (Gutman et al., 2016) . It comprises 1279 dermoscopic lesion images, divided into a training set of 900 images and a test set of 379 images. CANDID: We use a subset of the CANDID-PTX dataset (Feng et al., 2021) , consisting of 500 images, which consists of binary segmentation of lungs from chest X-rays. We use a randomized 80-20 train-test split.We use UNet++ (Zhou et al., 2018) as our base segmentation network. Please note that our method is agnostic to the choice of network; we picked UNet++ as it is stable in training on label-scarce conditions. We use a combination of DICE loss and binary cross-entropy loss with a scaling ratio of 1 and 0.5 respectively to train our models. We use the SGD optimizer with a learning rate of 10 -3 , and a cosine annealing learning rate scheduler with minimum learning rate 10 -5 . All images were also reshaped to a resolution of 256 × 256. We use a batch size of 4 and train for 100 epochs on NVIDIA A4000 GPUs.We provide the quantitative results for 3 datasets in Table 1 in label-scarce settings. The performance of models trained for datasets augmented with weak labels is compared against the performance of a model trained on just the base dataset D labeled , which consists of goldstandard (GS) labels. In addition, we compare against SAM's automatic option as well as UniverSeg (Butoi et al., 2023) , using D labeled as the support dataset. We provide results for both bounding-box and point prompts (automatically generated by our pipeline), observing that different prompts are more suited to different datasets. We achieve improvements of up to 73.3%, with more dramatic improvements where the initial DICE score was lower; though the weak labels may not be an exact match with the gold-standard labels, they are accurate enough to provide a boost in performance. More qualitative results can be seen in the appendix (App. A). Synthetic Data In addition to experiments using real data, we conduct experiments using data from a diffusion model trained on unlabeled samples from each of the datasets (Table 2 ). We train the diffusion model on the train splits of BUSI, ISIC, and CXR-COVID (Fraiwan et al., 2023) datasets respectively and sample it multiple times to generate the required number of synthetic images. As diffusion models can produce large amounts of synthetic data after being trained on just a few samples, the number of possible weak labels is higher.In addition to our main experiments, we perform ablation studies examining various elements of our pipeline and also discuss some limitations in this section.Label scarce settings We vary the number of images in the base dataset to observe the effect on the performance boost from weak labeling and present the results in Table 3 . Though increasing gold-standard labels is expected to reduce performance gains from weak labels, when starting with only 10 base images, the quality of coarse labels plummets, rendering MedSAM prompting ineffective in such extreme label scarcity. However, it can be noted that we always gain performance while using our pipeline. Limitations Our work focuses on 2D ultrasound, X-ray, and dermoscopic data. We recognize that tasks involving highly intricate structures may require different input selection approaches due to potential sensitivities in our method (App. D). We also did not touch on 3D segmentation tasks; future work could investigate extensions to address such scenarios. 
conclusion: 
 We introduce a new method for addressing label-scarce scenarios in medical image segmentation using recent advancements in vision foundation models. By selecting inputs to MedSAM from coarse labels trained on a small gold-standard dataset, we create augmented datasets with weak labels that can be auto-generated for any number of unlabeled data. Using these augmented datasets, we train models that obtain significant boosts in performance on label-scarce settings. Weak labels generated through our method can also be used to improve human-in-the-loop annotation processes.  
 
 Appendix:  


Decision: Reviewers appreciated the concept of the paper of using foundation models to generate labels, but they all raised initial concerns about this paper
- Benchmarking is insufficient with respect to well-performing baselines
- Insufficient analysis of the differences between pseudolabels generated by the different approaches
- Generalizability of foundation models to specific medical imaging tasks is not clear

The authors engaged in the rebuttal and modified significantly their paper, in my opinion addressing all concerns raised in the relatively short reviews satisfactorily.

Thus, while reviewers maintained or lowered their scores, I recommend acceptance.

Reviews: review: paper_summary: paper_summary: This work addresses the challenge of generating segmentation labels for medical imaging tasks. By leveraging foundation models like Segment Anything Model (SAM) and its medically adapted counterpart, MedSAM, the authors propose a pipeline capable of generating weak labels for unlabeled medical images, thereby augmenting datasets with limited labels. This method involves training a model on a few gold-standard labels, which then guides MedSAM in generating weak labels, streamlining the process for both real and synthetic images. Experiments across various medical imaging modalities, including ultrasound, dermatology, and X-rays, demonstrate the effectiveness of this approach, with notable improvements in model performance using the augmented datasets.
main_review: main_review: I am retaining my original rating for the paper. While I do appreciate the efforts of the authors for rebuttal, I do not believe that sufficient additions have been made to the paper to meet the required quality.
strength_weakness: strength_weakness: One potential weakness is the reliance on foundation models like SAM and MedSAM, which, despite their adaptability, may still not capture the nuances of highly specialized medical imaging tasks. 

Also, this work does not provide a comparison of the state-of-the-art segmentation tools from the respective modalities to proposed method.
questions: questions: How do the state-of-the-art, domain specific models correspond to being trained with the weak labels generated by the proposed method?
paper_summary: paper_summary: The paper presents an auto-prompting approach to generate pseudo labels for semi-supervised segmentation training applied to 2D medical image segmentation. Multiple 2D image modalities and segmentation tasks were evaluated and comparison of SAM and MedSAM approaches to generate the pseudo labels was performed.
main_review: main_review: The idea of generating auto-prompts to generate pseudo labels is conceptually novel, albeit the method itself is rather incremental. Moreover, the evaluations are limited and comparison to other baselines are not considered.
strength_weakness: strength_weakness: * Besides the generation of auto-prompts that basically perform simple morphological operations to extract bounding box on the MedSAM segmentations, the method is rather incremental. 
* Whereas the stated problem is the variability of generated pseudo labels due to prompts, a detailed analysis of variability in the generated pseudo labels due to different methods and to what extent they impact downstream accuracy is missing. Most of the analysis focused on the number of training examples used to train the initial model, which seems to have a minimal impact on the accuracy. 
* Baseline comparison to other methods such as UniverSeg are missing.
questions: questions: * Whereas the stated problem is the variability of generated pseudo labels due to prompts, a detailed analysis of variability in the generated pseudo labels due to different methods and to what extent they impact downstream accuracy is missing. Most of the analysis focused on the number of training examples used to train the initial model, which seems to have a minimal impact on the accuracy. 
* Baseline comparison to other methods such as UniverSeg are missing.
paper_summary: paper_summary: This paper proposed a method that leverages foundation models to generate weak labels and then augment label-scarce datasets. The validation experiments on three public datasets demonstrated the outperformance of the proposed method.
main_review: main_review: Thanks to the authors for the careful response. However, the innovation and content are still not strong enough for me to be accepted by MIDL. A more necessary motivation is expected, and more experiments are required.
strength_weakness: strength_weakness: The innovation of this paper is the proposed pipeline, which is quite tricky and not strong enough. 
Why do not use MedSAM in the first-round training as the initial model? What is the advantage of using the initial model compared with using MedSAM to generate coarse labels?
questions: questions: Please clarify the disadvantage of fine-tuning MedSAM to generate coarse labels compared with training the initial model.

Comments: comment: We believe that we have thoroughly addressed the reviewer's concerns and kindly request the reviewer to let us know if there are  any additional issues following our response. We thank the reviewer for their time.
We believe that we have thoroughly addressed the reviewer's concerns and kindly request the reviewer to let us know if there are  any additional issues following our response. We thank the reviewer for their time.
We believe that we have thoroughly addressed the reviewer's concerns and kindly request the reviewer to let us know if there are  any additional issues following our response. We thank the reviewer for their time.
We thank all of the reviewers for their feedback. We have addressed all of their comments in the following replies and made changes accordingly in the revised edition of the paper.
**1) Why not use MedSAM in the first-round training as the initial model?**

As demonstrated in our paper, MedSAM has no auto option; generating predictions on images without a prompt largely results in blank masks. While we could use MedSAM to generate coarse labels independent of our pipeline, doing so would require manual prompts (again from, for instance, medical experts), making it labor-intensive. Hence, our method of generating prompts automatically is advantageous, as it has no manual intervention, and so could be very fast to generate weak labels for any number of real or synthetic data. We have added baseline results from the automatic version of SAM in our paper; without prompts, SAM segments many regions of the image, in which it is not clear which corresponds to the target, and we select a random region as the target mask. Training using these weak labels results in performance drops for all three datasets, in comparison to the notable gains from using our method.
**1) Addressing the comment on methodology being incremental**

Please note that our work is on the Validation & Application track, which is intended for “well-validated applications of deep learning algorithms in medical imaging”. Our focus was on validating that foundation models can be used effectively in real-life label-scarce scenarios (which was shown through experiments on a variety of modalities and dataset sizes). We do propose a new pipeline and show that it is generalizable. While it is possible to develop specific heuristic-based methods (such as thresholding based on image contours or darkness) for each dataset, our focus was on creating a dataset- and modality-agnostic pipeline. The generalizability of our method can be noted from the performance gains on all datasets without changes in the method.

**2) Variability of generated pseudo labels**

Figures 1 and 5 qualitatively demonstrate how variable MedSAM and SAM are in generating pseudo labels. To further motivate this claim, we have now done experiments involving using SAM’s auto option, which is included in the main results table; without prompts, SAM segments many regions of the image; selecting one arbitrarily to use as the segmentation label results in performance drops. 
 
We also ablated on the type of auto-prompt we chose for our method in Appendix C. We provided comparisons of various input selection methods to our method in Figure 6.

**3) Comparisons with UniverSeg**

We have added this to our main results table, along with baseline results corresponding to labels from SAM in its auto option.  For UniverSeg, we evaluate on the test set using the gold-standard dataset as the support images and labels; our method outperforms UniverSeg on all 3 datasets. Without prompts, SAM segments many regions of the image, in which it is not clear which corresponds to the target, and we select a random region as the target mask. Training using these weak labels results in performance drops for all three datasets, in comparison to the notable gains from using our method. We include the table here.

| \# GS Labels | \# Weak Labels | |  |  |  | **DICE** |  |  |  |  |
|--------------|---------------|----------|----------|----------|----------|----------|----------|----------|----------|----------|
|              |               |  | **BUSI** |  |  | **ISIC** | |  | **CANDID-PTX** |  |
| UniverSeg    |               |   | 0.3681   |   |   | 0.5257   |   |   | 0.7700     |     |
|              |                        | Auto     | Box      | Points   | Auto     | Box      | Points   | Auto       | Box        | Points     |
| 25           | 0             | 0.3059   | 0.3059   | 0.3059   | 0.6123   | 0.6123   | 0.6123   | 0.8182   | 0.8182     | 0.8182     |
| 25           | 25            | 0.2629   | 0.3613   | 0.3777   | 0.5810   | 0.7367   | 0.7884   | 0.6396   | 0.8879     | 0.8726     |
| 25           | 50            | 0.2401   | 0.4326   | 0.4661   | 0.5907   | 0.7587   | 0.8087   | 0.5519   | 0.9044     | 0.8872     |
| 25           | 100           | 0.2124   | 0.4661   | **0.5302**| 0.5893   | 0.7424   | **0.8483**| 0.4115   | **0.9096**| 0.8443     |
**1) On Reliance of SAM and MedSAM:**

Please note that our work is not focused on solving the nuances that occur in highly specialized medical imaging tasks. Our main goal is to show how label-scarce tasks in medical imaging can be boosted using weak labels to augment training datasets; we observe notable performance gains (up to 73%). We would like to point out that our method being generalizable and adaptable by relying on MedSAM is in fact a strength and not a weakness.

**2) Comparison with domain-specific models:**

Our work is designed to be implemented on top of any state-of-the-art, domain-specific model. In our work, we considered UNet++ as our base network, as most domain-specific segmentation frameworks follow a UNet/UNet++ network architecture with domain-specific loss functions or modules. Our contribution is in boosting the size of the available dataset for training, which can be implemented on top of any domain-specific segmentation model.