When Iteration Cannot Help: Single-Pass Missing-Modality MRI Synthesis and What Actually Limits It

git clone https://github.com/genmilab/PRIMED.git
cd PRIMED
conda create -n primed python=3.10 -y && conda activate primed
pip install -e .

The code needs Python 3.9 or later, PyTorch 2.1 or later, MONAI 1.4 or later and einops. It was tested with MONAI 1.4.0.

Download the checkpoints from Hugging Face:

pip install -U huggingface_hub
hf download adidukre/PRIMED --local-dir weights
cd weights && sha256sum -c SHA256SUMS

Each file holds the network weights only. The architecture is read from the weights, so no config is needed for inference.

BraTS. Download the BraTS 2023 adult glioma (GLI) training data from the challenge page and keep its folder layout:

/path/to/ASNR-MICCAI-BraTS2023-GLI-Challenge-TrainingData/
└── BraTS-GLI-00000-000/
    ├── BraTS-GLI-00000-000-t1n.nii.gz
    ├── BraTS-GLI-00000-000-t1c.nii.gz
    ├── BraTS-GLI-00000-000-t2w.nii.gz
    ├── BraTS-GLI-00000-000-t2f.nii.gz
    └── BraTS-GLI-00000-000-seg.nii.gz

Volumes must be co-registered, skull-stripped and on the 1 mm BraTS grid of 240 × 240 × 155 voxels, as BraTS releases them. Preprocessing runs on the fly: - clip to the 0.1 and 99.9 percentiles of the volume; - rescale to [0, 1]; - pad the depth from 155 to 160 slices; - crop the in-plane axes from 240 to 224.

IXI. Arrange the IXI dataset as /path/to/IXI/<subject>/<subject>-{T1,T2,PD}.nii.gz, with the three contrasts co-registered and resampled to the same 240 × 240 × 155, 1 mm grid. The IXI volumes are not skull-stripped.

splits/ holds the subject lists used in the paper:

splits/
├── brats_gli.json           # 1032 train / 219 val; 3 val subjects select checkpoints, the other 216 are the test set
├── brats_men_60.txt         # BraTS-MEN training release, first 60 subjects
├── brats_men_val_141.txt    # BraTS-MEN validation release
├── brats_ped_99.txt         # BraTS-PED training release
└── ixi.json                 # 477 train / 100 val; 4 val subjects select checkpoints, the other 96 are the test set

Synthesise a missing contrast from NIfTI files:

python synthesize.py \
  --weights weights/primed_brats.safetensors \
  --target t1c \
  --inputs t1n=/path/to/subject-t1n.nii.gz t2w=/path/to/subject-t2w.nii.gz t2f=/path/to/subject-t2f.nii.gz \
  --output /path/to/subject-t1c-synthesised.nii.gz

The output is on the grid of the inputs, with intensities in the normalised [0, 1] range. With primed_brats_avail.safetensors, any one or two acquired contrasts may be given; the empty input slots are filled with their mean.

import torch
from primed import load_model
from primed.data import load_volume

device = "cuda" if torch.cuda.is_available() else "cpu"
model = load_model("weights/primed_brats.safetensors", device)
inputs = torch.stack(
    [load_volume(f"/path/to/subject-{m}.nii.gz") for m in ("t1n", "t2w", "t2f")]
).unsqueeze(0).to(device)
t1c = model.synthesize(inputs.mean(dim=1), inputs, torch.tensor([1], device=device))

Score the 216 held-out BraTS-GLI subjects and print the per-target table:

python evaluate.py \
  --config configs/brats_primed.yaml \
  --weights weights/primed_brats.safetensors \
  --output results/primed_brats_test.jsonl

Per-case results are written as JSON lines: whole-volume PSNR on the 224 × 224 × 155 crop, PSNR inside the brain and inside each tumour region, and 3D SSIM with a 7³ window.

For example, the 60 BraTS-MEN subjects:

python evaluate.py --config configs/brats_primed.yaml \
  --weights weights/primed_brats.safetensors \
  --split all --data_root /path/to/BraTS-MEN-Train --subjects splits/brats_men_60.txt \
  --output results/primed_brats_men60.jsonl

For IXI, use --config configs/ixi_primed.yaml with primed_ixi.safetensors.

Set data_root, output_dir and, where used, init_weights in the config, then run from the repository root:

python train.py --config configs/brats_primed.yaml
torchrun --nproc_per_node=4 train.py --config configs/brats_primed.yaml

Training writes best.safetensors (best validation PSNR, raw or EMA weights), last.pt (state to resume from), and final.safetensors / final_ema.safetensors to output_dir. It needs about 38 GB of GPU memory per process at batch size 2.