#### NOTE
New to DeepInverse? Get started with the basics with the
[5 minute quickstart tutorial](https://deepinv.org/auto_examples/basics/demo_quickstart.html.md#sphx-glr-auto-examples-basics-demo-quickstart-py)..
# Tour of MRI functionality in DeepInverse
This example presents the various datasets, forward physics and models
available in DeepInverse for Magnetic Resonance Imaging (MRI) problems:
- Physics: [`deepinv.physics.MRI`](https://deepinv.org/api/stubs/deepinv.physics.MRI.html.md#deepinv.physics.MRI),
[`deepinv.physics.MultiCoilMRI`](https://deepinv.org/api/stubs/deepinv.physics.MultiCoilMRI.html.md#deepinv.physics.MultiCoilMRI),
[`deepinv.physics.DynamicMRI`](https://deepinv.org/api/stubs/deepinv.physics.DynamicMRI.html.md#deepinv.physics.DynamicMRI)
- Datasets: raw kspace with the [FastMRI](https://fastmri.med.nyu.edu) dataset
[`deepinv.datasets.FastMRISliceDataset`](https://deepinv.org/api/stubs/deepinv.datasets.FastMRISliceDataset.html.md#deepinv.datasets.FastMRISliceDataset) and an in-memory easy-to-use version
[`deepinv.datasets.SimpleFastMRISliceDataset`](https://deepinv.org/api/stubs/deepinv.datasets.SimpleFastMRISliceDataset.html.md#deepinv.datasets.SimpleFastMRISliceDataset), and raw dynamic k-t-space data with the
[CMRxRecon](https://cmrxrecon.github.io) dataset.
- Models: [`deepinv.models.VarNet`](https://deepinv.org/api/stubs/deepinv.models.VarNet.html.md#deepinv.models.VarNet)
(VarNet Hammernik *et al.*[1](#footcite-hammernik2018learning), E2E-VarNet Sriram *et al.*[2](#footcite-sriram2020end)),
[`deepinv.models.MoDL`](https://deepinv.org/api/stubs/deepinv.models.MoDL.html.md#deepinv.models.MoDL) (a simple MoDL Aggarwal *et al.*[3](#footcite-aggarwal2018modl) unrolled model)
Contents:
1. Get started with FastMRI (singlecoil + multicoil)
2. Train an accelerated MRI with neural networks
3. Load raw FastMRI data (singlecoil + multicoil)
4. Train using raw data
5. Explore 3D MRI
6. Explore dynamic MRI
```Python
import deepinv as dinv
import torch, torchvision
from torch.utils.data import DataLoader
device = dinv.utils.get_device()
rng = torch.Generator(device=device).manual_seed(0)
```
```none
Selected GPU 0 with 27186.0 MiB free memory
```
## 1. Get started with FastMRI
You can get started with our simple
[FastMRI](https://fastmri.med.nyu.edu) mini slice subsets which provide
quick, easy-to-use, in-memory datasets which can be used for simulation
experiments.
#### IMPORTANT
By using this dataset, you confirm that you have agreed to and signed the [FastMRI data use agreement](https://fastmri.med.nyu.edu/).
#### SEE ALSO
Datasets [`deepinv.datasets.FastMRISliceDataset`](https://deepinv.org/api/stubs/deepinv.datasets.FastMRISliceDataset.html.md#deepinv.datasets.FastMRISliceDataset) [`deepinv.datasets.SimpleFastMRISliceDataset`](https://deepinv.org/api/stubs/deepinv.datasets.SimpleFastMRISliceDataset.html.md#deepinv.datasets.SimpleFastMRISliceDataset)
: We provide convenient datasets to easily load both raw and reconstructed FastMRI images.
You can download more data on the [FastMRI site](https://fastmri.med.nyu.edu/).
Load mini demo knee and brain datasets (original data is 320x320 but we resize to
128 for speed):
```Python
transform = torchvision.transforms.Resize(128)
knee_dataset = dinv.datasets.SimpleFastMRISliceDataset(
anatomy="knee",
transform=transform,
train=True,
download=True,
)
brain_dataset = dinv.datasets.SimpleFastMRISliceDataset(
anatomy="brain",
transform=transform,
train=True,
download=True,
)
img_size = knee_dataset[0].shape[-2:] # (128, 128)
dinv.utils.plot({"knee": knee_dataset[0], "brain": brain_dataset[0]})
```
```none
/local/jtachell/deepinv/deepinv/deepinv/datasets/fastmri.py:109: DeprecationWarning: The tuple format for dataset outputs is deprecated and will be removed in a future version.It is recommended to set `use_dict_output=True` for better readability and flexibility in returned outputs.The default is currently `False` for backward compatibility, but will be switched to `True` in a future version.
super().__init__(use_dict_output=use_dict_output)
```
Let’s start with single-coil MRI. We can define a constant Cartesian 4x
undersampling mask by sampling once from a physics generator. The mask,
data and measurements will all be of shape `(B, C, H, W)` where
`C=2` is the real and imaginary parts.
```Python
physics_generator = dinv.physics.generator.GaussianMaskGenerator(
img_size=img_size, acceleration=4, rng=rng, device=device
)
mask = physics_generator.step()["mask"]
physics = dinv.physics.MRI(mask=mask, img_size=img_size, device=device)
x = next(iter(DataLoader(knee_dataset))).to(device)
dinv.utils.plot(
{
"x": x,
"mask": mask,
"y": physics(x).clamp(-1, 1),
}
)
print("Shapes:", x.shape, physics.mask.shape)
```
```none
Shapes: torch.Size([1, 2, 128, 128]) torch.Size([1, 2, 128, 128])
```
We can next generate an accelerated single-coil MRI measurement dataset. Let’s use knees
for training and brains for testing.
We can also use the physics generator to randomly sample a new mask per
sample, and save the masks alongside the measurements.
Note that you could alternatively train using `online_measurements`, where you can generate
random measurements on the fly.
```Python
dataset_path = dinv.datasets.generate_dataset(
train_dataset=knee_dataset,
test_dataset=brain_dataset,
val_dataset=None,
physics=physics,
physics_generator=physics_generator,
save_physics_generator_params=True,
overwrite_existing=False,
device=device,
save_dir=dinv.utils.get_cache_home(),
batch_size=1,
)
train_dataset = dinv.datasets.HDF5Dataset(
dataset_path,
split="train",
load_physics_generator_params=True,
use_dict_output=True,
)
test_dataset = dinv.datasets.HDF5Dataset(
dataset_path, split="test", load_physics_generator_params=True, use_dict_output=True
)
train_dataloader = DataLoader(train_dataset)
iterator = iter(train_dataloader)
batch0 = next(iterator)
batch1 = next(iterator)
x0, y0, params0 = batch0["x"], batch0["y"], batch0["params"]
x1, y1, params1 = batch1["x"], batch1["y"], batch1["params"]
dinv.utils.plot(
{
"x0": x0,
"mask0": params0["mask"],
"x1": x1,
"mask1": params1["mask"],
}
)
```
```none
/local/jtachell/deepinv/deepinv/deepinv/datasets/datagenerator.py:630: UserWarning: Dataset /local/jtachell/.cache/deepinv/dinv_dataset0.h5 already exists, skipping...
warn(f"Dataset {hf_path} already exists, skipping...")
```
We can also simulate multicoil MRI data. Either pass in ground-truth
coil maps, or pass an integer to simulate simple birdcage coil maps. The
measurements `y` are now of shape `(B, C, N, H, W)`, where `N` is
the coil-dimension.
```Python
mc_physics = dinv.physics.MultiCoilMRI(img_size=img_size, coil_maps=3, device=device)
dinv.utils.plot(
{
"x": x,
"mask": mask,
"coil_map_0": mc_physics.coil_maps.abs()[:, 0, ...],
"coil_map_1": mc_physics.coil_maps.abs()[:, 1, ...],
"coil_map_2": mc_physics.coil_maps.abs()[:, 2, ...],
"RSS": mc_physics.A_adjoint_A(x, mask=mask, rss=True),
}
)
```
## 2. Train an accelerated MRI problem with neural networks
Next, we train a neural network to solve the MRI inverse problem. We provide various
models specifically used for MRI reconstruction. These are unrolled
networks which require a backbone denoiser, such as UNet or DnCNN:
```Python
denoiser = dinv.models.UNet(
in_channels=2,
out_channels=2,
scales=2,
)
denoiser = dinv.models.DnCNN(
in_channels=2,
out_channels=2,
pretrained=None,
depth=2,
)
```
These backbones can be used within specific MRI models, such as
VarNet Hammernik *et al.*[1](#footcite-hammernik2018learning), E2E-VarNet Sriram *et al.*[2](#footcite-sriram2020end) and MoDL Aggarwal *et al.*[3](#footcite-aggarwal2018modl),
for which we provide implementations:
```Python
model = dinv.models.VarNet(denoiser, num_cascades=2, mode="varnet").to(device)
model = dinv.models.MoDL(denoiser, num_iter=2).to(device)
```
Now that we have our architecture defined, we can train it with supervised or self-supervised (using Equivariant
Imaging) loss. We use the PSNR metric on the complex magnitude.
For the sake of speed in this example, we only use a very small 2-layer DnCNN inside an unrolled
network with 2 cascades, and train with 2 images for 1 epoch.
```Python
loss = dinv.loss.SupLoss()
loss = dinv.loss.EILoss(transform=dinv.transform.CPABDiffeomorphism())
trainer = dinv.Trainer(
model=model,
physics=physics,
optimizer=torch.optim.Adam(model.parameters()),
train_dataloader=train_dataloader,
metrics=dinv.metric.PSNR(complex_abs=True),
epochs=1,
show_progress_bar=False,
save_path=None,
)
```
To improve results in the case of this very short training, we start training from a pretrained model state (trained on 900 images):
```Python
url = dinv.models.utils.get_weights_url(
model_name="demo", file_name="demo_tour_mri.pth"
)
ckpt = torch.hub.load_state_dict_from_url(
url, map_location=lambda storage, loc: storage, file_name="demo_tour_mri.pth"
)
trainer.model.load_state_dict(ckpt["state_dict"]) # load the state dict
trainer.optimizer.load_state_dict(ckpt["optimizer"]) # load the optimizer state dict
model = trainer.train() # train the model
trainer.plot_images = True
```
```none
/local/jtachell/deepinv/deepinv/deepinv/training/trainer.py:1415: UserWarning: non_blocking_transfers=True but DataLoader.pin_memory=False; set pin_memory=True to overlap host-device copies with compute.
self.setup_train()
The model has 2376 trainable parameters
Train epoch 0: TotalLoss=0.0, PSNR=30.979
```
Now that our model is trained, we can test it. Notice that we improve the PSNR compared to the zero-filled
reconstruction, both on the train (knee) set and the test (brain) set:
```Python
_ = trainer.test(train_dataloader)
_ = trainer.test(DataLoader(test_dataset))
```
```none
/local/jtachell/deepinv/deepinv/deepinv/training/trainer.py:1611: UserWarning: non_blocking_transfers=True but DataLoader.pin_memory=False; set pin_memory=True to overlap host-device copies with compute.
self.setup_train(train=False)
Eval epoch 0: PSNR=31.005, PSNR no learning=29.991
Test results:
PSNR no learning: 29.991 +- 0.464
PSNR: 31.005 +- 0.718
/local/jtachell/deepinv/deepinv/deepinv/training/trainer.py:1611: UserWarning: non_blocking_transfers=True but DataLoader.pin_memory=False; set pin_memory=True to overlap host-device copies with compute.
self.setup_train(train=False)
Eval epoch 0: PSNR=28.879, PSNR no learning=28.03
Test results:
PSNR no learning: 28.030 +- 0.351
PSNR: 28.879 +- 0.327
```
## 3. Load raw FastMRI data
It is also possible to use the raw data directly.
The raw multi-coil FastMRI train/validation data is provided as pairs of `(x, y)` where
`y` are the fully-sampled k-space measurements of arbitrary size, and
`x` are the cropped root-sum-square (RSS) magnitude reconstructions.
Let’s download a sample volume and check out its middle slice.
```Python
dinv.datasets.download_archive(
dinv.utils.get_image_url("demo_fastmri_brain_multicoil.h5"),
dinv.utils.get_cache_home() / "brain" / "fastmri.h5",
)
dataset = dinv.datasets.FastMRISliceDataset(
dinv.utils.get_cache_home() / "brain", slice_index="middle", use_dict_output=True
)
batch = next(iter(DataLoader(dataset)))
x, y = batch["x"].to(device), batch["y"].to(device)
img_size, kspace_shape = x.shape[-2:], y.shape[-2:]
n_coils = y.shape[2]
print("Shapes:", x.shape, y.shape) # x (B, 1, W, W); y (B, C, N, H, W)
```
```none
File already downloaded: /local/jtachell/.cache/deepinv/brain/fastmri.h5. Skipping...
0%| | 0/1 [00:00
Note that we can relate `x` and fully-sampled `y` using our
[`deepinv.physics.MultiCoilMRI`](https://deepinv.org/api/stubs/deepinv.physics.MultiCoilMRI.html.md#deepinv.physics.MultiCoilMRI) (note that since we are not
provided with the ground-truth coil-maps, we can only perform the
adjoint operator).
```Python
physics = dinv.physics.MultiCoilMRI(
img_size=img_size,
mask=torch.ones(kspace_shape, device=device),
coil_maps=torch.ones(
(n_coils,) + kspace_shape, dtype=torch.complex64, device=device
),
device=device,
)
x_rss = physics.A_adjoint(y, rss=True)
x_rss = physics.crop(x_rss, shape=x.shape) # FastMRI provided RSS is cropped
assert torch.allclose(x, x_rss)
```
We can also pre-estimate coil sensitivity maps using ESPIRiT from the raw data.
```Python
dataset = dinv.datasets.FastMRISliceDataset(
dinv.utils.get_cache_home() / "brain",
slice_index="middle",
transform=dinv.datasets.MRISliceTransform(
estimate_coil_maps=True,
acs=15, # Num. low frequency, fix to 15
),
use_dict_output=True,
)
batch = next(iter(DataLoader(dataset)))
x, y, params = batch["x"].to(device), batch["y"].to(device), batch["params"]
physics.update(**params)
dinv.utils.plot(
{"x": x, "maps0": physics.coil_maps[:, 0], "maps1": physics.coil_maps[:, 1]}
)
```
```none
0%| | 0/1 [00:00
## 4. Train using raw data
For training with multicoil raw data, we can simulate random masks **on-the-fly**:
```Python
dataset = dinv.datasets.FastMRISliceDataset(
dinv.utils.get_cache_home() / "brain",
slice_index="middle",
transform=dinv.datasets.MRISliceTransform(
mask_generator=dinv.physics.generator.GaussianMaskGenerator(
img_size=kspace_shape,
acceleration=4,
),
seed_mask_generator=False, # More diversity during training
estimate_coil_maps=False, # Set to true if coil maps are not already set in physics.
# This will use ACS size from mask generator. If mask generator is None, then try find ACS size from metadata.
),
use_dict_output=True,
)
```
```none
0%| | 0/1 [00:00
Note if the data is already undersampled raw kspace data (e.g. FastMRI test set)
you can also easily directly load it and their associated masks for testing or training
(optionally specify separate target folder if targets are in a different folder):
```default
dataset = dinv.datasets.FastMRISliceDataset(
root=root,
target_root=target_root,
transform=dinv.datasets.MRISliceTransform()
)
```
We use the E2E-VarNet model designed for
multicoil MRI. For this example, we do not perform joint coil sensitivity map estimation and
simply assume they are flat. If you want to estimate the maps, either pass a model
as the `sensitivity_model` parameter, or use a different model which uses precomputed maps.
```Python
model = dinv.models.VarNet(denoiser, num_cascades=2, mode="e2e-varnet").to(device)
```
We also need to modify the metrics used to crop the model output and take the magnitude when
comparing to the cropped magnitude RSS targets:
```Python
def crop(x_net, x):
"""Crop to GT shape then take magnitude."""
return dinv.utils.MRIMixin().rss(
dinv.utils.MRIMixin().crop(x_net, shape=x.shape), multicoil=False
)
class CropPSNR(dinv.metric.PSNR):
def forward(self, x_net=None, x=None, *args, **kwargs):
return super().forward(crop(x_net, x), x, *args, **kwargs)
class CropMSE(dinv.metric.MSE):
def forward(self, x_net=None, x=None, *args, **kwargs):
return super().forward(crop(x_net, x), x, *args, **kwargs)
trainer = dinv.Trainer(
model=model,
physics=physics,
losses=dinv.loss.SupLoss(metric=CropMSE()),
metrics=CropPSNR(),
optimizer=torch.optim.Adam(model.parameters()),
train_dataloader=DataLoader(dataset),
epochs=1,
save_path=None,
show_progress_bar=False,
device=device,
)
_ = trainer.train()
```
```none
/local/jtachell/deepinv/deepinv/deepinv/training/trainer.py:1415: UserWarning: non_blocking_transfers=True but DataLoader.pin_memory=False; set pin_memory=True to overlap host-device copies with compute.
self.setup_train()
The model has 2372 trainable parameters
Train epoch 0: TotalLoss=0.004, CropPSNR=24.31
```
## 5. Explore 3D MRI
We can also simulate 3D MRI data.
Here, we use a demo 3D brain volume of shape `(181, 217, 181)` from the
[BrainWeb](https://brainweb.bic.mni.mcgill.ca/brainweb/) dataset
and simulate 3D single-coil or multi-coil Fourier measurements using
[`deepinv.physics.MRI`](https://deepinv.org/api/stubs/deepinv.physics.MRI.html.md#deepinv.physics.MRI) or
[`deepinv.physics.MultiCoilMRI`](https://deepinv.org/api/stubs/deepinv.physics.MultiCoilMRI.html.md#deepinv.physics.MultiCoilMRI).
```Python
x = (
dinv.utils.load_np_url(
"https://huggingface.co/datasets/deepinv/images/resolve/main/brainweb_t1_ICBM_1mm_subject_0.npy?download=true"
)
.unsqueeze(0)
.unsqueeze(0)
.to(device)
)
x = torch.cat([x, torch.zeros_like(x)], dim=1) # add imaginary dimension
print(x.shape) # (B, C, D, H, W) where D is depth
physics = dinv.physics.MultiCoilMRI(img_size=x.shape[1:], three_d=True, device=device)
physics = dinv.physics.MRI(img_size=x.shape[1:], three_d=True, device=device)
dinv.utils.plot_ortho3D([x, physics(x)], titles=["x", "y"])
```
```none
torch.Size([1, 2, 181, 217, 181])
/local/jtachell/deepinv/deepinv/deepinv/utils/plotting.py:1373: UserWarning: This figure was using a layout engine that is incompatible with subplots_adjust and/or tight_layout; not calling subplots_adjust.
plt.subplots_adjust(hspace=0.05, wspace=0.05)
```
## 6. Explore dynamic MRI
Finally, we show how to use the dynamic MRI for image sequence data of
shape `(B, C, T, H, W)` where `T` is the time dimension. Note that
this is also compatible with 3D MRI. We use dynamic MRI data from the
[CMRxRecon](https://cmrxrecon.github.io/) challenge of cardiac cine
sequences and load them using [`deepinv.datasets.CMRxReconSliceDataset`](https://deepinv.org/api/stubs/deepinv.datasets.CMRxReconSliceDataset.html.md#deepinv.datasets.CMRxReconSliceDataset)
provided in deepinv. We download demo data from the first patient
including ground truth images, undersampled kspace, and associated masks:
```Python
dinv.datasets.download_archive(
dinv.utils.get_image_url("CMRxRecon.zip"),
dinv.utils.get_cache_home() / "CMRxRecon.zip",
extract=True,
)
dataset = dinv.datasets.CMRxReconSliceDataset(
dinv.utils.get_cache_home() / "CMRxRecon", use_dict_output=True
)
batch = next(iter(DataLoader(dataset)))
x, y, params = batch["x"].to(device), batch["y"].to(device), batch["params"]
print(f"""
Ground truth: {x.shape} (B, C, T, H, W)
Measurements: {y.shape}
Acc. mask: {params["mask"].shape}
""")
```
```none
File already downloaded: /local/jtachell/.cache/deepinv/CMRxRecon.zip. Skipping...
0%| | 0/1 [00:00
Dynamic MRI data is directly compatible with existing functionality.
For example, you can train with this data by passing the dataset to
[`deepinv.Trainer`](https://deepinv.org/api/stubs/deepinv.Trainer.html.md#deepinv.Trainer), which will automatically load in the data
`x, y, params`. Or, you can use the data directly with the physics
[`deepinv.physics.DynamicMRI`](https://deepinv.org/api/stubs/deepinv.physics.DynamicMRI.html.md#deepinv.physics.DynamicMRI).
You can also pass in a custom k-t acceleration mask generator to
generate random time-varying masks:
```Python
physics_generator = dinv.physics.generator.EquispacedMaskGenerator(
img_size=x.shape[1:],
acceleration=16,
)
physics = dinv.physics.DynamicMRI(img_size=(512, 256), device=device)
dataset = dinv.datasets.CMRxReconSliceDataset(
dinv.utils.get_cache_home() / "CMRxRecon",
mask_generator=physics_generator,
mask_dir=None,
use_dict_output=True,
)
batch = next(iter(DataLoader(dataset)))
x, y, params = batch["x"].to(device), batch["y"].to(device), batch["params"]
x = x.to(device)
y = y.to(device)
params = {k: v.to(device) for k, v in params.items()}
```
```none
0%| | 0/1 [00:00
We provide a video plotting function, [`deepinv.utils.plot_videos`](https://deepinv.org/api/stubs/deepinv.utils.plot_videos.html.md#deepinv.utils.plot_videos). Here, we
visualize t=5 frames of the ground truth `x`, the mask, and the zero-filled
reconstruction `x_zf` (and crop to square for better visibility):
```Python
x_zf = physics.A_adjoint(y, **params)
dinv.utils.plot(
{
f"t={i}": torch.cat([x[:, :, i], params["mask"][:, :, i], x_zf[:, :, i]])[
..., 128:384, :
]
for i in range(5)
}
)
```
* **References:**
* Kerstin Hammernik, Teresa Klatzer, Erich Kobler, Michael P Recht, Daniel K Sodickson, Thomas Pock, and Florian Knoll. Learning a variational network for reconstruction of accelerated mri data. *Magnetic resonance in medicine*, 79(6):3055–3071, 2018.
* Anuroop Sriram, Jure Zbontar, Tullie Murrell, Aaron Defazio, C Lawrence Zitnick, Nafissa Yakubova, Florian Knoll, and Patricia Johnson. End-to-end variational networks for accelerated mri reconstruction. In *Medical image computing and computer assisted intervention–MICCAI 2020: 23rd international conference, Lima, Peru, October 4–8, 2020, proceedings, part II 23*, 64–73. Springer, 2020.
* Hemant K Aggarwal, Merry P Mani, and Mathews Jacob. Modl: model-based deep learning architecture for inverse problems. *IEEE transactions on medical imaging*, 38(2):394–405, 2018.
**Total running time of the script:** (0 minutes 3.776 seconds)
🏷 Tags: MRI
[`Download Jupyter notebook: demo_mri_tour.ipynb`](https://deepinv.org/auto_examples/physics/_downloads/eb647c1bf8ad5ea4472f9e1873a36999/demo_mri_tour.ipynb)
[`Download Python source code: demo_mri_tour.py`](https://deepinv.org/auto_examples/physics/_downloads/5d281bf14be29ad8b15d60df168778a9/demo_mri_tour.py)
[`Download zipped: demo_mri_tour.zip`](https://deepinv.org/auto_examples/physics/_downloads/5a66fae685eee96e53e1e57f2ed44add/demo_mri_tour.zip)
[Gallery generated by Sphinx-Gallery](https://sphinx-gallery.github.io)