A leakage-free, long-tailed ensemble pipeline using frozen pathology foundation models + chunked XGBoost.
Classify 512×512 H&E histopathology patches into ten brain-tumor histologic sub-region categories for the BraTS 2026 Pathology (PATH) challenge. Three frozen self-supervised pathology encoders (UNI2-h, Virchow2, Phikon-v2) produce complementary patch embeddings; an ensemble of XGBoost classifiers trained on fixed-width chunks of the embedding space turns those embeddings into class probabilities, which are then refined with class-balanced weighting, rare-class boosting, and per-class threshold optimisation.
Out-of-fold accuracy 0.762 • MCC 0.713 • macro-F1 0.541 • AUC 0.788
Patient-grouped 5-fold StratifiedGroupKFold, pooled OOF, 1,042,640 training patches.
- Highlights
- The task & dataset
- Method overview
- Repository layout
- Installation
- Data preparation
- Training (end-to-end)
- Inference & submission
- Building the submission container (Docker / Apptainer)
- Results
- Running on an HPC / SLURM cluster
- Configuration reference
- Reproducibility & the leakage fix
- References
- Two-stage design — frozen foundation encoders (no fine-tuning) + gradient-boosted trees. Cheap to train, robust, and trivially parallelisable.
- Three heterogeneous encoders — UNI2-h (1536-d), Virchow2 (2560-d), Phikon-v2 (1024-d), trained with different objectives/corpora for complementary features.
- Chunked XGBoost — each embedding is split into 512-d chunks, one classifier per
(fold, model, chunk); predictions are averaged. 50 predictors in total. - Leakage-free validation — whole patients (never individual patches) are assigned to folds with
StratifiedGroupKFold. This corrects a severe defect in the naive split that leaked 124/126 patients across train/val and inflated metrics to ~0.999. - Long-tail handling — class-balanced sample weights, rare-class probability boosting, and grid-searched per-class decision thresholds.
- Offline, reproducible submission container — Docker and Apptainer recipes with foundation-model weights baked in (no network needed at evaluation time).
The BraTS-Path taxonomy has ten histologic sub-region classes. Nine are inherited from the 2024 edition; BraTS 2026 adds NOTA ("None of the Above") for tissue that matches no defined sub-region. The dataset is strongly long-tailed — the rarest class (DM) is ~14× rarer than the most common (CT).
| Id | Code — histologic sub-region | # patches | Share |
|---|---|---|---|
| 0 | CT — Cellular Tumor | 275,909 | 26.46 % |
| 1 | DM — Dense Macrophages | 19,596 | 1.88 % |
| 2 | IC — Infiltration into Cortex | 111,194 | 10.66 % |
| 3 | LI — Leptomeningeal Infiltration | 26,423 | 2.53 % |
| 4 | MP — Microvascular Proliferation | 53,248 | 5.11 % |
| 5 | NC — Necrosis | 214,884 | 20.61 % |
| 6 | PL — Presence of Lymphocytes | 21,113 | 2.02 % |
| 7 | PN — Pseudopalisading Necrosis | 92,109 | 8.83 % |
| 8 | WM — White Matter | 68,360 | 6.56 % |
| 9 | NOTA — None of the Above | 159,804 | 15.33 % |
| — | Total | 1,042,640 | 100 % |
Class ids follow
Dataset/class_map.jsonexactly. The submission CSV uses these integer ids (see Inference & submission). Rare classes (< 4 % share) — DM, LI, PL — are bold above and are the targets of rare-class boosting.
Example patches (one per class):
| CT | DM | IC | LI | MP |
|---|---|---|---|---|
 |
 |
 |
 |
 |
| NC | PN | WM | PL | NOTA |
 |
 |
 |
 |
 |
┌──────────────────────────────────────────────┐
512×512 H&E patch ──► │ Stage 1 — frozen pathology foundation models │
│ • UNI2-h ViT-h/14 → 1536-d (CLS) │
│ • Virchow2 ViT-h → 2560-d (CLS ⊕ mean)│
│ • Phikon-v2 ViT → 1024-d (CLS) │
└───────────────┬──────────────────────────────┘
│ embeddings cached to HDF5
▼
┌──────────────────────────────────────────────┐
│ Stage 2 — chunked XGBoost ensemble │
│ • split each embedding into 512-d chunks │
│ • 1 XGBoost per (fold, model, chunk) │
│ • average predict_proba over all 50 models │
└───────────────┬──────────────────────────────┘
▼
┌──────────────────────────────────────────────┐
│ Long-tail post-processing │
│ • class-balanced sample weights (train) │
│ • rare-class boost ×2.5 (DM, LI, PL) │
│ • per-class threshold optimisation (macro-F1)│
└───────────────┬──────────────────────────────┘
▼
argmax → class id (0–9)
Why these choices
- Frozen encoders — pathology foundation models already encode rich tissue morphology; freezing them avoids over-fitting on a long-tailed set and is far cheaper than end-to-end fine-tuning.
- Token pooling — UNI2-h and Phikon-v2 use the CLS token directly; Virchow2 concatenates its CLS token with the mean of its patch tokens (→ 2560-d).
- 512-d chunking — instead of one classifier on the full 5,120-d concatenation, each 512-d slice gets its own XGBoost (UNI2-h→3, Virchow2→5, Phikon-v2→2 chunks; zero-padded if the last chunk is short). Averaging the chunk predictors acts as an implicit regulariser.
- Rare-class boost — minority-class probabilities are scaled by 2.5 then re-normalised, a cheap training-free recall lever for DM/LI/PL.
- Threshold optimisation — per-class decision thresholds are grid-searched on pooled OOF predictions to maximise macro-F1.
A full write-up (architecture diagrams, ablations, per-fold tables) is in
report/brats_path_report.pdf.
BraTS2026_PATH/
├── README.md ← you are here
├── setup.py ← pip-installable package (entry points)
├── requirements.txt
├── codes/
│ ├── config.py ← central configuration (paths, models, hyperparams)
│ ├── foundation_models.py ← load encoders + extract embeddings
│ ├── data_loader.py ← manifest, dataset, dataloaders
│ ├── patient_split.py ← leakage-free patient-grouped CV folds
│ ├── build_manifest.py ← one-time dataset index
│ ├── extract_features.py ← extract train/val embeddings → HDF5
│ ├── extract_val_features.py ← extract OFFICIAL val embeddings → HDF5
│ ├── train_xgboost.py ← train one XGBoost per (model, chunk)
│ ├── run_cv.py ← full patient-grouped 5-fold CV (OOF)
│ ├── ensemble_predict.py ← average chunk predictors + rare-class boost
│ ├── optimize_thresholds.py ← per-class threshold grid search
│ ├── evaluate.py ← the 6 challenge metrics + reports
│ ├── predict_val.py ← OFFICIAL val → submission_val.csv
│ ├── inference.py ← predict on arbitrary images/dirs
│ ├── plot_training.py ← training/eval plots
│ ├── saved_models/ ← trained XGBoost fold models + thresholds.npy (committed)
│ │ ├── cv/fold_{0..4}/*.json
│ │ └── thresholds.npy
│ ├── results/ ← cv_metrics.json, plots, submission_val.csv (committed)
│ ├── features/ ← embedding HDF5s (git-ignored — regenerate)
│ └── *.sh ← SLURM submission scripts
├── docker/
│ ├── Dockerfile ← submission image (offline)
│ ├── brats.def ← Apptainer/Singularity equivalent
│ ├── infer.py ← container entrypoint (/input → /output/predictions.csv)
│ ├── build.sh / run.sh ← Docker build/run helpers
│ ├── build_apptainer.sh ← Apptainer build helper
│ ├── requirements-docker.txt
│ └── README.md ← container-specific docs
└── report/
├── brats_path_report.pdf ← full technical report
├── brats_path_report.tex
└── figures/ ← plots + sample patches used in this README
Not in git (regenerate locally): the raw Dataset/, the 21 GB of embedding
HDF5s in codes/features/, large .npy result arrays, and the 8.9 GB built SIF.
See .gitignore.
Requires Python ≥ 3.10 and a CUDA GPU for feature extraction (XGBoost training and inference can fall back to CPU).
git clone https://github.com/ujjwalbaid0408/BraTS2026_PATH.git
cd BraTS2026_PATH
# (recommended) a fresh environment
python -m venv .venv && source .venv/bin/activate
# or: conda create -n brats-path python=3.10 && conda activate brats-path
# install dependencies + the package (exposes brats-* console scripts)
pip install -e .
# or, just the deps: pip install -r requirements.txtFoundation-model access. UNI2-h and Virchow2 are gated on the Hugging Face Hub — request access on their model pages, then authenticate once:
huggingface-cli login # paste an HF token with the accepted licenses
# optional pre-download (otherwise downloaded lazily on first extract):
python - <<'PY'
from huggingface_hub import snapshot_download
for r in ["MahmoodLab/UNI2-h", "paige-ai/Virchow2", "owkin/phikon-v2"]:
snapshot_download(r)
PYCore dependencies: torch, torchvision, timm, transformers,
huggingface_hub, safetensors, xgboost, h5py, numpy, pandas,
scikit-learn, scipy, Pillow, matplotlib, seaborn, tqdm (full
versions in requirements.txt).
Download the BraTS-Path data from the challenge portal
and lay it out as config.py expects (override BASE_DIR in config.py if your
paths differ):
Dataset/
├── extracted/ ← {sample_id}.jpg + {sample_id}.cls (train patches + labels)
├── validation/images/ ← official val patches val_<hash>.jpg
├── class_map.json ← class name → id
└── BraTS-Path-2026-Train-Patch-Patient-Slide-Mapping.csv ← Name,Patient,Slide (for leakage-free splits)
Build the one-time manifest (scans extracted/ for .cls label files):
python codes/build_manifest.py # → codes/manifest.csvInspect/verify the patient-grouped folds (asserts no patient spans two folds):
python codes/patient_split.py # prints per-fold patient/patch/class countsThe pipeline is staged so every step is cached and resumable.
# all 3 active models × {train, val}; HDF5 written to codes/features/
python codes/extract_features.py --model all --split all
# or a single combination
python codes/extract_features.py --model uni2h --split trainEach output is codes/features/{model}_{split}.h5 with datasets
embeddings (N,D) float32, labels (N,) int32, sample_ids (N,) str. Existing
files are skipped, so re-runs are cheap.
run_cv.py is the recommended path: it runs the full patient-grouped 5-fold CV by
reusing the cached embeddings, writes every fold model, and reports pooled
out-of-fold metrics. Early stopping uses an inner fold carved from the training
folds, so the held-out fold is never seen during fitting.
python codes/run_cv.py # all models with embeddings
python codes/run_cv.py --models uni2h virchow2 phikon_v2 --device autoOutputs:
codes/saved_models/cv/fold_{0..4}/{model}_chunk_{NNN}.json— 50 fold modelscodes/saved_models/thresholds.npy— optimised per-class thresholdscodes/results/cv/cv_metrics.json— all metrics (argmax + thresholded + per-fold)codes/results/cv/plots/{oof_confusion_matrix,oof_per_class}.png
A simpler one-shot train/val path also exists:
python codes/train_xgboost.py --model all # trains on the patient-disjoint train split
python codes/ensemble_predict.py --split val # average chunk predictors (+ rare-class boost)
python codes/optimize_thresholds.py # grid-search per-class thresholds
python codes/evaluate.py --split val --thresh # report the 6 metricsFirst extract embeddings for the official (unlabeled) val patches, then predict by
averaging all 5 CV-fold models per (model, chunk):
python codes/extract_val_features.py # → codes/features/{model}_valofficial.h5
python codes/predict_val.py # → codes/results/submission_val.csvThe submission CSV is exactly:
SubjectID,Prediction
val_00002994ff324576,5
val_00007e874ce61d7c,9
...where Prediction is the integer class id (0=CT … 8=WM, 9=NOTA) and SubjectID
is the patch filename without extension.
python codes/inference.py --input /path/to/patch.jpg
python codes/inference.py --input /path/to/dir_of_jpgs/
# → CSV with sample_id, predicted_class, confidence, and per-class probabilitiesThe container reads every image in /input and writes a single
/output/predictions.csv — the exact challenge contract. Foundation-model weights
are baked in, so it runs fully offline (no HF token / network at eval time).
See docker/README.md for full details.
# weights must be in the local HF cache first (gated repos — see Installation)
cd docker
./build.sh # → brats2026-path:latest (~9 GB)
IMAGE=team/brats-path:v1 ./build.sh # custom tag
./run.sh /abs/path/to/input_images /abs/path/to/output_dir
# → /abs/path/to/output_dir/predictions.csvcd docker
./build_apptainer.sh # → docker/brats.sif
apptainer run --nv -B /abs/in:/input:ro -B /abs/out:/output brats.sifBoth recipes install the same deps, bake the same weights + trained fold
models, and run the same infer.py. Tunables (via -e / env):
BRATS_BATCH_SIZE (default 256), BRATS_NUM_WORKERS (default 8). GPU is strongly
recommended (≈114 k val patches through three ViTs); CPU-only also works (drop
--gpus all / --nv).
Patient-grouped 5-fold StratifiedGroupKFold, pooled out-of-fold (OOF) — 1,042,640 patches.
| Metric | OOF (argmax) | OOF (+ thresholds) | per-fold mean ± std |
|---|---|---|---|
| Accuracy | 0.7428 | 0.7618 | 0.750 ± 0.102 |
| AUC (macro) | 0.7885 | 0.7885 | 0.884 ± 0.065 |
| F1 (macro) | 0.5184 | 0.5411 | 0.477 ± 0.058 |
| F1 (weighted) | 0.7274 | 0.7455 | 0.716 ± 0.130 |
| MCC | 0.6900 | 0.7126 | 0.684 ± 0.112 |
| Sensitivity | 0.5124 | 0.5378 | 0.490 ± 0.042 |
| Specificity | 0.9694 | 0.9716 | 0.968 ± 0.012 |
Per-class OOF (argmax):
| Class | Sensitivity | Specificity | F1 |
|---|---|---|---|
| CT | 0.840 | 0.895 | 0.789 |
| DM | 0.278 | 0.989 | 0.299 |
| IC | 0.863 | 0.968 | 0.811 |
| LI | 0.023 | 0.969 | 0.021 |
| MP | 0.177 | 0.991 | 0.263 |
| NC | 0.907 | 0.949 | 0.862 |
| PL | 0.000 | 0.992 | 0.000 |
| PN | 0.541 | 0.968 | 0.578 |
| WM | 0.573 | 0.983 | 0.632 |
| NOTA | 0.923 | 0.989 | 0.929 |
The rare classes (DM, LI, PL, and MP) remain the hardest — they dominate the gap between weighted and macro F1. Threshold optimisation recovers ~+0.02 macro-F1 / MCC over plain argmax. Common, morphologically distinct classes (NC, NOTA, IC, CT) are recognised well.
| OOF confusion matrix | OOF per-class F1 / sensitivity |
|---|---|
 |
 |
Predicted class distribution vs. training prior (official val submission):
Numbers above are reproduced from codes/results/cv/cv_metrics.json.
Ready-to-edit SLURM batch scripts live in codes/ and docker/:
| Script | Purpose |
|---|---|
codes/submit_uni2h_extract_*.sh |
extract embeddings on GPU nodes (B200 / RTX PRO 6000) |
codes/submit_val_extract.sh |
extract official-val embeddings |
codes/submit_cv_cpu.sh |
run the 5-fold CV on CPU partitions (XGBoost is CPU-friendly) |
codes/submit_cv_pipeline.sh |
full CV pipeline |
codes/run_after_extract.sh |
chain CV + threshold opt after extraction finishes |
docker/build_apptainer.sbatch |
build the SIF on a compute node with egress |
docker/smoke_test.sbatch |
container smoke test on a few sample patches |
Useful environment overrides (see config.py): BRATS_NUM_WORKERS,
BRATS_BATCH_SIZE. Feature extraction is GPU-bound; XGBoost CV runs comfortably on
CPU partitions, which is handy when GPU queues are congested.
All knobs live in codes/config.py:
| Setting | Default | Meaning |
|---|---|---|
ACTIVE_MODELS |
FOUNDATION_MODELS_V2 |
UNI2-h + Virchow2 + Phikon-v2 (switch sets here) |
CV_FOLDS / CV_SEED |
5 / 42 | StratifiedGroupKFold config |
HOLDOUT_VAL_FOLD |
0 | which fold is "val" for the single-split path |
CHUNK_SIZE |
512 | embedding dims per XGBoost model |
XGBOOST_PARAMS |
depth 6, lr 0.1, 500 trees, hist | gradient-boosting hyperparameters |
RARE_THRESHOLD / RARE_BOOST_FACTOR |
0.04 / 2.5 | rare-class boost trigger & multiplier |
THRESHOLD_METRIC / THRESHOLD_STEPS |
macro_f1 / 50 | threshold grid search |
BATCH_SIZE / NUM_WORKERS |
512 / 8 | extraction (env-overridable) |
IMAGE_SIZE |
224 | encoder input resolution |
To add GigaPath (a 4th encoder), accept its HF license and uncomment its block
in FOUNDATION_MODELS_V2.
The original split hashed each patch to a shard via md5(patch_name) % NUM_SHARDS,
scattering patches from one slide/patient across both train and val. 124 of 126
patients and 248 of 255 slides appeared on both sides, so the model was validated
on near-identical neighbours of its own training patches — inflating metrics to
~0.999 (memorisation, not generalisation).
codes/patient_split.py is the single source of truth for a correct split:
StratifiedGroupKFold(n_splits=5, shuffle=True, random_state=42) grouped by
patient, so no subject ever appears in two folds. The assignment is cached to
Dataset/cv_folds.json and a hard assertion (_verify_no_overlap) fails the run
if any patient spans multiple folds. No embedding re-extraction is needed — the
per-patch embeddings are split-independent; only the partition was wrong.
Determinism: fixed seeds (CV_SEED=42, XGBOOST_PARAMS.random_state=42) and a
cached fold file mean run_cv.py reproduces the reported numbers.
- R. Chen et al. Towards a general-purpose foundation model for computational pathology (UNI). Nature Medicine, 2024. — MahmoodLab/UNI2-h
- E. Vorontsov et al. Virchow: A million-slide digital pathology foundation model. 2024. — paige-ai/Virchow2
- A. Filiot et al. Phikon / Phikon-v2: Scaling self-supervised learning for histopathology. Owkin, 2023–2024. — owkin/phikon-v2
- T. Chen and C. Guestrin. XGBoost: A scalable tree boosting system. KDD, 2016.
- M. Oquab et al. DINOv2: Learning robust visual features without supervision. 2023.
- F. Pedregosa et al. Scikit-learn: Machine learning in Python (StratifiedGroupKFold). JMLR, 2011.
- BraTS-Path Challenge. Patch-Level Classification of Histopathological Subregions in Glioblastoma. BraTS-Path 2026 Proceedings. — source of the official class-name key.
- Patch-Level Brain-Tumor Sub-region Classification Using Foundation Models Under Long-Tailed Data Distributions. BraTS-Path winning solution (method this pipeline builds on).
Built for the BraTS 2026 PATH challenge. The raw data is governed by the challenge's data-use agreement and is not redistributed here. Foundation-model weights belong to their respective authors under their own licenses. Please cite the references above if you build on this work.