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")
)

/volatile/github-ci-mind-inria/gpu_mind_runner/_work/mri-nufft/venv/lib/python3.10/site-packages/torch/cuda/__init__.py:827: UserWarning: Can't initialize NVML
  warnings.warn("Can't initialize NVML")

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/venv/lib/python3.10/site-packages/mrinufft/_utils.py:76: 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/venv/lib/python3.10/site-packages/mrinufft/_utils.py:81: UserWarning: Samples will be rescaled to [-0.5, 0.5), assuming they were in [-pi, pi)
  warnings.warn(

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, 34.26it/s]
  9%|▉         | 9/100 [00:00<00:02, 37.96it/s]
 13%|█▎        | 13/100 [00:00<00:02, 38.78it/s]
 18%|█▊        | 18/100 [00:00<00:02, 39.31it/s]
 23%|██▎       | 23/100 [00:00<00:01, 39.69it/s]
 28%|██▊       | 28/100 [00:00<00:01, 39.83it/s]
 33%|███▎      | 33/100 [00:00<00:01, 40.02it/s]
 38%|███▊      | 38/100 [00:00<00:01, 40.24it/s]
 43%|████▎     | 43/100 [00:01<00:01, 40.23it/s]
 48%|████▊     | 48/100 [00:01<00:01, 40.30it/s]
 53%|█████▎    | 53/100 [00:01<00:01, 40.40it/s]
 58%|█████▊    | 58/100 [00:01<00:01, 40.39it/s]
 63%|██████▎   | 63/100 [00:01<00:00, 40.27it/s]
 68%|██████▊   | 68/100 [00:01<00:00, 40.30it/s]
 73%|███████▎  | 73/100 [00:01<00:00, 40.26it/s]
 78%|███████▊  | 78/100 [00:02<00:00, 36.04it/s]
 82%|████████▏ | 82/100 [00:02<00:00, 34.57it/s]
 86%|████████▌ | 86/100 [00:02<00:00, 33.56it/s]
 90%|█████████ | 90/100 [00:02<00:00, 29.99it/s]
 94%|█████████▍| 94/100 [00:02<00:00, 26.13it/s]
 97%|█████████▋| 97/100 [00:02<00:00, 23.64it/s]
100%|██████████| 100/100 [00:02<00:00, 22.72it/s]
100%|██████████| 100/100 [00:02<00:00, 33.93it/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/venv/lib/python3.10/site-packages/mrinufft/operators/base.py:1032: 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/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()