Deep learning model for predicting brain tumour contrast enhancement from non-contrast MRI sequences.
This repository contains the code and model weights for the paper:
"Predicting brain tumour enhancement from non-contrast MR imaging with artificial intelligence"
Ruffle JK, Mohinta S, Pombo G, Biswas A, Campbell A, Davagnanam I, Doig D, Hamman A, Hyare H, Jabeen F, Lim E, Mallon D, Owen S, Wilkinson S, Brandner S, Nachev P.
- 83% balanced accuracy
- 91.5% sensitivity
- 74.4% specificity
- Developed using 11,089 brain MRI studies (9,980 training + 1,109 held-out test) from 10 international datasets
- Validated across glioma, meningioma, metastases, and post-resection cases in both adult and paediatric populations
The model is trained to predict enhancing intracranial disease from T1-, T2-, and FLAIR-weighted imaging alone. Ground truth is derived from the full multimodal acquisition including T1c. Downstream evaluation covers radiologist benchmarking, equity calibration, and head-to-head model comparison.
Representative test-set cases that expert radiologists, when reviewing the non-contrast imaging only, wrongly judged to contain no post-contrast enhancement. The model's predicted enhancing-tumour mask is shown alongside ground truth.
Model performance in detecting enhancing tumour from non-contrast sequences across (a) the entire cohort, (b) individual datasets, (c) pathologies, (d) countries of origin, (e) patient ages, and (f) patient sexes.
Figures reproduced from the paper. Full captions and additional content are available in the preprint PDF.
pip install nnunetv2Download from Zenodo and extract:
# Download from Zenodo (link above) and extract to your nnUNet results folder
unzip enhancement_segmenter_weights.zip -d /path/to/nnUNet_results/export nnUNet_results="/path/to/nnUNet_results"Place your MRI sequences in a folder with this naming convention:
{subject_id}_0000.nii.gz→ FLAIR{subject_id}_0001.nii.gz→ T1 (non-contrast){subject_id}_0002.nii.gz→ T2
Try it on a public dataset: the BraTS challenge releases free, registration-only multi-modal brain tumour MRI (FLAIR, T1, T1c, T2). Use only FLAIR (
_0000), T1 (_0001), and T2 (_0002) — drop T1c, since this model predicts enhancement from non-contrast sequences only.
nnUNetv2_predict \
-d Dataset003_enhance_and_abnormality_batchconfig \
-i /path/to/input_folder \
-o /path/to/output_folder \
-f 0 1 2 3 4 \
-tr nnUNetTrainer_4000epochs \
-c 3d_fullres \
-p nnUNetResEncUNetPlans_80G \
-chk checkpoint_best.pthThe output segmentation will contain: 0=background, 1=brain, 2=non-enhancing abnormality, 3=predicted enhancing tumour.
- Framework: nnU-Net v2
- Architecture: ResidualEncoderUNet (3D fullres)
- Input: 3 non-contrast MRI sequences (FLAIR, T1, T2)
- Output: 4-class segmentation
- 0: Background
- 1: Brain parenchyma
- 2: Abnormality (non-enhancing lesion)
- 3: Enhancing tumour (predicted)
- Training: 5-fold cross-validation
- Test-set Dice distribution: 76.8% of patients with Dice > 0.3, 67.5% > 0.5, 50.2% > 0.7
- Python 3.9 or newer
- PyTorch 2.0 or newer (CUDA build matching your GPU driver)
- CUDA-capable GPU with ≥ 8 GB VRAM (inference)
- ≥ 48 GB GPU VRAM recommended for training (a single 48 GB GPU is sufficient for the configuration shipped here)
The released model was developed and validated under the following environment. The model is also expected to run on any compatible newer version, but the configuration below is what we have explicitly verified.
| Component | Version |
|---|---|
| Operating system | Ubuntu 22.04.5 LTS (Linux kernel 6.8.0, x86_64) |
| GPU hardware | NVIDIA RTX 6000 Ada Generation (48 GB VRAM) |
| NVIDIA driver | 580.105.08 |
| CUDA toolkit | 13.0 (build cu130) |
| cuDNN | 9.1.x |
| Python | 3.10.12 |
| PyTorch | 2.11.0 (+cu130 build) |
| nnU-Net | v2 (≥ 2.2) |
| numpy | ≥ 1.21 |
| pandas | ≥ 1.3 |
| nibabel | ≥ 3.2 |
| scipy | ≥ 1.7 |
| scikit-image | ≥ 0.19 |
| matplotlib | ≥ 3.5 |
| tqdm | ≥ 4.62 |
The Python dependencies are also pinned in requirements.txt. nnU-Net v2 itself transitively pulls in PyTorch and most heavyweight ML dependencies; the table above lists the upstream toolchain that was confirmed working on a clean install.
pip install nnunetv2export nnUNet_raw="/path/to/nnUNet_raw"
export nnUNet_preprocessed="/path/to/nnUNet_preprocessed"
export nnUNet_results="/path/to/nnUNet_results"Download the pretrained model weights from Zenodo:
After downloading, place the model files in the following structure:
$nnUNet_results/
└── Dataset003_enhance_and_abnormality_batchconfig/
└── nnUNetTrainer_4000epochs__nnUNetResEncUNetPlans_80G__3d_fullres/
├── fold_0/
│ └── checkpoint_best.pth
├── fold_1/
│ └── checkpoint_best.pth
├── fold_2/
│ └── checkpoint_best.pth
├── fold_3/
│ └── checkpoint_best.pth
├── fold_4/
│ └── checkpoint_best.pth
├── dataset.json
├── plans.json
├── dataset_fingerprint.json
└── crossval_results_folds_0_1_2_3_4/
├── postprocessing.pkl
├── postprocessing.json
├── plans.json
└── dataset.json
Prepare your input data in the nnU-Net format:
input_folder/
├── subject001_0000.nii.gz # FLAIR
├── subject001_0001.nii.gz # T1
├── subject001_0002.nii.gz # T2
├── subject002_0000.nii.gz # FLAIR
├── subject002_0001.nii.gz # T1
├── subject002_0002.nii.gz # T2
└── ...
Channel mapping:
_0000.nii.gz: FLAIR sequence_0001.nii.gz: T1-weighted sequence (non-contrast)_0002.nii.gz: T2-weighted sequence
nnUNetv2_predict \
-d Dataset003_enhance_and_abnormality_batchconfig \
-i /path/to/input_folder \
-o /path/to/output_folder \
-f 0 1 2 3 4 \
-tr nnUNetTrainer_4000epochs \
-c 3d_fullres \
-p nnUNetResEncUNetPlans_80G \
-chk checkpoint_best.pthThe model outputs a segmentation mask with the following labels:
| Label | Description |
|---|---|
| 0 | Background |
| 1 | Brain parenchyma |
| 2 | Non-enhancing abnormality |
| 3 | Predicted enhancing tumour |
-
Organise your raw data with the following structure:
- 4D NIFTI files containing [FLAIR, T1, T1CE, T2] sequences
- Enhancement masks (ground truth segmentations)
-
Split sequences using FSL tools:
# Example: Split 4D sequences into individual 3D volumes
fslsplit input_4d.nii.gz output_prefix_ -t
# Rename channels to match nnUNet format
# _0000 = FLAIR, _0001 = T1, _0002 = T2
# Note: T1CE (channel 3) is excluded as we predict enhancement from non-contrast onlySee split_sequences_abnormality_class.sh for a complete data preparation example.
Create a dataset.json file:
{
"channel_names": {
"0": "FLAIR",
"1": "T1",
"2": "T2"
},
"labels": {
"background": 0,
"brain": 1,
"abnormality": 2,
"ET": 3
},
"numTraining": 9980,
"file_ending": ".nii.gz"
}# Plan with ResEncL planner and 48GB GPU memory target
nnUNetv2_plan_and_preprocess -d 003 -pl nnUNetPlannerResEncL -np 32 --verify_dataset_integrity
# Generate experiment plan with custom memory target
nnUNetv2_plan_experiment -d 3 -pl nnUNetPlannerResEncL -gpu_memory_target 48 -overwrite_plans_name nnUNetResEncUNetPlans_48GTrain all 5 folds (multi-GPU example):
# Train folds 0-2 in parallel on 3 GPUs
CUDA_VISIBLE_DEVICES=0 nnUNetv2_train 003 3d_fullres 0 --npz -p nnUNetResEncUNetPlans_48G &
CUDA_VISIBLE_DEVICES=1 nnUNetv2_train 003 3d_fullres 1 --npz -p nnUNetResEncUNetPlans_48G &
CUDA_VISIBLE_DEVICES=2 nnUNetv2_train 003 3d_fullres 2 --npz -p nnUNetResEncUNetPlans_48G &
wait
# Train folds 3-4
CUDA_VISIBLE_DEVICES=0 nnUNetv2_train 003 3d_fullres 3 --npz -p nnUNetResEncUNetPlans_48G &
CUDA_VISIBLE_DEVICES=1 nnUNetv2_train 003 3d_fullres 4 --npz -p nnUNetResEncUNetPlans_48G &
waitSee train_nnunet_abnormality_class_bigger_batch_L_mod.sh for the exact training configuration used in the paper.
nnUNetv2_find_best_configuration 003 -p nnUNetResEncUNetPlans_48GA full model card following the Hugging Face model-card schema — covering intended use, out-of-scope use, biases and limitations, training procedure, evaluation, and technical specifications — is available at MODEL_CARD.md.
enhancement_segmenter/
├── README.md # This file
├── MODEL_CARD.md # Hugging Face-style model card
├── LICENSE # Apache 2.0 license
├── requirements.txt # Python dependencies
├── predict.py # Simplified inference wrapper
├── train.sh # Training script with exact parameters
└── split_sequences_abnormality_class.sh # Data preparation example
If you use this code or model in your research, please cite:
@article{ruffle2025predicting,
title={Predicting brain tumour enhancement from non-contrast MR imaging with artificial intelligence},
author={Ruffle, James K and Mohinta, Samia and Pombo, Guilherme and Biswas, Asthik and Campbell, Alan and Davagnanam, Indran and Doig, David and Hamman, Ahmed and Hyare, Harpreet and Jabeen, Farrah and Lim, Emma and Mallon, Dermot and Owen, Stephanie and Wilkinson, Sophie and Brandner, Sebastian and Nachev, Parashkev},
journal={arXiv preprint arXiv:2508.16650},
year={2025}
}Please also cite nnU-Net:
@article{isensee2021nnu,
title={nnU-Net: a self-configuring method for deep learning-based biomedical image segmentation},
author={Isensee, Fabian and Jaeger, Paul F and Kohl, Simon AA and Petersen, Jens and Maier-Hein, Klaus H},
journal={Nature methods},
volume={18},
number={2},
pages={203--211},
year={2021},
publisher={Nature Publishing Group}
}This work was supported by:
- European Society of Radiology (ESR) / European Institute for Biomedical Imaging Research (EIBIR) Seed Grant
- Medical Research Council (MR/X00046X/1)
- British Society of Neuroradiology
- National Brain Appeal
- Wellcome Trust (213038/Z/18/Z)
- UCLH NIHR Biomedical Research Centre
This project is licensed under the Apache License 2.0 - see the LICENSE file for details.
For questions about the code or the bundled data, please open an issue on GitHub or contact the corresponding author, Dr James K. Ruffle.