Conditional DDPM for Brain MRI Modality Translation (IXI)

A conditional denoising diffusion probabilistic model (DDPM) that translates brain MRI slices between T1, T2, and PD modalities, trained on the IXI dataset. Six directional models are included (t1_to_t2, t2_to_t1, t1_to_pd, pd_to_t1, t2_to_pd, pd_to_t2), each with a built-in uncertainty map — a pixel-wise confidence estimate for every generated image, produced from the spread of multiple stochastic samples rather than a fixed post-hoc heuristic.

Built as a course project for Advanced Deep Learning, focused equally on translation quality and on explainability (XAI) for a generative model in a medical-imaging setting.

How it works

Conditioning: the source modality slice is channel-concatenated onto the noisy target at every diffusion step (2-channel input: [x_noisy ‖ source]), so the U-Net has direct access to the source anatomy throughout the reverse process.
Sampling: supports full DDPM ancestral sampling and fast DDIM sampling (deterministic eta=0 or stochastic eta=1).
Uncertainty estimation: run the stochastic DDIM sampler N times on the same input; the pixel-wise standard deviation across runs is the uncertainty map (bright = model is unsure, dark = confident), and the pixel-wise mean is the reported prediction. No MC-dropout is involved — the trained models use dropout=0, so variability comes purely from sampling stochasticity.

Uncertainty overlay for T1→T2 (subject IXI019): the heatmap sits on top of the mean prediction, brightest at tissue boundaries and structures the model is least certain how to reconstruct.

Architecture


Backbone	Conditional U-Net, residual blocks + self-attention
Parameters	~242M per direction
Base channels	128, channel multipliers `[1, 2, 4, 8]`
Residual blocks / level	2
Attention resolutions	32×32, 16×16 (8 heads)
Time conditioning	Sinusoidal embedding → 2-layer MLP, injected as an additive shift in every ResBlock
Image size	256×256, single channel per modality
Diffusion steps (training)	`T=1000`, cosine beta schedule (Nichol & Dhariwal 2021)
Sampling	DDPM (exact, O(T)) or DDIM (O(S), S≪T; Song et al. 2021)

Full definition is in modeling/ (unet.py, diffusion.py, config.py) — a dependency-free copy of the training repo's model code, so this repo is runnable standalone.

Training

Data: IXI brain MRI, resampled to isotropic 256×256 2D slices, skull-stripped, registered per-subject across modalities, intensities normalised to [-1, 1].
Schedule: 50 epochs, effective batch size 48 (3× RTX A5000, batch 16/GPU), AdamW, lr 1e-4 with 500-step warmup, EMA decay 0.9999, mixed precision.
Each direction was trained independently (its own U-Net, not a shared multi-task model).

Validation metrics (epoch 50, held-out IXI subjects)

Direction	Val loss	SSIM	PSNR (dB)
T1 → T2	0.00614	0.883	25.43
T1 → PD	0.00590	0.827	24.90
T2 → T1	0.01141	0.813	22.57
T2 → PD	0.00343	0.772	29.14
PD → T1	0.01050	0.781	25.35
PD → T2	0.00334	0.813	27.93

Usage

pip install torch safetensors pyyaml numpy

import numpy as np
import torch
from inference import load_model, load_source

diffusion = load_model("t1_to_t2", device="cpu")

source = load_source("t1_slice.npy")          # (1, 1, 256, 256), float32 in [-1, 1]

with torch.no_grad():
    # Fast deterministic sample (~20-30 DDIM steps is a good quality/speed trade-off)
    prediction = diffusion.ddim_sample(source, num_steps=30, eta=0.0)

    # Or: prediction + pixel-wise uncertainty map from stochastic sampling
    mean, uncertainty = diffusion.sample_with_uncertainty(
        source, num_samples=10, ddim_steps=30, eta=1.0,
    )

See inference.py for a complete CLI example (including saving both the prediction and the uncertainty map).

Weights format

Each checkpoints/<direction>/ema_fp16.safetensors holds the EMA weights of the full GaussianDiffusion module (U-Net params + the diffusion buffers), cast to fp16 to keep the repo compact. Reported metrics above were computed with the original fp32 training checkpoint; fp16 inference should be within noise of those numbers but hasn't been re-benchmarked separately.

Compute cost

This is a pixel-space diffusion model (no latent compression), so sampling cost scales with DDIM steps × image count. On a desktop-class CPU (14-core/20-thread), one DDIM step takes roughly 2 seconds for a single 256×256 image — a 30-step sample is well under a minute. It also runs on GPU for much faster sampling. It's compute-bound, not memory-bound: the model + one image's activations fit comfortably under 4GB RAM/VRAM.

Limitations

Trained on a single-institution dataset (IXI: 3 London hospitals, specific scanner protocols) — not validated for generalisation to other scanners, field strengths, or populations.
Operates on individual 2D slices, not full 3D volumes; no explicit cross-slice consistency constraint.
This is a course project, not a clinical tool. Outputs should not be used for diagnosis or any clinical decision-making.
The uncertainty map reflects sampling variance under the trained model, not calibrated epistemic or aleatoric uncertainty in a statistical sense — treat it as a relative, qualitative confidence signal rather than a calibrated probability.

License and attribution

Model weights and code released under CC BY-NC 4.0. Trained on the IXI dataset, licensed under CC BY-SA 3.0 — please review and respect that license if you use this model or redistribute derivatives.

Citation

If you reference this project, please cite:

@misc{ixi-ddpm-mri-translation,
  title  = {Uncertainty-Aware Conditional Diffusion Model for Brain MRI Modality Translation},
  author = {Pheng, Samnang},
  year   = {2026},
  note   = {Advanced Deep Learning course project},
  url    = {https://huggingface.co/SamnangPheng/ixi-mri-diffusion-translation}
}

Downloads last month: -; Downloads are not tracked for this model. How to track