# =============================================================================
# This file is part of LION library
# License : GPL-3
#
# Author : Ander Biguri, Thanh Trung (Troy) Vu
# =============================================================================
"""Plug-and-Play (PnP) Reconstructor using a prior function."""
from __future__ import annotations
import numpy as np
import torch
from tqdm import tqdm
from LION.classical_algorithms.conjugate_gradient import conjugate_gradient
from LION.models.LIONmodel import LIONmodel
from LION.reconstructors.LIONreconstructor import LIONReconstructor
from LION.utils.math import power_method
from LION.utils.parameter import LIONParameter
from LION.utils.citation import citation
# The settings of each algorithm, and their defaults.
# ADMM's coupling weight, relative to ||A||^2: the data step's AᵀA is in the
# operator's units (for CT, ||A||^2 is ~1e5), so an absolute weight works for one
# operator only. 1e-2 was best over doses on a CT problem at LIDC-IDRI's scale.
ADMM_ETA = 1e-2
ALGORITHMS = {
# eta None: ADMM_ETA * ||A||^2, relative to the operator's scale
"ADMM": dict(eta=None, max_iter=10, cg_max_iter=100, cg_tol=1e-7, noise_level=None),
"HQS": dict(
lambda_=0.23,
mu=0.1,
max_iter=100,
noise_level=None,
cg_max_iter=100,
cg_tol=1e-7,
),
"FBS": dict(step_size=None, max_iter=10, noise_level=None),
}
[docs]
class PnP(LIONReconstructor):
"""Plug-and-Play reconstruction with a denoiser (``model``) as the prior.
``PnP(operator, model, parameters)``, with ``parameters`` from
``PnP.default_parameters(algorithm)``:
- ``algorithm``: "ADMM" (Alternating Direction Method of Multipliers),
"HQS" (Half Quadratic Splitting) or "FBS" (Forward-Backward Splitting);
- that algorithm's settings (see ``admm_algorithm``, ``hqs_algorithm``,
``forward_backward_splitting``). A setting missing from the parameters
has its default.
``model``: the denoiser. A LIONmodel prior is called on a batch of one
image (with its normalisation, and the noise level if it uses one). Any
other callable is called on the image as it is (the operator's domain
shape) and must return the same shape.
``noise_level``: in the units of the data the denoiser sees, after its
normalisation, as in training (``GaussianDenoiserSolver``): for images in
[0, 1], the literature's k/255.
The algorithms run without gradients (``torch.no_grad``): the
reconstruction can't be differentiated, and the iterations don't keep
their autograd graph. (A gradient-step denoiser still computes its own
gradient.)
"""
[docs]
def __init__(self, operator, model=None, parameters=None):
super().__init__(operator, model, parameters)
algorithm = getattr(self.parameters, "algorithm", None)
if algorithm not in ALGORITHMS:
raise ValueError(
f"Unknown algorithm: {algorithm!r} (one of {', '.join(ALGORITHMS)})"
)
self._op_norm = None # ||A||, computed once when first needed (FBS, ADMM)
[docs]
@staticmethod
def default_parameters(algorithm="ADMM") -> LIONParameter:
"""``algorithm`` and its settings, with their defaults."""
if algorithm not in ALGORITHMS:
raise ValueError(
f"Unknown algorithm: {algorithm!r} (one of {', '.join(ALGORITHMS)})"
)
return LIONParameter(
algorithm=algorithm,
# where the algorithm starts, f(y, op), when not given x0.
# "default": the operator's default reconstruction (FDK for CT),
# replaced by that function when the reconstructor is made.
initial_reconstruction="default",
**ALGORITHMS[algorithm],
)
[docs]
def setting(self, name):
"""A setting of the algorithm: from the parameters, or its default."""
return getattr(
self.parameters, name, ALGORITHMS[self.parameters.algorithm][name]
)
[docs]
def reconstruct_sample(self, y, x0=None, verbose: int = 1):
"""Reconstruct one measurement (e.g. a (C, angles, detectors)
sinogram) with the algorithm and settings in ``self.parameters``."""
if not isinstance(y, torch.Tensor):
raise TypeError("Sinogram must be a torch.Tensor")
algorithm = self.parameters.algorithm
settings = {k: self.setting(k) for k in ALGORITHMS[algorithm]}
if algorithm == "HQS":
return self.hqs_algorithm(y, x0=x0, verbose=verbose, **settings)
if algorithm == "FBS":
return self.forward_backward_splitting(
y, x0=x0, verbose=verbose, **settings
)
return self.admm_algorithm(y, x0=x0, verbose=verbose, **settings)
# ------------------------------------------------------------------ helpers
def _apply_prior(self, x: torch.Tensor, noise_level=None) -> torch.Tensor:
"""Apply the prior to an image ``x`` with the operator's domain shape.
LIONmodel priors get a batch of one (with the model's normalisation and,
if the model takes one, the noise level). Other callables get ``x`` as it
is and must return the same shape.
"""
if not isinstance(self.model, LIONmodel):
return self.model(x)
x_in = self.model.normalise(x.unsqueeze(0))
if self.model.takes_noise_level():
out = self.model(x_in, noise_level=noise_level)
else:
out = self.model(x_in)
return self.model.unnormalise(out, x.unsqueeze(0)).squeeze(0)
def _operator_norm(self, device):
"""||A||, by the power method, computed once per reconstructor."""
if self._op_norm is None:
self._op_norm = power_method(self.op, device=device)
return self._op_norm
def _solve_data_step(self, y, z, weight, x0, cg_max_iter, cg_tol):
"""argmin_x ||A x - y||^2 / 2 + weight ||x - z||^2 / 2, by conjugate gradient."""
def normal_operator(v: torch.Tensor) -> torch.Tensor:
return self.op.adjoint(self.op(v)) + weight * v
rhs = self.op.adjoint(y) + weight * z
return conjugate_gradient(
normal_operator, rhs, x0, max_iter=cg_max_iter, tol=cg_tol
)
# --------------------------------------------------------------- algorithms
[docs]
@torch.no_grad()
def hqs_algorithm(
self,
y,
*,
lambda_=0.23,
mu=0.1,
max_iter=100,
noise_level=None,
cg_max_iter=100,
cg_tol=1e-7,
x0=None,
verbose: int = 1,
):
"""Half Quadratic Splitting (Zhang, Li, Zuo, Zhang, Van Gool and
Timofte, "Plug-and-Play Image Restoration with Deep Denoiser Prior",
IEEE TPAMI 44(10), 2022): alternates::
x = argmin ||A x - y||² / 2 + mu ||x - z||² / 2 (conjugate gradient)
z = prior(x)
and returns z. ``lambda_``: the regularisation weight; ``mu``: the
weight of the coupling term; ``noise_level``: if given, ``mu = lambda_ /
noise_level²``, else the prior's noise level is ``sqrt(lambda_ / mu)``;
``max_iter``: the iterations; ``cg_max_iter``, ``cg_tol``: the
conjugate gradient's; ``x0``: the start (``initial_guess(y)``, the
operator's default reconstruction, FDK for CT, if None); ``verbose``: 0
nothing, 1 a progress bar."""
if noise_level is not None:
mu = lambda_ / noise_level**2
else:
noise_level = float(np.sqrt(lambda_ / mu))
x = self.initial_guess(y) if x0 is None else x0
z = x.clone()
iterator = range(max_iter)
if verbose >= 1:
iterator = tqdm(iterator, desc="HQS iterations")
for _ in iterator:
x = self._solve_data_step(y, z, mu, x, cg_max_iter, cg_tol)
z = self._apply_prior(x, noise_level)
return z
[docs]
@torch.no_grad()
def forward_backward_splitting(
self,
y,
step_size=None,
max_iter=10,
noise_level=None,
x0=None,
verbose: int = 1,
):
"""Forward-backward splitting: x = prior(x - step_size Aᵀ(A x - y)),
``max_iter`` times. ``step_size``: 1 / ||A||² if None (||A|| computed
once per reconstructor, on the first call); ``noise_level``: the
prior's; ``x0``: the start (``initial_guess(y)``, the operator's
default reconstruction, FDK for CT, if None); ``verbose``: 0 nothing, 1
a progress bar."""
if step_size is None:
step_size = 1.0 / (self._operator_norm(y.device) ** 2)
x = self.initial_guess(y) if x0 is None else x0
iterator = range(max_iter)
if verbose >= 1:
iterator = tqdm(iterator, desc="FBS iterations")
for _ in iterator:
# TODO: would it not make sense to have an adaptive noise_level?
step = x - step_size * self.op.adjoint(self.op(x) - y)
x = self._apply_prior(step, noise_level)
return x
[docs]
@torch.no_grad()
def admm_algorithm(
self,
measurement: torch.Tensor,
eta: float | None = None,
max_iter: int = 10,
cg_max_iter: int = 100,
cg_tol: float = 1e-7,
noise_level=None,
x0=None,
verbose: int = 1,
) -> torch.Tensor:
"""ADMM with the prior as the proximal step::
x = argmin ||A x - y||² / 2 + eta ||x - (v - u)||² / 2 (conjugate gradient)
v = prior(x + u)
u = u + x - v
``max_iter`` times, returning x. ``measurement``: y; ``eta``: the
coupling weight (None: ``ADMM_ETA`` ||A||², 1e-2 of it, ||A||
computed once per reconstructor, on the first call; a larger weight
trusts the prior more); ``cg_max_iter``, ``cg_tol``: the conjugate gradient's;
``noise_level``: the prior's; ``x0``: the start of x and v
(``initial_guess(measurement)``, the operator's default reconstruction,
FDK for CT, if None); ``verbose``: 0 nothing, 1 a progress bar."""
if eta is None:
eta = ADMM_ETA * self._operator_norm(measurement.device) ** 2
x = self.initial_guess(measurement) if x0 is None else x0
v = x.clone()
u = torch.zeros(self.op.domain_shape, device=measurement.device)
iterator = range(max_iter)
if verbose >= 1:
iterator = tqdm(iterator, desc="ADMM iterations")
for _ in iterator:
x = self._solve_data_step(measurement, v - u, eta, x, cg_max_iter, cg_tol)
v = self._apply_prior(x + u, noise_level)
u = u + (x - v)
return x
[docs]
@staticmethod
def cite(cite_format="bib"):
"""Where Plug-and-Play comes from: its citation, BibTeX ("bib") or "MLA". The denoiser has its own."""
return citation(
cite_format,
bib="""
@inproceedings{venkatakrishnan2013plug,
title={Plug-and-Play priors for model based reconstruction},
author={Venkatakrishnan, Singanallur V and Bouman, Charles A and Wohlberg, Brendt},
booktitle={2013 IEEE Global Conference on Signal and Information Processing (GlobalSIP)},
pages={945--948},
year={2013},
doi={10.1109/GlobalSIP.2013.6737048}
}
""",
mla="""
Venkatakrishnan, Singanallur V., Charles A. Bouman, and Brendt Wohlberg. "Plug-and-Play priors for model based reconstruction." 2013 IEEE Global Conference on Signal and Information Processing (GlobalSIP) (2013): 945-948.
""",
)