AI Research Portal

This paper introduces MedSMILE, a novel framework for medical image segmentation that achieves near fully supervised performance using only a single annotated image and abundant unlabeled data. MedSMILE integrates transductive and inductive learning in an iterative loop, where a foundation model bootstraps and refines pseudo-labels for an inductive segmentation model. Experiments on seven datasets across four modalities show MedSMILE recovers 90%-95% of the fully supervised Dice score, outperforming existing semi-supervised techniques that require substantially more annotations.

The paper addresses the significant bottleneck of annotation cost in medical image segmentation, which typically requires large datasets with pixel-level annotations. This expense hinders the development and deployment of automated segmentation tools, especially for diverse medical conditions and imaging modalities. Existing solutions like Few-Shot Learning (FSL) and Semi-Supervised Learning (SSL) have limitations: FSL/OSL often require specific support/query structures, and SSL typically assumes access to more than one labeled sample and necessitates careful handling of pseudo-label noise.

Introduction of MedSMILE, a novel framework that synergistically integrates transductive and inductive learning for the extreme one-label semi-supervised setting.
Demonstration that competitive segmentation performance is achievable from a single annotated sample, recovering 90%-95% of fully supervised performance across diverse medical datasets.
Establishment of a new benchmark for annotation efficiency in medical image segmentation, significantly outperforming state-of-the-art semi-supervised methods that use orders of magnitude more labeled data.

MedSMILE is a multi-round pipeline for medical image segmentation training, designed for the extreme one-label semi-supervised setting. It utilizes a foundation model (FM), specifically DINOv2, and an iterative self-training process.

The pipeline consists of:

Optional Self-Supervised Encoder Fine-tuning: Adapts the DINOv2 encoder to the target unlabeled medical images (Du) using a contrastive learning framework (MoCo-style) and InfoNCE loss.
Round 0: Initial Pseudo-Label Generation: Leverages the DINOv2 encoder (potentially fine-tuned) to extract features from the single labeled image (Iref, Mref) and unlabeled images (Du). Class prototypes are calculated from Iref, and cosine similarity is used to compute low-resolution similarity maps. These are upsampled and converted to initial soft pseudo-labels using a pixel-wise softmax.
Rounds r = 1...R: Iterative Model Training and Pseudo-Label Refinement:
- Dataset Preparation: Soft pseudo-labels are converted to hard labels. For r > 1, the input is the refined pseudo-labels from the previous round.
- Model Initialization and Training: A Mask2Former model (MM2F) is initialized with pre-trained weights and re-initialized at the start of each round. It is trained on the current dataset using a segmentation loss (e.g., cross-entropy and Dice loss).
- Pseudo-Label Generation and Uncertainty Estimation: The trained MM2F performs inference on unlabeled images to generate raw pseudo-labels. Uncertainty is estimated using the average pixel-wise entropy of the probability map.
- Pseudo-Label Refinement via KNN Propagation: Uncertain pseudo-labels are refined by propagating information from high-confidence samples in the certain set. KNN is used to find nearest neighbors from the certain set, and their labels are weighted-averaged based on feature similarity.
- Prepare for Next Round: The refined pseudo-labels are saved and used as input for the next round.

The process repeats for R rounds, ultimately outputting a trained segmentation model.

Architecture:
- Foundation Model: DINOv2 Vision Transformer (potentially fine-tuned).
- Segmentation Model: Mask2Former (MM2F).
- Self-supervised Fine-tuning: Uses a student-teacher network with projection heads (shallow MLP).
Key Equations:
- Class Prototype Calculation: p'c = Σ_{h,w} I(Mref[h, w] = c) · Fref[:, h, w] / (Σ_{h,w} I(Mref[h, w] = c) + ε) pc = p'c / ||p'c||2 + ε
- Similarity Computation: Slow [h, w] = Fu[:, h, w] · Pc
- Initial Soft Pseudo-Labels: PL(0) [k, h, w] = exp(Shigh [h, w]) / Σ_{j=0}^{C-1} exp(Shigh [h, w])
- Hard Pseudo-Labels (Round r): Mraw [h, w] = argmax L(r) [c, h, w]
- Uncertainty Metric: u(r) = -1/(HW) Σ_{h,w} Σ_{c} P(r) [c, h, w] log P(r) [c, h, w]
- Refined Soft Pseudo-Labels: Mr) [k, h, w] = Σ_{j=1}^{Q} wj * OHE(Mraw)[k, h, w] / (Σ_{j=1}^{Q} wj + ε)
Loss/Objective:
- Self-supervised Fine-tuning: InfoNCE loss.
- Segmentation Training: Multi-class segmentation loss (e.g., combined cross-entropy and Dice loss).
Optimization: AdamW optimizer with a Cosine Annealing learning rate scheduler.
Augmentation: Multi-crop augmentation strategy for self-supervised fine-tuning.
Data Handling:
- Input: One labeled image (Iref, Mref) and N unlabeled images (Du).
- Pseudo-label conversion: Soft pseudo-labels converted to hard labels via argmax.
- KNN Propagation: Uses a pre-defined uncertainty quantile (0.9) to identify uncertain samples. K=5 for KNN.
Implementation Details:
- Number of Rounds (R): 3 iterative refinement rounds after Round 0.
- Learning Rate: 5 × 10⁻⁵.
- Batch Size: 4.
- Gradient Clipping: Maximum norm of 1.0.
- Uncertainty Quantile Threshold: 0.9.
- KNN Neighbors (K): 5.
- DINOv2 Layer for Features: x_norm_patchtokens.

Datasets:
- Polyp Segmentation: CVC-ColonDB, Kvasir, CVC-ClinicDB, ETIS.
- Skin Lesion Segmentation: ISIC 2018.
- Brain MRI Segmentation: LGG dataset.
- Ultrasound Segmentation: BUSI dataset.
Metrics: mDice, mIoU.
Baselines: MT, DAN, UA-MT, URPC, CLCC, SLC-Net, MC-Net+, CDMA, SCP-Net, MCF, KnowSAM, DEC-Seg, GenSeg, Polyp-PVT, PraNet-V2, iMAML, SkinFormer, Ms RED.
Numeric Results:
- Polyp Segmentation (CVC-ClinicDB, Kvasir): MedSMILE achieved mDice of 0.843 (vs. 0.836 for DEC-Seg) and 0.868 (vs. 0.859 for DEC-Seg), using only 1 labeled sample compared to 145 for others.
- Polyp Segmentation (CVC-ColonDB, ETIS): MedSMILE achieved mDice of 0.8255 (vs. 0.648 for DEC-Seg) and 0.7705 (vs. 0.592 for DEC-Seg).
- Skin Lesion Segmentation (ISIC-2018): MedSMILE achieved mIoU of 0.792 using 1 labeled image, outperforming methods with 259 labeled samples. It achieved 93% of the Dice score of fully supervised SkinFormer.
- LGG Brain MRI: MedSMILE achieved mDice of 0.882 and mIoU of 0.862, recovering approximately 92.3% of fully supervised performance.
- BUSI Ultrasound: MedSMILE achieved mDice of 0.765 and mIoU of 0.672, recovering approximately 95.1% of fully supervised mDice.
- Comparison to Fully Supervised: MedSMILE recovered 90%-95% of fully supervised Dice scores on polyp datasets. On CVC-ColonDB, it outperformed Polyp-PVT.
Runtime/Training Cost: The full MedSMILE pipeline on the ISIC dataset completed in under an hour on a single consumer-grade GPU (RTX 4070 Ti SUPER). The main computational cost comes from Mask2Former training, with constant overhead per round for inference and KNN propagation.

When the foundation model (FM) feature encoder fails to extract robust features, self-supervised fine-tuning on the unlabeled dataset becomes necessary, which can be time-consuming.
In fine-grained segmentation tasks requiring distinction between visually similar structures, the encoder may still struggle to capture subtle differences even after fine-tuning, potentially leading to unreliable initial pseudo-labels and limiting effectiveness.

Method evidence: "Its core novelty lies in the iterative loop where a foundation model both bootstraps and refines pseudo-labels for an inductive segmentation model." (section: Key Contributions)
Result evidence: "Across diverse medical datasets, MedSMILE recovers 90%-95% of the fully supervised performance (in terms of Dice score)." (section: Key Contributions)
Limitation evidence: "When the foundation model (FM) feature encoder (e.g., DINOv2) fails to extract robust features, self-supervised fine-tuning on the unlabeled dataset becomes necessary. While such fine-tuning can improve domain adaptation and performance, it can be time-consuming." (section: Limitations)

Medical Image Segmentation with Minimal Labeling Effort: How Far Can We Push the Limits?

Summary

Problem

Key Contributions

Method

Technical Details

Results

Weaknesses / Limitations

Evidence

Additional Figures