SPIRiT-Diffusion: Self-Consistency Driven Diffusion Model for Accelerated MRI

Official code for the paper "SPIRiT-Diffusion: Self-Consistency Driven Diffusion Model for Accelerated MRI", published in TMI 2024.

by Zhuo-Xu Cui*, Chentao Cao*, Yue Wang*, Sen Jia, Jing Cheng, Xin Liu, Hairong Zheng, Dong Liang, and Yanjie Zhu⁺ (* denotes equal contribution, + denotes corresponding author).

Illustration

Abstract

Diffusion models are leading methods for image generation and have shown success in MRI reconstruction. However, existing diffusion-based methods primarily work in the image domain, making them vulnerable to inaccuracies in coil sensitivity maps (CSMs). While k-space interpolation methods address this issue, conventional diffusion models are not directly applicable to k-space interpolation. To tackle this, we introduce SPIRiT-Diffusion, a k-space interpolation diffusion model inspired by the iterative SPIRiT method. By utilizing SPIRiT’s self-consistent term (k-space physical prior), we formulate a novel stochastic differential equation (SDE) for the diffusion process, enabling k-space data interpolation. This approach highlights the importance of optimization models in designing SDEs, aligning the diffusion process with underlying physical principles, known as model-driven diffusion. Evaluated on a 3D joint intracranial and carotid vessel wall imaging dataset, SPIRiT-Diffusion outperforms image-domain methods, achieving high-quality reconstruction even at an acceleration rate of 10.

Adv

Checkout our

• 🔥🔥🔥TMI'24 Work HFS-SDE tailored specifically for MR reconstruction with the diffusion process in high-frequency space.

Setup

This section covers environment setup, data preparation, usage instructions, experiment weights, and a quick start guide. In the quick start, we provide experiment weights and synthetic data for easy model validation, without requiring lengthy data downloads. Parameter settings for generating high-quality reconstruction samples are also included.

Dependencies

Install necessary Python packages:

pip install -r requirements.txt

Note: For conda users, follow the setup instructions in previous work such as HFS-SDE.

Data Preparation

Example Dataset

We provide instructions to download the example data and place it in example_data, with the data class in utils/datasets.py/ExampleDataSet.

Preprocessing

In the image domain, crop the data to 320x320 and estimate sensitivity maps using the BART toolbox. The Auto-Calibration Signal (ACS) region is set to 48x48.

• Sensitivity Estimation: Use the BART toolbox to estimate sensitivity maps. Code is in data_prepare/generate_smaps_fastMRI.py. Code for generating Sum of Squares (SOS) coefficients is available in data_prepare/estimate_sos_map.py.

• Kernel Estimation: The code is located in data_prepare/estimate_kernel.py.

Usage

Train and evaluate models via main.py:

main.py:
  --config: Training configuration (default: 'None')
  --mode: <train|sample> Running mode: train or sample
  --workdir: Working directory

• config: Path to config file, formatted with ml_collections.
• workdir: Path to store experiment artifacts (e.g., checkpoints, samples, results).
• mode: "train" to start training, "sample" to run sampling.

Pretrained Checkpoints

Checkpoints are available here.

Quick Start Guide

Training

Use the training script train_fastMRI.sh:

• Set the following parameters based on the table below:

Parameter	Meaning	Example Value
training.estimate_csm	Type of coil sensitivity map	bart or sos

Train the model:

bash train_fastMRI.sh spirit

Sampling

Apply our model using phantom data for high-quality reconstruction. Set the sampling parameters as shown:

Parameter	Meaning	Example Value
sampling.snr	Adjusts noise level; higher values reduce noise, lower values reduce artifacts	0 - 1
sampling.mse	Controls overall predictor error; larger values increase the weight of the score	Non-negative
sampling.corrector_mse	Controls overall corrector error; larger values give more weight to the score	Non-negative
model.eta	Step size; affects image details and smoothness	0 - 1

Note: Prioritize adjusting model.eta to ensure it is within an appropriate range.

Run sampling:

bash test_fastMRI.sh spirit

When encountering a process lock, run:

bash rm_lock.sh

References

If you find the code useful, please cite:

@article{cui2024spirit,
  title={SPIRiT-Diffusion: Self-consistency driven diffusion model for accelerated MRI},
  author={Cui, Zhuo-Xu and Cao, Chentao and Wang, Yue and Jia, Sen and Cheng, Jing and Liu, Xin and Zheng, Hairong and Liang, Dong and Zhu, Yanjie},
  journal={IEEE Transactions on Medical Imaging},
  year={2024},
  publisher={IEEE}
}

Our implementation is based on Score-based SDE by Dr. Yang Song, with additional code from csgm-mri-langevin. Many thanks to their authors!