Model-based iterative reconstruction

This example demonstrates how to reconstruct image from non-Cartesian k-space data with a regularization prior, using deepinv.

Imports

import numpy as np
import matplotlib.pyplot as plt
from brainweb_dl import get_mri
from deepinv.optim.prior import WaveletPrior
from deepinv.optim.data_fidelity import L2
from deepinv.optim.optimizers import optim_builder

from mrinufft import get_operator
from mrinufft.trajectories import initialize_3D_cones
import torch
import os

BACKEND = os.environ.get("MRINUFFT_BACKEND", "cufinufft")

Get MRI data, 3D FLORET trajectory, and simulate k-space data

samples_loc = initialize_3D_cones(32 * 32, Ns=256, nb_zigzags=16, width=3)
# Load and downsample MRI data for speed
mri = (
    torch.Tensor(np.ascontiguousarray(get_mri(0)[::2, ::2, ::2][::-1, ::-1]))
    .to(torch.complex64)
    .to("cuda")
)

Simulate k-space data

fourier_op = get_operator(BACKEND)(
    samples_loc,
    shape=mri.shape,
    density="pipe",
)
y = fourier_op.op(mri)  # Simulate k-space data
noise_level = y.abs().max().item() * 0.0002
y += noise_level * (torch.randn_like(y) + 1j * torch.randn_like(y))

/volatile/github-ci-mind-inria/gpu_mind_runner/_work/mri-nufft/mri-nufft/src/mrinufft/_utils.py:67: UserWarning: Samples will be rescaled to [-pi, pi), assuming they were in [-0.5, 0.5)
  warnings.warn(
/volatile/github-ci-mind-inria/gpu_mind_runner/_work/mri-nufft/mri-nufft/src/mrinufft/_utils.py:72: UserWarning: Samples will be rescaled to [-0.5, 0.5), assuming they were in [-pi, pi)
  warnings.warn(
/volatile/github-ci-mind-inria/gpu_mind_runner/_work/mri-nufft/mri-nufft/src/mrinufft/_array_compat.py:248: UserWarning: data is on gpu, it will be moved to CPU.
  warnings.warn("data is on gpu, it will be moved to CPU.")

Setup the physics and prior

physics = fourier_op.make_deepinv_phy()
wavelet = WaveletPrior(
    wv="sym8",
    wvdim=3,
    level=3,
    is_complex=True,
)

Initial reconstruction with adjoint

x_dagger = physics.A_dagger(y)

  0%|          | 0/100 [00:00<?, ?it/s]
  4%|▍         | 4/100 [00:00<00:02, 35.64it/s]
  8%|▊         | 8/100 [00:00<00:02, 35.77it/s]
 12%|█▏        | 12/100 [00:00<00:02, 35.80it/s]
 16%|█▌        | 16/100 [00:00<00:02, 35.74it/s]
 20%|██        | 20/100 [00:00<00:02, 35.69it/s]
 24%|██▍       | 24/100 [00:00<00:02, 35.69it/s]
 28%|██▊       | 28/100 [00:00<00:02, 35.60it/s]
 32%|███▏      | 32/100 [00:00<00:01, 35.62it/s]
 36%|███▌      | 36/100 [00:01<00:01, 35.70it/s]
 40%|████      | 40/100 [00:01<00:01, 35.69it/s]
 44%|████▍     | 44/100 [00:01<00:01, 35.69it/s]
 48%|████▊     | 48/100 [00:01<00:01, 35.70it/s]
 52%|█████▏    | 52/100 [00:01<00:01, 29.25it/s]
 56%|█████▌    | 56/100 [00:01<00:01, 27.57it/s]
 59%|█████▉    | 59/100 [00:01<00:01, 23.03it/s]
 62%|██████▏   | 62/100 [00:02<00:01, 21.26it/s]
 65%|██████▌   | 65/100 [00:02<00:01, 18.65it/s]
 68%|██████▊   | 68/100 [00:02<00:01, 17.14it/s]
 70%|███████   | 70/100 [00:02<00:01, 16.33it/s]
 72%|███████▏  | 72/100 [00:02<00:01, 15.57it/s]
 74%|███████▍  | 74/100 [00:02<00:01, 15.27it/s]
 76%|███████▌  | 76/100 [00:03<00:01, 15.03it/s]
 79%|███████▉  | 79/100 [00:03<00:01, 16.20it/s]
 81%|████████  | 81/100 [00:03<00:01, 14.83it/s]
 83%|████████▎ | 83/100 [00:03<00:01, 14.03it/s]
 85%|████████▌ | 85/100 [00:03<00:01, 13.36it/s]
 87%|████████▋ | 87/100 [00:03<00:00, 13.23it/s]
 89%|████████▉ | 89/100 [00:04<00:00, 13.25it/s]
 91%|█████████ | 91/100 [00:04<00:00, 13.38it/s]
 93%|█████████▎| 93/100 [00:04<00:00, 12.92it/s]
 95%|█████████▌| 95/100 [00:04<00:00, 12.84it/s]
 97%|█████████▋| 97/100 [00:04<00:00, 13.46it/s]
100%|██████████| 100/100 [00:04<00:00, 15.94it/s]
100%|██████████| 100/100 [00:04<00:00, 20.96it/s]

Setup and run the reconstruction algorithm Data fidelity term

data_fidelity = L2()
# Algorithm parameters
lamb = 1e1
stepsize = 0.8 * float(1 / fourier_op.get_lipschitz_cst().get())
params_algo = {"stepsize": stepsize, "lambda": lamb, "a": 3}
max_iter = 100
early_stop = True

/volatile/github-ci-mind-inria/gpu_mind_runner/_work/mri-nufft/mri-nufft/src/mrinufft/operators/base.py:1075: UserWarning: Lipschitz constant did not converge
  warnings.warn("Lipschitz constant did not converge")

Instantiate the algorithm class to solve the problem.

wavelet_recon = optim_builder(
    iteration="FISTA",
    prior=wavelet,
    data_fidelity=data_fidelity,
    early_stop=early_stop,
    max_iter=max_iter,
    params_algo=params_algo,
)
x_wavelet = wavelet_recon(y, physics)

/volatile/github-ci-mind-inria/gpu_mind_runner/_work/mri-nufft/mri-nufft/.venv/lib/python3.10/site-packages/cufinufft/_plan.py:393: UserWarning: Argument `data` does not satisfy the following requirement: C. Copying array (this may reduce performance)
  warnings.warn(f"Argument `{name}` does not satisfy the "

Display results

plt.figure(figsize=(12, 6))
plt.subplot(1, 3, 1)
plt.imshow(torch.abs(mri[..., mri.shape[2] // 2 - 5]).cpu(), cmap="gray")
plt.title("Ground truth")
plt.axis("off")
plt.subplot(1, 3, 2)
plt.imshow(
    torch.abs(x_dagger[0, 0, ..., x_dagger.shape[2] // 2 - 5]).cpu(), cmap="gray"
)
plt.title("Adjoint reconstruction")
plt.axis("off")
plt.subplot(1, 3, 3)
plt.imshow(
    torch.abs(x_wavelet[0, 0, ..., x_wavelet.shape[2] // 2 - 5]).cpu(), cmap="gray"
)
plt.title("Reconstruction with wavelet prior")
plt.axis("off")
plt.show()