Simplicity Wins: Foundation Models for Few-Shot Medical Segmentation

Adapts the OOAL architecture (DINOv2 + CLIP) for medical image segmentation with minimal training data.

Paper: Simplicity Wins: Foundation Models for Few-Shot Medical Segmentation (2025)

Key Results

Architecture

Input Image (224x224)
        |
   [DINOv2 ViT-B/14]  <-- Frozen backbone
        |
   Multi-Layer Fusion (layers 3, 6, 9, 11)  <-- Trainable
        |
   Cross-Attention  <-----  [CLIP + CoOp]  <-- Text prompt
        |
   U-Net Decoder  <-- Trainable
        |
   Segmentation Mask

Key insight: Frozen DINOv2 features with a simple decoder achieve strong results with only 100 training samples.

Installation

git clone https://github.com/YOUR_USERNAME/OOAL-Medical.git
cd OOAL-Medical
pip install -r requirements.txt

Requirements: PyTorch 1.12+, Python 3.7+, CUDA 11.3+

Dataset Setup

Download datasets to the data/ directory:

Training

Basic training (100-shot):

python train_unet.py --dataset kvasir --num_shots 100 --iters 3000

With LoRA adaptation:

python train_unet.py --dataset kvasir --num_shots 100 --use_lora --lora_rank 4

MoE-LoRA (multi-domain):

python train_moe_lora.py --datasets kvasir isic --num_shots 100 --moe_num_experts 2

Training Arguments

Inference

Single model:

python inference_unet.py \
    --model_path save_models_unet/kvasir_100shot_.../best_model.pth \
    --dataset kvasir \
    --output_dir inference_outputs/

All models (batch):

python run_all_inference.py --num_samples 50

Project Structure

OOAL-Medical/
├── models/                 # Model implementations
│   ├── ooal_medical_unet.py   # Main model (DINOv2 + U-Net)
│   ├── moe_lora.py            # MoE-LoRA adapter
│   ├── lora_adapter.py        # LoRA implementation
│   ├── coop.py                # CoOp text prompting
│   └── seg_decoder_unet.py    # U-Net decoder
├── utils/                  # Utilities
│   ├── evaluation.py          # Metrics (Dice, IoU)
│   └── tta.py                 # Test-time augmentation
├── data/                   # Datasets (download separately)
├── paper/                  # Paper materials
├── train_unet.py           # Main training script
├── train_moe_lora.py       # MoE training
├── inference_unet.py       # Inference with visualization
└── run_all_inference.py    # Batch inference

Citation

@article{ooal-medical-2025,
  title={Simplicity Wins: Foundation Models for Few-Shot Medical Segmentation},
  author={[Authors]},
  year={2025}
}

Based on OOAL:

@inproceedings{li2024ooal,
  title={One-Shot Open Affordance Learning with Foundation Models},
  author={Li, Bo and Huang, Jiapeng and Mu, Qiang and Yang, Jingyi and Fang, Jiasheng and Chen, Yao and Chen, Qing and Liu, Shuguang},
  booktitle={CVPR},
  year={2024}
}

Acknowledgements

• OOAL - Original architecture
• CoOp - Context optimization
• DINOv2 - Vision backbone

License

MIT License - see LICENSE for details.