BraTS-Path 2026: Rare-Class Conditional Routing with Foundation Model Representations

Official code repository for the MICCAI BraTS-Path 2026 challenge submission:

"Rare-Class Conditional Routing with Foundation Model Representations for Glioma Histologic Sub-region Classification"

Shubham Innani, Suhang You, Carla Pitarch-Abaigar, Dimitrios Makris, Spyridon Bakas

---

Overview

We systematically evaluate fully-supervised and semi-supervised training paradigms on top of four pathology foundation model backbones for 10-class H&E patch classification of glioma sub-regions. Our best system (BraTS-Path Score = 0.750, MCC = 0.858, F1 = 0.641) uses a class-conditional routing ensemble: LoRA-UNI2 handles rare classes {DM, LI, PL} and XGBoost-UNI2 handles all remaining classes.

Method	Backbone	MCC	F1	BPS
XGBoost	UNI2	0.847	0.475	0.661
LoRA	UNI2	0.822	0.580	0.701
CDMA (semi-sup)	H-Optimus-1	0.800	0.573	0.687
XGBoost ensemble (3 models)	UNI2+Virchow2+H-Opt1	0.865	0.500	0.683
LoRA ensemble (3 models)	UNI2+Virchow2+H-Opt1	0.838	0.607	0.723
Routing ensemble (Config 02)	LoRA-UNI2 + XGB-UNI2	0.858	0.641	0.750

---

Repository Structure

.
├── src/                        # Training and evaluation code
│   ├── main_v2.py              # Supervised MLP (frozen backbone + head)
│   ├── main_lora.py            # Supervised LoRA fine-tuning
│   ├── main_cdma.py            # Semi-supervised CDMA (3-head cross-distillation)
│   ├── main_cdma_lora.py       # Semi-supervised CDMA + LoRA
│   ├── extract_embeddings.py   # Extract frozen FM embeddings (labeled set)
│   ├── extract_unlabelled_emb.py # Extract FM embeddings (unlabeled WSI patches)
│   ├── train_xgb.py            # XGBoost training with chunking + rare-boost
│   ├── pseudo_label_xgb.py     # Centroid-based pseudo-labeling for XGBoost
│   ├── build_ensembles.py      # Class-routing ensemble construction
│   └── predict.py              # Inference / prediction CSV export
│
├── docker/                     # Docker submission (Config 02 — best system)
│   ├── Dockerfile
│   ├── requirements.txt
│   ├── run.py                  # Challenge entrypoint (stable, do not modify)
│   ├── src/
│   │   ├── inference.py        # High-level inference flow
│   │   ├── inference_dependencies.py  # Model loading, embedding, routing logic
│   │   └── webdataset_loader.py
│   └── scripts/
│       ├── 01_build_image.sh   # Build Docker image + save .tar archive
│       ├── 02_run_docker_image.sh  # Run built image on local data
│       └── 03_convert_tar_to_sif.sh  # Convert to Apptainer SIF
│
├── train.sh                    # SLURM dispatch for all training jobs
└── BraTS-Path-2026-Train-Patch-Patient-Slide-Mapping.csv

---

Data Download

Challenge page (Task 5): https://challenges.synapse.org/Challenges/DetailsPage/Task5?id=syn74274097

You must register for the challenge and accept the data use agreement on Synapse before data access is granted.

1. Install the Synapse client and log in

pip install synapseclient
synapse login --authToken YOUR_SYNAPSE_PAT

Get a Personal Access Token at: https://www.synapse.org/#!PersonalAccessTokens:

2. Download data

All data is distributed as WebDataset .tar shards. Each shard contains .jpg patch images (512×512 px, 20×, JPEG quality 90) and .cls label files (integer 0–9).

# Labeled training shards  (~1.6M patches, 40 shards)
synapse get -r syn74274097 --include "BraTS-Path2026-Train-TARS*" --downloadLocation data/train_shards/

# Unlabeled WSI large images (.tiff + foreground masks, 80 WSIs)
# Scripts for patching these into shards are provided by the organizers on Synapse.
synapse get -r syn74274097 --include "*Unlabeled*" --downloadLocation data/unlabeled/

# Validation shards (~114K patches, no labels released)
synapse get -r syn74274097 --include "BraTS-Path2026-Valid*" --downloadLocation data/val_shards/

A small debug set (3,200 patches randomly sampled from training) is available without registration at: https://www.synapse.org/Synapse:syn74488419

3. Class map

Label	Class
0	CT — Cellular Tumor
1	DM — Dense Macrophages
2	IC — Infiltration into Cortex
3	LI — Leptomeningeal Infiltration
4	MP — Microvascular Proliferation
5	NC — Geographic Necrosis
6	PL — Presence of Lymphocytes
7	PN — Pseudopalisading Necrosis
8	WM — Penetration into White Matter
9	NOTA — None of the Above

4. Expected layout after download

data/
├── train_shards/
│   └── shard-{000000..000039}.tar    # labeled training shards
├── unlabeled/
│   └── *.tiff                        # unlabeled WSIs for semi-supervised use
└── val_shards/
    └── *.tar                         # validation shards (Docker inference input)

The BraTS-Path-2026-Train-Patch-Patient-Slide-Mapping.csv file in this repo maps each training patch name to its patient ID (0–125) and slide ID (0–254), which we used for patient-stratified splits.

5. Loading shards with WebDataset

import webdataset as wds
from torchvision import transforms

transform = transforms.Compose([transforms.Resize((224, 224)), transforms.ToTensor()])

dataset = (
    wds.WebDataset("data/train_shards/shard-{000000..000039}.tar", shardshuffle=True)
    .shuffle(5000)
    .decode("pil")
    .to_tuple("__key__", "jpg", "cls")
    .map_tuple(
        lambda k: k.decode() if isinstance(k, bytes) else k,
        transform,
        lambda c: int(c.decode() if isinstance(c, bytes) else c),
    )
)

Note: Challenge data is subject to a data use agreement. Do not redistribute.

---

Environment Setup

conda create -n bratspath python=3.10
conda activate bratspath

pip install torch==2.3.1 torchvision==0.18.1 --index-url https://download.pytorch.org/whl/cu121

pip install timm==1.0.12 scikit-learn xgboost peft huggingface_hub \
            webdataset tqdm tensorboard transformers numpy

---

Foundation Models

We evaluate four pathology foundation model backbones:

Backbone	Dim	Source
H-Optimus-1	1536	Bioptimus (HuggingFace: bioptimus/H-optimus-1)
UNI2-h	1536	MahmoodLab (HuggingFace: MahmoodLab/UNI2-h)
Virchow2	2560	Paige (HuggingFace: paige-ai/Virchow2)
DINOv3-GBM	1024	In-house, pretrained on TCGA-GBM/LGG (~2K WSIs)

All pan-pathology models download automatically at first run. For dinov3_custom, place weights at:

dinov3_weights/gbm_lgg_all/clean_dino_state_dict.pth

---

Feature Extraction

Extract frozen embeddings for XGBoost training:

# Labeled training set
python src/extract_embeddings.py \
    --model uni2 \
    --data_dir /path/to/BraTS-Path2026-Train-JPG \
    --output_dir embeddings/uni2

# Unlabeled WSI patches (for pseudo-labeling / CDMA)
python src/extract_unlabelled_emb.py \
    --model uni2 \
    --unlabeled_dir /path/to/unlabeled_patches \
    --output_dir embeddings/uni2_unlabeled

---

Training

1. Supervised MLP (frozen backbone + linear head)

python src/main_v2.py --model uni2 --epochs 100 --seed 42

2. Supervised LoRA

python src/main_lora.py \
    --model uni2 \
    --use_lora --lora_rank 4 --lora_alpha 8 \
    --epochs 100 --min_epochs 50 --patience 15 \
    --lr 5e-5 --backbone_lr 5e-6 \
    --seed 42

3. Semi-supervised CDMA (3-head cross-distillation)

CDMA uses 649,670 unlabeled patches extracted from the 80 provided WSIs. Three identical MLP heads operate on the backbone embedding; multiplicative noise U[−0.3, 0.3] is applied to heads 2 and 3 during training. Cross-knowledge distillation via symmetric pairwise KL divergence (T=10) runs over all six ordered head pairs on both labeled and unlabeled samples.

python src/main_cdma.py \
    --model uni2 \
    --unlabeled_dir /path/to/unlabeled_patches \
    --seed 42

4. Semi-supervised CDMA + LoRA

python src/main_cdma_lora.py \
    --model uni2 \
    --unlabeled_dir /path/to/unlabeled_patches \
    --lora_rank 4 --lora_alpha 8 \
    --seed 42

5. XGBoost (chunked training + rare-boost + threshold calibration)

Embeddings must be extracted first (see Feature Extraction above).

The training set is split into 10 stratified chunks. Rare classes (count < 10,000: DM, LI, PL) are replicated into every chunk so all 10 models see rare-class examples. Per-chunk class weights are set to inverse class frequency. After averaging 10-chunk probabilities, a rare-boost multiplier (max_count / count_c)^0.3 is applied and renormalized. Per-class thresholds are then calibrated by coordinate ascent (3 rounds, 30 grid points in [0.2, 5.0]) to maximize macro-MCC.

python src/train_xgb.py \
    --embeddings_dir embeddings/uni2 \
    --output_dir output_xgb/uni2 \
    --n_chunks 10

6. Centroid-based pseudo-labeling (exploratory)

Assigns pseudo-labels to unlabeled patches based on cosine distance to per-class centroids of training embeddings. Evaluated at the 35th, 60th, and 80th confidence percentiles. Note: this was a negative result — all thresholds degraded XGBoost performance relative to the supervised baseline.

python src/pseudo_label_xgb.py \
    --embeddings_dir embeddings/uni2 \
    --unlabeled_embeddings embeddings/uni2_unlabeled \
    --output_dir output_pseudo_label \
    --percentile 35   # or 60, 80

---

Ensemble: Class-Routing (Config 02 — Best System)

Routes rare classes {DM=1, LI=3, PL=6} to LoRA-UNI2 and all other classes to XGBoost-UNI2. If either model predicts a rare class for a patch, the LoRA prediction is used.

python src/build_ensembles.py \
    --lora_preds output_lora/uni2/predictions.csv \
    --xgb_preds  output_xgb/uni2/predictions.csv \
    --output     submissions/ensemble02_predictions.csv

Full routing sweep (all 12 configurations):

python src/build_ensembles.py --sweep

---

SLURM / HPC Training

train.sh dispatches to the correct script based on three positional arguments: <model> <lora|nolora> <supervised|semi>

sbatch --job-name=uni2_lora    train.sh uni2 lora   supervised
sbatch --job-name=uni2_cdma    train.sh uni2 nolora semi
sbatch --job-name=hopt_lora    train.sh h_optimus_1 lora supervised

---

Docker Inference (Challenge Submission)

The Docker image packages Config 02 (LoRA-UNI2 + XGB-UNI2 routing ensemble) for offline inference on WebDataset .tar shards.

Architecture

• Input: WebDataset .tar shards mounted at /input
• Output: predictions.csv at /output
• Pipeline: extract UNI2 embeddings → XGBoost probabilities → LoRA classifier logits → rare-class routing → write CSV

Build

cd docker
IMAGE_NAME=brats_ensemble02 IMAGE_TAG=latest ./scripts/01_build_image.sh

The build saves docker_archives/brats_ensemble02_latest.tar (~6.4 GB).

Run locally

INPUT_DIR=/path/to/input_shards \
OUTPUT_DIR=/path/to/output \
IMAGE_NAME=brats_ensemble02 IMAGE_TAG=latest \
./scripts/02_run_docker_image.sh

To override batch size (default 1536):

docker run --gpus all -e BATCH_SIZE=1024 ...

Convert to Apptainer SIF (for cluster validation)

IMAGE_NAME=brats_ensemble02 IMAGE_TAG=latest ./scripts/03_convert_tar_to_sif.sh

apptainer exec --nv \
  --bind /path/to/input:/input:ro \
  --bind /path/to/output:/output \
  sif/brats_ensemble02_latest.sif \
  python run.py --input /input --output /output

Adding checkpoints

Model weights are not included in this repository (.pth files are git-ignored). Place them before building:

docker/src/ckpts/
    lora_uni2_best.pth          # LoRA-UNI2 classifier checkpoint
    xgb_chunk_models.pkl        # XGBoost 10-model ensemble
    xgb_thresholds.npy          # Per-class decision thresholds

docker/src/foundation_model_weights/
    uni2_model.safetensors      # UNI2-h backbone weights (offline copy)

---

Key Hyperparameters

LoRA

Parameter	Value
Rank	4
Alpha	8
Dropout	0.05
Head LR	5×10⁻⁵
Backbone LR	5×10⁻⁶
Schedule	OneCycleLR, 5% warm-up
Min epochs	50
Early stopping patience	15
Weight decay	0.05
Label smoothing	0.1
Seed	42

XGBoost (per chunk)

Parameter	Value
n_estimators	300
max_depth	6
learning_rate	0.05
subsample	0.8
colsample_bytree	0.4
min_child_weight	5
gamma	0.1
objective	multi:softprob
Rare-boost exponent	0.3
Threshold calibration	Coordinate ascent, 3 rounds, 30 pts/class in [0.2, 5.0]

CDMA

Parameter	Value
Heads	3 × (Linear→GELU→Dropout(0.5)→Linear)
Noise (heads 2,3)	U[−0.3, 0.3] multiplicative
KD temperature	10
Contrastive temperature	0.07
Projection head	embed → 256 → 128
Epochs	50
Seed	42

---

Citation

If you use this code, please cite:

@inproceedings{innani2026bratspath,
  title     = {Rare-Class Conditional Routing with Foundation Model Representations
               for Glioma Histologic Sub-region Classification},
  author    = {Innani, Shubham and You, Suhang and Pitarch-Abaigar, Carla
               and Makris, Dimitrios and Bakas, Spyridon},
  booktitle = {MICCAI BraTS-Path 2026},
  year      = {2026}
}