Domain and Task-Focused Example Selection for Data-Efficient Contrastive Medical Image Segmentation

1. Problem Statement

Medical image segmentation is critical for diagnosis and treatment planning but relies heavily on fully supervised learning, which requires massive amounts of pixel-level annotated data. This annotation process is expensive, time-consuming, and prone to inter-observer variability. While Self-Supervised Learning (SSL) and Contrastive Learning (CL) offer alternatives by pre-training on unlabeled data, existing CL methods for medical imaging face specific challenges:

• Over-reliance on Augmentation: Standard CL (e.g., SimCLR) creates positive pairs via geometric/color augmentations, which often fail to capture the nuanced anatomical relationships inherent in medical datasets.
• Lack of Task Relevance: Generic pre-training strategies do not necessarily learn features relevant to specific downstream segmentation tasks (e.g., organ boundaries).
• Data Scarcity: Models struggle to generalize when fine-tuned on very limited labeled data or when applied to cross-domain datasets (different scanners or protocols).

2. Methodology: The PolyCL Framework

The authors propose PolyCL, a novel self-supervised contrastive learning framework designed specifically for medical image segmentation. It follows a two-stage training process:

A. Pre-training (Self-Supervised Contrastive Learning)

Instead of relying solely on data augmentation, PolyCL introduces Domain and Task-Focused Example Selection strategies to construct positive and negative pairs from unlabeled CT scans. Three strategies are proposed:

1. PolyCL-S (Scan-based): Selects positive pairs from the same scan (intra-scan) and negative pairs from different scans. This learns intra-scan consistency without needing organ labels.
2. PolyCL-O (Organ-based): Requires knowing which slices contain the target organ. Positive pairs are slices containing the same organ status (both contain or both lack the organ), while negatives have opposite status. This forces the encoder to learn organ-specific features.
3. PolyCL-M (Mixed): Combines the above. Positive pairs must be from the same scan and share the same organ status. This learns both intra-scan coherence and inter-scan invariance while discriminating based on anatomical content.

The model uses a standard contrastive loss (InfoNCE) to align embeddings of positive pairs and push apart negative pairs.

B. Fine-tuning (Supervised Segmentation)

After pre-training, the projection head is discarded. A U-Net style decoder is attached to the pre-trained encoder. The model is fine-tuned on a small set of labeled data using a Dice loss function.

C. Integration with Segment Anything Model (SAM)

To further enhance performance with limited labels, the authors integrate SAM in two innovative ways:

1. SAM Refinement (Post-processing): Coarse segmentation masks generated by PolyCL are converted into bounding box prompts. These prompts are fed into SAM to refine the mask boundaries, leveraging SAM's zero-shot capabilities for anatomical accuracy.
2. SAM 2 Propagation: Utilizing SAM 2's video segmentation capabilities, the framework propagates a segmentation mask from a single annotated 2D slice across an entire 3D CT volume. This enables volumetric segmentation with minimal manual annotation (as few as one slice).

3. Key Contributions

• Novel Pre-training Strategy: Introduction of PolyCL, which moves beyond augmentation-based CL to use anatomical and scan-level context for creating meaningful positive/negative pairs.
• Three Selection Paradigms: The definition and evaluation of Scan-based (S), Organ-based (O), and Mixed (M) example selection strategies, demonstrating that task-aware selection improves feature learning.
• SAM Integration: A dual approach to leveraging foundation models: using SAM for boundary refinement and SAM 2 for 3D volume propagation from single slices.
• Comprehensive Evaluation: Extensive experiments on multiple public CT datasets (LiTS, TotalSegmentator, MSD, KiTS, BTCV) covering in-domain, cross-domain, and low-data regimes.

4. Experimental Results

The framework was evaluated on liver and kidney segmentation tasks across various encoder backbones (U-Net, ResU-Net, TransUNet).

• In-Domain Performance: PolyCL variants consistently matched or slightly trailed the best fully supervised baselines in Dice scores but significantly outperformed them in Hausdorff Distance (boundary accuracy). For example, PolyCL-M with a ResU-Net backbone improved Hausdorff distance by ~13% over the best fully supervised baseline.
• Cross-Domain Generalization: PolyCL demonstrated superior robustness when tested on unseen datasets (e.g., training on LiTS, testing on TotalSegmentator). PolyCL-S with ResU-Net achieved a 10.7% improvement in Dice over the best fully supervised model in cross-domain settings.
• Low-Data Efficiency: In reduced-label scenarios (5%–50% labeled data), PolyCL significantly outperformed standard SSL methods (SupCon, SimCLR).
  • At 5% labeled data, PolyCL-M + SAM achieved a Dice of 0.682, outperforming the baseline by ~19%.
  • SAM Refinement drastically reduced boundary errors (Hausdorff distance), with PolyCL-O + SAM achieving a distance of 4.695 compared to much higher values for unrefined models.
• 3D Propagation: Using SAM 2 to propagate a single annotated slice allowed for high-quality 3D segmentation, with PolyCL-M + SAM 2 achieving a Dice of 0.765 at 50% labeling, proving effective even with ultra-low annotation.

5. Significance and Clinical Relevance

• Data Efficiency: PolyCL significantly reduces the dependency on large annotated datasets, making advanced segmentation feasible for rare conditions or institutions with limited resources.
• Annotation Reduction: The integration of SAM 2 allows for 3D volumetric segmentation from a single 2D slice, potentially reducing radiologist annotation time by orders of magnitude.
• Boundary Precision: The combination of PolyCL pre-training and SAM refinement addresses the common issue of "fuzzy" boundaries in deep learning segmentation, which is critical for surgical planning and dosimetry.
• Generalizability: The framework's ability to transfer across different CT scanners and protocols suggests it is a robust solution for real-world clinical deployment where data distribution shifts are common.

In conclusion, PolyCL represents a significant step forward in medical image analysis by aligning self-supervised pre-training with domain-specific anatomical knowledge and leveraging foundation models to bridge the gap between limited data and high-precision clinical requirements.