Reconstructors

A reconstructor turns measurements into images with a trained model: the model inside an algorithm (a denoiser as a prior, a learned regulariser), or the model alone. It is made with everything that decides its result (the operator, the model and the algorithm’s settings), then run on data:

reconstructor = PnP(operator, model, parameters)
images = reconstructor.reconstruct(measurements)

So a reconstructor with its settings is one method: evaluations score it, records name it, and the same settings give the same result. This tutorial puts LION’s side by side, then saves one. Reconstructors are optional: a model can always be called directly (Supervised training).

Previous: Saving, loading and records. Next: Evaluation and comparison.

Needs a GPU and the processed LIDC-IDRI data in LION_DATA_PATH; the PnP part downloads the authors’ trained DRUNet (130 MB).

Note

The models trained in these tutorials are toys: a few epochs on a few slices, to show how LION is used, so their results are poor (often worse than FDK). Trained as their papers train them (many epochs, all the data), they do what the papers report.

Classical reconstructions

The simplest are functions of the measurements and the operator (or its parameters): fdk, sirt, tv_min for CT, and FISTA, conjugate gradient and SPGL1 for any operator (LION.classical_algorithms). Evaluations take such functions as methods too (Evaluation and comparison).

import tempfile
from pathlib import Path

import matplotlib.pyplot as plt
import torch
from torch.utils.data import DataLoader, Subset

from LION.experiments import ct_experiments
from LION.metrics.psnr import PSNR
from LION.operators.ct import fdk, sirt

experiment = ct_experiments.LowDoseCTRecon()
geometry = experiment.operator_params
sinogram, image = experiment.get_dataset("validation")[0]

reconstructions = {
    "FDK": fdk(sinogram, geometry),
    "SIRT": sirt(sinogram, geometry, num_iterations=100),
}

The forward reconstructor

ForwardReconstructor runs a model trained to reconstruct (FBPConvNet, LPD, …) as its solver did: the initial reconstruction for image models, the model’s normalisation, without gradients. Here a FBPConvNet trained for a moment:

from LION.models.post_processing.FBPConvNet import FBPConvNet
from LION.optimizers.SupervisedSolver import SupervisedSolver
from LION.reconstructors.ForwardReconstructor import ForwardReconstructor

training = Subset(experiment.get_dataset("train"), range(40))
training_loader = DataLoader(training, batch_size=4, shuffle=True)

folder = Path(tempfile.mkdtemp())
fbpconvnet = FBPConvNet(geometry)
solver = SupervisedSolver(
    fbpconvnet, torch.optim.Adam(fbpconvnet.parameters()), torch.nn.MSELoss()
)
solver.set_training(training_loader)
solver.set_saving(folder, "FBPConvNet")
solver.train(2)
solver.save_final_results()

fbpconvnet, _, _ = FBPConvNet.load(folder / "FBPConvNet")  # from its file
forward = ForwardReconstructor(geometry, fbpconvnet)
reconstructions["FBPConvNet"] = forward.reconstruct(sinogram)

/mhome/damtp/r/ab2860/AItomotools/LION/data_loaders/LIDC_IDRI.py:487: UserWarning: LIDC-IDRI: left out 2 damaged or unusable files (re-run the pre-processing for these patients): ['LIDC-IDRI-0238 (not resampled)', 'LIDC-IDRI-0585 (not resampled)']
  warnings.warn(

Plug-and-Play

PnP uses a denoiser as the prior of an iterative algorithm (ADMM, HQS or FBS). Here the DRUNet its authors trained on natural images in [0, 1], so the images are normalised to [0, 1] for it (Denoisers and Plug-and-Play trains one):

from LION.models.CNNs.drunet import DRUNet
from LION.reconstructors.PnP import PnP

denoiser = DRUNet.pretrained()
denoiser.set_normalisation("dataset", training)

parameters = PnP.default_parameters("ADMM")
parameters.max_iter = 20
parameters.noise_level = 10 / 255
pnp = PnP(geometry, denoiser, parameters)
reconstructions["PnP-ADMM"] = pnp.reconstruct(sinogram)

Learned regularisers

LearnedRegularizerReconstructor minimises ½||Ax - y||² + α R(x) by gradient descent, for a regulariser R: a trained network (Learned regularisers) or any function of an image, e.g. a smoothed total variation:

def smoothed_tv(x, eps=1e-3):
    dx = x[..., 1:, :-1] - x[..., :-1, :-1]
    dy = x[..., :-1, 1:] - x[..., :-1, :-1]
    return torch.sqrt(dx**2 + dy**2 + eps**2).sum()


from LION.reconstructors.LearnedRegularizerReconstructor import (
    LearnedRegularizerReconstructor,
)

parameters = LearnedRegularizerReconstructor.default_parameters()
parameters.alpha = 10.0
parameters.n_steps = 50
regularised = LearnedRegularizerReconstructor(geometry, smoothed_tv, parameters)
reconstructions["smoothed TV"] = regularised.reconstruct(sinogram)

psnr = PSNR()
fig, axes = plt.subplots(1, len(reconstructions), figsize=(4 * len(reconstructions), 4))
for ax, (name, x) in zip(axes, reconstructions.items()):
    ax.imshow(x[0].cpu(), cmap="gray")
    ax.set_title(f"{name}, {psnr(x, image, batched=False):.1f} dB")
    ax.axis("off")

What they share

Every reconstructor:

• is made with (operator, model, parameters): the operator, or its parameters (a CT geometry); the model; the algorithm’s settings, from its class’s default_parameters();

• reconstructs with reconstruct(y, x0=None, verbose=1): one sample (the operator’s range shape) or a batch, each on its own; x0 a start, for algorithms that take one;

• starts, when it starts from an image, from its parameters’ initial_reconstruction (the operator’s default, FDK for CT):

print(pnp.parameters.initial_reconstruction)
batch = torch.stack([sinogram, sinogram])
print(forward.reconstruct(batch).shape)
start = sirt(sinogram, geometry, num_iterations=20)
x = pnp.reconstruct(sinogram, x0=start, verbose=0)

<function fdk at 0x7f2b3c260220>
torch.Size([2, 1, 512, 512])

In a solver’s validation

A solver validates through a reconstructor when the model is used inside one: e.g. a learned regulariser’s (Learned regularisers), or a denoiser validated by its PnP reconstructions. The solver puts the model being trained into it:

solver.set_validation(validation_loader, 1, PSNR(), reconstructor=pnp)

Saving and loading

A reconstructor saves its settings and its model by the file it was loaded from (the weights stay there); load makes both again. instantiate makes it without its model, to give it another:

from LION.reconstructors.LIONreconstructor import LIONReconstructor

forward.save(folder / "FBPConvNet_reconstructor.json")
again = LIONReconstructor.load(folder / "FBPConvNet_reconstructor.json")
print(type(again).__name__, type(again.model).__name__)

other_model = FBPConvNet(geometry)
again = LIONReconstructor.instantiate(folder / "FBPConvNet_reconstructor.json")
again.model = other_model

/mhome/damtp/r/ab2860/AItomotools/LION/reconstructors/LIONreconstructor.py:211: UserWarning: the model given has other weights than the recorded one: the reconstructions won't be the recorded ones
  reconstructor = cls.instantiate(saved, model, allow=allow)
ForwardReconstructor FBPConvNet

Your own

A reconstructor derives from LIONReconstructor and defines default_parameters() and reconstruct_sample(y, x0, verbose); the tutorial Extending LION writes one.