Physics-Informed Deep Learning for Geometric Distortion Correction in Prostate DWI
Distortion-guided restoration: a physics-informed learning framework to correct prostate diffusion MRI artifacts Ziyang Long, Nader Binesh, Lixia Wang, Archana Vadiraj Malaji, Chia-Chi Yang, Haoran Sun, Rola Saouaf, Timothy Daskivich, Hyung Kim, Yibin Xie, Debiao Li, Hsin-Jung Yang. Radiology Advances 3(4), 2026. doi:10.1093/radadv/umag031
Prostate multiparametric MRI (mpMRI) is the clinical gold standard for prostate cancer detection and PI-RADS grading. However, the diffusion-weighted imaging (DWI) component — acquired using single-shot echo-planar imaging (ssEPI) — is highly vulnerable to B0 field inhomogeneities, causing severe geometric distortions in the form of spatial warping, pixel pile-up, and signal dropout. These artifacts are dramatically worse in patients with hip prostheses or bowel distension, precisely the demographic most at risk for prostate cancer.
DGR addresses this without any additional scan acquisition. By learning to invert a physics-based forward distortion simulator, DGR corrects severe geometric distortions using only routinely acquired DWI and T2-weighted images.
Distorted ssEPI DWI ──────────────────────────────► Corrected DWI
(geometric warping, (anatomically
pixel pile-up, faithful,
signal dropout) diagnostic quality)
This repository ships both halves of the method — you can regenerate the training data, not just run the network:
| Half | What it does | Where |
|---|---|---|
| Forward — simulation | Takes a measured ΔB0 field, fits it, perturbs it into a family of physically plausible variants, converts each to a voxel displacement map from the EPI readout geometry, and warps undistorted DWI into distorted DWI. This manufactures the paired training data that does not exist clinically. | dgr/physics/, dgr/utils/, scripts/simulation/ |
| Reverse — restoration | Inverts that forward model: a CNN front-end produces the geometric correction, then a conditional diffusion module refines it on the anatomical manifold under T2w guidance. | dgr/models/, dgr/inference/, scripts/restoration/ |
- No extra acquisitions required — works with the standard clinical DWI + T2W protocol; no B0 field maps, no reverse phase-encoded scans
- Physics-informed training — forward ssEPI distortion simulator driven by 11 real B0 field maps from hip-prosthesis patients, augmented to 110 maps by perturbing the high-order terms of a 12th-order 2-D polynomial fit (>40,000 paired training samples)
- Hybrid CNN–Diffusion architecture — two-stage pipeline: coarse geometric correction via CNN, fine texture restoration via conditional diffusion refinement (SDEdit-style)
- T2W anatomical conditioning — uses the distortion-free T2W scan as an anatomical reference via deformable cross-attention
- Clinically validated — prospective cohort of 34 subjects with severe baseline distortion; blinded radiologist scoring shows significant improvement in geometric fidelity, image quality, and diagnostic confidence
── FORWARD (simulation) ──
DICOM / .mat ─► b0_field_read ─► b0_registration ─► ΔB0 in T2 space
│
generate_b0_variants_{poly,sh} ◄──┘
│ (12th-order polynomial coeff. perturbation)
▼
B0 variant fields
│
undistorted DWI (b=50, b=1400) ────┼──► generate_dwi_pairs ──► paired (distorted, GT) NPZ
│ (VDM → EPI forward splat)
── REVERSE (restoration) ──
▼
stage 1: train_stage1_cnn distorted DWI ──► geometry-corrected DWI
▼
stage 2: train_stage2_diffusion + T2w guidance ──► refined DWI
▼
infer_dgr ──► evaluate_distortion_correction
dgr/utils/warp.py :: compute_vdm_from_b0_2d_ESP turns a ΔB0 field in Hz into a voxel
displacement map given the echo spacing, phase-encode line count, partial-Fourier factor and
in-plane acceleration. dgr/utils/epi_warp.py :: forward_splat_with_fallback then applies it as a
conservative forward splat along the PE axis, so signal pile-up and stretching are both
modelled rather than approximated by an interpolating pull-warp.
A 2.5D multi-scale encoder–decoder with residual blocks processes the distorted DWI alongside co-registered T2W. A contrast-aware deformable cross-attention module treats distorted DWI as queries and T2W as keys/values, computing adaptive spatial offsets to bridge the geometric mismatch between the two modalities. Feature Pyramid Network (FPN)-style top-down aggregation recovers fine spatial details.
A conditional diffusion UNet refines the coarse CNN output, conditioned on the T2W image and on
the frozen stage-1 result. It is trained with prediction_type="sample" — the network predicts
the clean image rather than the noise, so learning happens on the anatomical manifold. Inference
uses SDEdit-style img2img initialization with DPM-Solver for fast sampling.
Distorted low-b DWI ─┐
ADC map ─┼─► [Stage 1: CNN Backbone] ──► Coarse-corrected DWI + ADC
Co-reg. T2W ─┘ (deformable │
cross-attn) │
▼
[Stage 2: Diffusion Refinement]
(T2W + CNN output conditioning,
SDEdit img2img, DPM-Solver)
│
▼
Final corrected DWI + ADC
│
▼
High-b DWI (derived from ADC)
git clone https://github.com/Albertlongzi/DGR
cd DGR
conda create -n dgr python=3.8 -y && conda activate dgr
pip install -r requirements.txt
pip install -e .The published results were produced with Python 3.8.20, PyTorch 2.4.1+cu121, diffusers 0.35.1,
NumPy 1.24.3. monai and torchmetrics are optional; the code falls back cleanly without them.
Weights live on Hugging Face under a gated repository — access is reviewed and approved manually, and the request form asks what you intend to use them for:
https://huggingface.co/Zylong/DGR
| Stage | File | Params | Size |
|---|---|---|---|
| 1 — CNN | stage1_cnn/stage1_cnn.safetensors |
32.1 M | 128 MB |
| 2 — diffusion | stage2_diffusion/stage2_diffusion.safetensors |
299.6 M | 1.20 GB |
hf auth login # required: the repository is gated
python - <<'PY'
from huggingface_hub import hf_hub_download
hf_hub_download("Zylong/DGR", "stage1_cnn/stage1_cnn.safetensors")
hf_hub_download("Zylong/DGR", "stage2_diffusion/stage2_diffusion.safetensors")
PYEach weight file ships with a config.json recording the architecture hyperparameters and the
noise-scheduler settings. scripts/restoration/infer_dgr.py accepts either a .safetensors or a
raw training .pt; given the former it reads the sibling config.json, which is where
prediction_type="sample" comes from — the diffusers default of "epsilon" would silently return
noise. To produce your own release files from a training checkpoint:
python tools/export_checkpoint.py --ckpt runs/stage2/diff_epoch_092.pt \
--out_dir hf_export/stage2 --kind diffusion --name stage2_diffusion# a) fit + perturb the measured B0 field into physically plausible variants
python scripts/simulation/generate_b0_variants_poly.py --help # 2-D polynomial perturbation
python scripts/simulation/generate_b0_variants_sh.py --help # spherical-harmonic, orders >= 3 only
# b) warp undistorted DWI through the forward EPI model to build training pairs
python scripts/simulation/generate_dwi_pairs.py \
--input_roots /path/to/preprocessed_local /path/to/preprocessed_fastmri /path/to/preprocessed_disease \
--b0_root /path/to/B0_variants_poly/order_12 \
--output_root /path/to/dwi_pair \
--max_b0_subjects_per_dwi 11 --smooth_sigma 1.5 --seed 123 --num_workers 6
# c) held-out test set, built the same way from a disjoint B0 pool
python scripts/simulation/generate_dwi_testset.py --helpTwo perturbation bases are provided. The released training pairs used the 12th-order 2-D
polynomial basis (generate_b0_variants_poly.py, fitted per slice), scaling the high-order
coefficients by factors of 0–0.15 to turn 11 measured field maps into 110. The spherical-harmonic
generator is included as an alternative; it perturbs only orders ≥ 3, leaving orders 0–2 intact so
the shim-correctable component stays fixed.
torchrun --nproc-per-node=6 scripts/restoration/train_stage1_cnn.py \
--npz_root /path/to/dwi_pair/pe_axis0 \
--npz_root2 /path/to/dwi_pair/pe_axis1 \
--out_dir runs/stage1_cnn \
--radius 2 --batch_size 6 --epochs 25 --lr 3e-4 --warmup_steps 3054 \
--base_channels 64 --latent_dim 8 --prompt_k 8 --prompt_temp 1.0 \
--use_ssim_loss --ssim_weight 0.25 --ms_w1 0.2 --ms_w2 0.05torchrun --nproc-per-node=4 scripts/restoration/train_stage2_diffusion.py \
--npz_root /path/to/dwi_pair/pe_axis0 \
--out_dir runs/stage2_diffusion \
--cnn_ckpt runs/stage1_cnn/mageultra_best.pt \
--batch_size 6 --epochs 100 --lr 1e-4 --warmup_steps 1000 \
--radius 2 --t2_cond_channels 64 --t2_contrast_mod none \
--val_interval 2 --val_steps 50 --val_strength 0.3 --use_dpm_solver_validationpython scripts/restoration/infer_dgr.py \
--cnn_ckpt stage1_cnn/stage1_cnn.safetensors \
--ckpt stage2_diffusion/stage2_diffusion.safetensors \
--test_root /path/to/preprocessed_test_npz \
--out_dir outputs/dgr_infer \
--steps 100 --strength 0.3 --eta 0.0 --sampler dpmsolver \
--radius 2 --t2_cond_channels 64 --b_low 50 --b_high 1400 \
--slice_mode all --save_npz --save_slices--strength is the SDEdit refinement strength. Use 0.1–0.3: 0.1 is the conservative
setting reported in the paper, and 0.3 is the default in this script and the value every final
inference run used. Above that range the diffusion prior starts inventing structure rather than
refining, so do not raise it without checking outputs against a reference.
python scripts/evaluation/evaluate_distortion_correction.py --helpReports PSNR / SSIM / NMSE / MAE (whole-FOV and prostate-centred) plus wall-clock timing, against FUGUE and TOPUP baselines.
SLURM job templates for all four steps are in slurm/; they take DGR_ROOT and DGR_DATA from
the environment and bake in no absolute paths. configs/*.yaml record the exact settings behind
the released checkpoints.
Table 2 of the paper. FUGUE and TOPUP are given the inputs they need — an acquired field map and a reverse phase-encode acquisition respectively — while DGR uses neither.
| Method | Contrast | PSNR (dB) ↑ | NMSE ↓ | SSIM ↑ |
|---|---|---|---|---|
| Baseline (no correction) | b = 50 | 17.67 ± 4.06 | 0.447 ± 0.355 | 0.564 ± 0.146 |
| Baseline (no correction) | ADC | 15.60 ± 2.51 | 0.351 ± 0.172 | 0.281 ± 0.115 |
| FUGUE + fieldmap | b = 50 | 17.92 ± 7.11 | 0.903 ± 1.334 | 0.770 ± 0.094 |
| FUGUE + fieldmap | ADC | 22.37 ± 2.51 | 0.073 ± 0.040 | 0.736 ± 0.065 |
| TOPUP + reverse PE | b = 50 | 22.45 ± 4.83 | 0.175 ± 0.181 | 0.705 ± 0.115 |
| TOPUP + reverse PE | ADC | 19.42 ± 2.92 | 0.157 ± 0.107 | 0.409 ± 0.140 |
| DGR (ours) | b = 50 | 23.88 ± 2.93 | 0.089 ± 0.049 | 0.706 ± 0.079 |
| DGR (ours) | ADC | 22.99 ± 1.97 | 0.062 ± 0.028 | 0.624 ± 0.107 |
Improvements in PSNR and NMSE were statistically significant against TOPUP and/or FUGUE across both contrasts (paired Wilcoxon signed-rank, all P < .001 except low-b PSNR vs TOPUP, P = .02). DGR leads on PSNR and NMSE for both contrasts; FUGUE retains the higher SSIM, though at a much worse low-b NMSE (0.903 ± 1.334) and with a field map DGR never sees.
| Criterion | Original ssEPI | DGR | p-value |
|---|---|---|---|
| Geometric fidelity | 2.6 | 3.3 | < 0.001 |
| Overall image quality | 2.5 | 2.9 | < 0.001 |
| Diagnostic confidence | 2.5 | 3.0 | < 0.001 |
- Zero false negatives and zero false positives in lesion analysis (n=18 with histopathology)
- Inference time: 13–15 seconds per subject on NVIDIA H100
DGR/
├── dgr/ # library
│ ├── physics/ # B0 field I/O, DICOM handling, SH & polynomial fitting, B0→T2 registration
│ ├── utils/ # forward EPI model: VDM computation, splat / warp / resample kernels
│ ├── models/ # phc_net → phc_e2e_mega_net → phc_e2e_mageultra_net (stage 1)
│ │ # diffusion_unet_diffusers (stage 2)
│ ├── data/ # paired dual-b NPZ dataset with 2.5-D slice stacking
│ ├── inference/ # DDIM / DDPM / DPM-Solver samplers with T2 + CNN conditioning
│ ├── losses/ # SSIM, ADC consistency, relative intensity, TV / Jacobian penalties
│ └── conditioning/ # T2W conditioning channel construction
├── scripts/
│ ├── simulation/ # forward-model entry points
│ ├── restoration/ # stage-1 / stage-2 training and DGR inference
│ └── evaluation/ # quantitative comparison against FUGUE / TOPUP
├── configs/ # YAML records of the settings behind the released checkpoints
├── slurm/ # portable SLURM job templates
└── tools/ # checkpoint export for Hugging Face
This work uses two datasets:
| Dataset | Examinations | Usage |
|---|---|---|
| fastMRI Prostate | 312 | Training / Test |
| In-house (Cedars-Sinai Medical Center) | 130 | Training / Test |
| Total | 442 | 408 training, 34 clinical test (5 fastMRI + 29 Cedars-Sinai) |
B0 field maps were acquired from 11 patients with unilateral or bilateral hip prostheses and augmented to 110 maps by perturbing the high-order terms of a 12th-order 2-D polynomial fit, driving the forward distortion simulator.
The clinical source data cannot be redistributed. The simulation half of this repository lets the
training pairs be regenerated from any DWI + T2w + ΔB0 source, including the public fastMRI
Prostate dataset. configs/simulation.yaml documents the expected NPZ keys.
@article{long2026dgr,
title = {Distortion-guided restoration: a physics-informed learning framework
to correct prostate diffusion MRI artifacts},
author = {Long, Ziyang and Binesh, Nader and Wang, Lixia and Malaji, Archana Vadiraj
and Yang, Chia-Chi and Sun, Haoran and Saouaf, Rola and Daskivich, Timothy
and Kim, Hyung and Xie, Yibin and Li, Debiao and Yang, Hsin-Jung},
journal = {Radiology Advances},
volume = {3},
number = {4},
year = {2026},
doi = {10.1093/radadv/umag031}
}Code is released under the MIT license. The pretrained weights are distributed separately under a research-only license — see the Hugging Face repository.
Not a medical device. Not cleared or approved for clinical use, diagnosis, or treatment planning by any regulatory body.
This work was supported by NIH grants R01NS121544, R01HL156818, R01HL165211, R01HL181091, and R43NS120795. We thank the Research Imaging Core (RIC) at Cedars-Sinai Medical Center, MRI Technologist Mike Ngo, Irene Lee, and nurses Catherine Ubaldo-Prado and Lee Hyae for their support in data acquisition.