By leveraging diffusion models, low-resolution MRI scans are upgraded to high-resolution images. The model is trained on image pairs, enabling it to learn noise reduction and generate clearer, more detailed MRI scans.
Code : https://github.com/chichonnade/MRI-Super-Resolution/tree/main
Demo : Low Resolution 1.5T to High Resolution 3T MRI
Overview
Latest diffusion models have shown promising results in super-resolution tasks. This project aims to leverage these models to enhance the resolution of MRI images. It involves training a diffusion model on HR / LR MRI pairs to learn the noise distribution and generate high-resolution images from low-resolution inputs. Using a cascade of these models, we can achieve significant improvements in image quality.
Forward noise process
Using a cosine noise schedule, we can gradually add noise to the input image, allowing the model to learn the noise distribution and generate high-resolution images. The noise variance is increased over 1000 steps with a maximum variance of 0.3 at the final step to prevent adding out-of-bound noise values.
Here are the formulas used:
Cosine Schedule for Alphas:
$$ \alpha_t = \cos^2\left(\frac{\pi t}{2T}\right) $$Cumulative Product of Alphas:
$$\bar{\alpha_{t}} = \prod_{i=1}^{t}\alpha_i$$Forward Noise Process:
$$ q(x_t | x_0) = \mathcal{N}(x_t; \sqrt{\bar{\alpha}_t} x_0, 0.3(1 - \bar{\alpha}_t) I) $$
\( x_t \) can be sampled as:
\[ x_t = \sqrt{\bar{\alpha}_t} x_0 + \sqrt{(1 - \bar{\alpha}_t)} \varepsilon_{t}, \text{ where } \varepsilon_{t} \sim \mathcal{N}(0, \mathbf{0.3I}) \]Where:
- \( \alpha_t \) represents the amount of signal retained at time step \( t \).
- \( \bar{\alpha}_t \) is the cumulative product of the alpha values up to time step \( t \).
- \( x_t \) is the noisy image at step \( t \).
- \( x_0 \) is the original image.
- \( q(x_t \mid x_0) \) represents the probability distribution of the noisy image given the original image.
LR/HR pairs
We use the Human Connectome Project (HCP) dataset, which contains both 1.5T and 3T MRI scans. The LR images are generated by downsampling the HR images using the spline interpolation method. Then we train the diffusion model on these LR/HR pairs to learn the noise distribution and generate high-resolution images.
Reverse process using a 3D-Unet
Our model is based on the original DDPM architecture with 3D convolutions to handle the 3D MRI data. The model consists of a series of diffusion steps, each with a series of 3D convolutions, normalization layers, and attention blocks.
Reverse Noise Schedule for Betas:
$$ \beta_t = \frac{1 - \alpha_{t+1}}{1 - \bar{\alpha}_{t+1}} $$Variance of the Reverse Process:
$$ \sigma_t^2 = \frac{1 - \bar{\alpha}_t}{1 - \bar{\alpha}_{t+1}} \beta_t $$Reverse Process Sampling:
$$ p(x_{t-1} \mid x_t, x_0) = \mathcal{N}(x_{t-1}; \mu_t(x_t, x_0), \sigma_t^2 I) $$Mean of the Reverse Process:
$$ \mu_t(x_t, x_0) = \frac{\sqrt{\bar{\alpha}_{t-1}} \beta_t}{1 - \bar{\alpha}_t} x_0 + \frac{\sqrt{\alpha_t}(1 - \bar{\alpha}_{t-1})}{1 - \bar{\alpha}_t} x_t $$
Where:
- \( \beta_t \) represents the noise variance for the reverse process at time step \( t \).
- \( \sigma_t^2 \) is the variance of the reverse process, which determines how much noise is reduced at each reverse step.
- \( x_t \) is the noisy image at step \( t \).
- \( x_{t-1} \) is the denoised image at the previous step \( t-1 \).
- \( \mu_t(x_t, x_0) \) is the mean of the reverse process, guiding the denoising from \( x_t \) to \( x_{t-1} \).
- \( p(x_{t-1} \mid x_t, x_0) \) represents the probability distribution of the denoised image at step \( t-1 \) given the noisy image at step \( t \) and the original image \( x_0 \).
Updates
Currently assessing the performance of two reverse process models. A 3D-Unet for best performance and a combination of 3 2D-Unets to reduce training cost.
Credits:
This project was conducted in collaboration with GENCI who provided the computational ressources. Authored by Hendrik Chiche, Ludovic Corcos and Logan Rouge.
