MRI-DiffRecon is a research framework for accelerated MRI reconstruction using score-based diffusion models. By treating MRI reconstruction as a conditional inverse problem, we leverage the powerful prior learned by score-based generative models to recover high-quality images from heavily undersampled k-space measurements — achieving 4x–8x scan acceleration without perceptual quality loss.

This repository provides complete implementations of:

• Score-based diffusion reconstruction (VP-SDE and VE-SDE formulations)
• Noise-conditioned score network (NCSN++) with complex-valued support
• Data consistency projection integrated into the reverse diffusion process
• Predictor-corrector sampling with adaptive step sizes
• Competitive baselines: cascaded U-Net, compressed sensing (TV + wavelet)
• Comprehensive evaluation: SSIM, PSNR, NMSE, LPIPS, FID, and clinical quality metrics

Standard MRI acquisitions require patients to remain motionless in a scanner for 20–60 minutes per session. This creates significant challenges:

By undersampling k-space (acquiring fewer frequency measurements) and using deep learning reconstruction, we can reduce scan times by 4x–16x while preserving the diagnostic quality required for clinical decision-making. At 4x acceleration, a 40-minute knee protocol becomes 10 minutes; at 8x, it becomes 5 minutes — transforming patient experience and enabling high-throughput clinical workflows.

This work is directly relevant to the active commercialization of AI reconstruction in clinical MRI:

• Siemens Healthineers — Deep Resolve (generative AI reconstruction, SNR boosting)
• GE HealthCare — AIR Recon DL (deep learning reconstruction, FDA 510k cleared)
• Philips — SmartSpeed (AI-powered acceleration, up to 4x speed improvement)

Diffusion models represent the next generation beyond current CNN-based approaches, offering principled uncertainty quantification and superior perceptual quality at high acceleration factors.

Undersampled k-space  ──► Zero-filled IFFT ──► Aliased image x₀
                                                       │
                              ┌────────────────────────┘
                              ▼
                    ┌─────────────────────┐
                    │   Score Network     │  ← Noise level σ(t)
                    │   (NCSN++ U-Net)    │
                    │  Complex-valued     │
                    │  + Time embedding   │
                    └────────┬────────────┘
                             │  ∇_x log p_σ(x)
                             ▼
                    Reverse SDE Step
                    (Euler-Maruyama / PC)
                             │
                             ▼
                    ┌─────────────────────┐
                    │  Data Consistency   │  ← k-space measurements y
                    │  Projection         │
                    │  x ← x - λ·A†(Ax-y)│
                    └────────┬────────────┘
                             │
                         (iterate T steps)
                             │
                             ▼
                    High-quality reconstruction x̂

1. Score Network (NCSN++): A U-Net backbone conditioned on noise level via sinusoidal time embeddings and FiLM (Feature-wise Linear Modulation) layers. Operates on complex-valued MRI data by treating real and imaginary channels independently with shared weights.

2. Diffusion Process: We implement both VP-SDE (variance-preserving) and VE-SDE (variance-exploding) stochastic differential equations. The forward process gradually adds noise; the reverse process denoises while maintaining fidelity to measured k-space data.

3. Data Consistency: At each reverse diffusion step, a gradient step projects the current estimate back toward the measured k-space data: x ← x - λ · A†(Ax - y), where A is the undersampling operator (FFT + mask) and λ is a consistency weight.

4. Predictor-Corrector Sampling: Combines a numerical SDE solver (predictor) with Langevin MCMC correction steps, dramatically improving sample quality over pure ancestral sampling.

See RESULTS.md for full ablation studies, inference timing, and clinical quality analysis.

# Clone the repository
git clone https://github.com/yourusername/mri-diffusion-recon.git
cd mri-diffusion-recon

# Create conda environment (recommended)
conda create -n mri-diffusion python=3.9
conda activate mri-diffusion

# Install PyTorch (CUDA 11.8)
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

# Install project dependencies
pip install -e .

# Verify installation
python -c "import src; print('MRI-DiffRecon installed successfully')"

Download the fastMRI dataset from fastmri.org. This requires registering and agreeing to the data use agreement.

# Set data directory
export FASTMRI_DATA=/path/to/fastmri

# Verify dataset structure
python scripts/verify_data.py --data_dir $FASTMRI_DATA

# Expected structure:
# $FASTMRI_DATA/
#   knee_singlecoil_train/
#   knee_singlecoil_val/
#   knee_singlecoil_test/
#   brain_multicoil_train/
#   brain_multicoil_val/

Download from the Calgary-Campinas Public Brain MR Dataset.

export CC_DATA=/path/to/calgary-campinas

# Train on fastMRI knee, 4x acceleration
python scripts/train.py \
    --config configs/fastmri_knee_config.yaml \
    --output_dir runs/knee_4x_diffusion \
    --acceleration 4 \
    --gpus 4

# Train on fastMRI brain, multi-coil
python scripts/train.py \
    --config configs/fastmri_brain_config.yaml \
    --output_dir runs/brain_4x_diffusion \
    --acceleration 4 \
    --gpus 8

python scripts/train.py \
    --config configs/fastmri_knee_config.yaml \
    --model unet \
    --output_dir runs/knee_4x_unet

Training logs are automatically saved to TensorBoard:

tensorboard --logdir runs/

# Reconstruct a single volume
python scripts/reconstruct.py \
    --checkpoint runs/knee_4x_diffusion/best_model.pt \
    --input data/sample_knee.h5 \
    --output results/reconstruction.h5 \
    --acceleration 4 \
    --num_steps 1000 \
    --method pc  # predictor-corrector

# Fast inference with DDIM (50 steps)
python scripts/reconstruct.py \
    --checkpoint runs/knee_4x_diffusion/best_model.pt \
    --input data/sample_knee.h5 \
    --output results/fast_reconstruction.h5 \
    --num_steps 50 \
    --method ddim

# Batch reconstruction
python scripts/reconstruct.py \
    --checkpoint runs/knee_4x_diffusion/best_model.pt \
    --input_dir data/test_volumes/ \
    --output_dir results/reconstructions/ \
    --batch_size 4

# Full evaluation against all baselines
python scripts/evaluate.py \
    --predictions results/reconstructions/ \
    --ground_truth data/test_volumes/ \
    --output_dir results/metrics/ \
    --metrics ssim psnr nmse lpips fid

# Clinical quality assessment
python scripts/evaluate.py \
    --predictions results/reconstructions/ \
    --ground_truth data/test_volumes/ \
    --clinical_quality \
    --pathology_map data/annotations/

mri-diffusion-recon/
├── src/
│   ├── models/
│   │   ├── score_network.py      # NCSN++ score estimation network
│   │   ├── diffusion_mri.py      # VP-SDE / VE-SDE diffusion process
│   │   ├── unet_baseline.py      # U-Net and cascaded U-Net baselines
│   │   └── compressed_sensing.py # TV + wavelet CS baselines
│   ├── data/
│   │   ├── fastmri_dataset.py    # fastMRI dataset loader
│   │   └── kspace_transforms.py  # FFT, masking, coil combination
│   ├── training/
│   │   └── train_score.py        # Denoising score matching trainer
│   ├── inference/
│   │   └── reconstruct.py        # PC sampling + data consistency
│   └── evaluation/
│       ├── mri_metrics.py        # SSIM, PSNR, NMSE, LPIPS, FID
│       └── clinical_quality.py   # Clinical assessment framework
├── configs/
│   ├── fastmri_knee_config.yaml
│   └── fastmri_brain_config.yaml
├── scripts/
│   ├── train.py
│   ├── reconstruct.py
│   └── evaluate.py
├── docs/
│   └── MRI_RECONSTRUCTION.md    # Background on MRI physics and k-space
├── notebooks/                   # Jupyter demo notebooks
├── tests/                       # Unit tests
├── RESULTS.md                   # Detailed experimental results
├── requirements.txt
└── setup.py

This project builds upon the following key works:

@inproceedings{chung2022score,
  title={Score-based diffusion models for accelerated {MRI}},
  author={Chung, Hyungjin and Ye, Jong Chul},
  booktitle={Medical Image Analysis},
  year={2022}
}

@inproceedings{song2021score,
  title={Score-Based Generative Modeling through Stochastic Differential Equations},
  author={Song, Yang and Sohl-Dickstein, Jascha and Kingma, Diederik P and Kumar, Abhishek and Ermon, Stefano and Poole, Ben},
  booktitle={ICLR},
  year={2021}
}

@inproceedings{jalal2021robust,
  title={Robust Compressed Sensing {MRI} with Deep Generative Priors},
  author={Jalal, Ajil and Arvinte, Marius and Daras, Giannis and Price, Eric and Dimakis, Alexandros G and Tamir, Jon I},
  booktitle={NeurIPS},
  year={2021}
}

@article{zbontar2018fastmri,
  title={fast{MRI}: An Open Dataset and Benchmarks for Accelerated {MRI}},
  author={Zbontar, Jure and Knoll, Florian and Sriram, Anuroop and others},
  journal={arXiv preprint arXiv:1811.08839},
  year={2018}
}

@article{souza2018open,
  title={An Open, Multi-Vendor, Multi-Field-Strength Brain {MR} Dataset and Analysis of Publicly Available Skull Stripping Methods},
  author={Souza, Roberto and Lucena, Oeslle and others},
  journal={NeuroImage},
  year={2018}
}

If you use this code in your research, please cite:

@misc{mri-diffrecon2024,
  title={{MRI-DiffRecon}: Score-Based Diffusion Models for Accelerated {MRI} Reconstruction},
  author={Anonymous},
  year={2024},
  url={https://github.com/yourusername/mri-diffusion-recon}
}

This project is licensed under the MIT License — see LICENSE for details.

• fastMRI dataset provided by NYU Langone Health and Facebook AI Research
• Score-SDE codebase by Yang Song (Stanford)
• Calgary-Campinas dataset provided by the University of Calgary