UNIStainNet: Foundation-Model-Guided Virtual Staining of H&E to IHC

UNIStainNet is a SPADE-UNet generator conditioned on dense spatial tokens from a frozen UNI pathology foundation model (ViT-L/16). A single 42M-parameter model generates multiple IHC stains (HER2, Ki67, ER, PR) from H&E input via learned stain embeddings with FiLM modulation.

• Dense UNI spatial conditioning: 32×32 = 1,024 spatial tokens from a frozen UNI ViT-L/16, providing 64× more spatial tokens than CLS-token approaches
• Misalignment-aware losses: Designed for consecutive-section training pairs with 10–50 px misalignment (perceptual loss, unconditional discriminator, DAB intensity supervision)
• Unified multi-stain model: One model serves four IHC markers at 4× fewer parameters than per-stain specialists, with comparable stain accuracy
• Tissue-type stratified failure analysis: Zero-shot tissue classification reveals where generation fails (non-tumor tissue) vs. succeeds (invasive carcinoma: <3% failure rate on MIST)

A single model produces stain-specific expression patterns — membrane (HER2), punctate nuclear (Ki67), and diffuse nuclear (ER/PR):

The same H&E input translated to four IHC stains by changing only the stain embedding. ER and PR produce visually similar diffuse nuclear staining, consistent with their co-expression in most breast cancers:

MIST — HER2 comparison with prior methods (click to expand)

BCI — HER2 class-wise comparison (click to expand)

pip install -r requirements.txt
pip install -e .

UNI model weights are downloaded automatically via timm from MahmoodLab/uni. You will need a Hugging Face token with access granted.

Download from the BCI Challenge and organize as:

BCI_dataset/
├── HE/
│   ├── train/       # H&E training images (1024×1024)
│   ├── test/        # H&E test images
│   └── val/         # (optional)
└── IHC/
    ├── train/       # IHC (HER2) training images
    ├── test/
    └── val/         # (optional)

Images are paired by filename. HER2 class labels (0, 1+, 2+, 3+) are inferred from filename prefixes.

Download from MIST and organize as:

MIST/
├── HER2/
│   └── TrainValAB/
│       ├── trainA/   # H&E images (512×512)
│       ├── trainB/   # IHC HER2 images
│       ├── testA/
│       └── testB/
├── Ki67/
│   └── TrainValAB/ ...
├── ER/
│   └── TrainValAB/ ...
└── PR/
    └── TrainValAB/ ...

python scripts/train/train_bci.py \
    --data_dir /path/to/BCI_dataset \
    --batch_size 4 \
    --max_epochs 100

python scripts/train/train_mist.py \
    --data_dir /path/to/MIST \
    --batch_size 16 \
    --max_epochs 75

Train on a subset of stains:

python scripts/train/train_mist.py \
    --data_dir /path/to/MIST \
    --stains HER2 Ki67

python scripts/eval/eval_bci.py \
    --checkpoint checkpoints/bci/last.ckpt \
    --data_dir /path/to/BCI_dataset

python scripts/eval/eval_mist.py \
    --checkpoint checkpoints/mist_multistain/last.ckpt \
    --data_dir /path/to/MIST

Evaluate specific stains:

python scripts/eval/eval_mist.py \
    --checkpoint checkpoints/mist_multistain/last.ckpt \
    --data_dir /path/to/MIST \
    --stains HER2 Ki67

DAB p90 score: Mean intensity of pixels at or above the 90th percentile of the DAB channel (mean of top-10%), computed via Beer–Lambert color deconvolution.

├── src/
│   ├── models/
│   │   ├── generator.py         # SPADE-UNet generator
│   │   ├── blocks.py            # SPADE+FiLM, ResBlk, attention blocks
│   │   ├── discriminator.py     # Multi-scale PatchGAN
│   │   ├── trainer.py           # PyTorch Lightning training module
│   │   ├── uni_processor.py     # UNI feature extraction + processing
│   │   ├── edge_encoder.py      # Sobel edge encoder
│   │   └── losses.py            # Perceptual, adversarial, DAB losses
│   ├── data/
│   │   ├── bci_dataset.py       # BCI dataset loader
│   │   └── mist_dataset.py      # MIST multi-stain dataset loader
│   └── utils/
│       ├── dab.py               # DAB color deconvolution
│       └── metrics.py           # Evaluation metrics (FID, KID, LPIPS, DAB, IOD)
├── scripts/
│   ├── train/                   # Training scripts
│   └── eval/                    # Evaluation scripts
└── requirements.txt

This project is for research purposes only. UNI model weights are subject to their own license terms.