+"""Total generalized variation denoising prior."""
+
+__all__ = ["TGVDenoiser", "tgv_denoise"]
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+
+[docs]class TGVDenoiser(nn.Module):
+
r"""
+
Proximal operator of (2nd order) Total Generalised Variation operator.
+
+
(see K. Bredies, K. Kunisch, and T. Pock, "Total generalized variation," SIAM J. Imaging Sci., 3(3), 492-526, 2010.)
+
+
This algorithm converges to the unique image :math:`x` (and the auxiliary vector field :math:`r`) minimizing
+
+
.. math::
+
+
\underset{x, r}{\arg\min} \; \frac{1}{2}\|x-y\|_2^2 + \lambda_1 \|r\|_{1,2} + \lambda_2 \|J(Dx-r)\|_{1,F}
+
+
where :math:`D` maps an image to its gradient field and :math:`J` maps a vector field to its Jacobian.
+
For a large value of :math:`\lambda_2`, the TGV behaves like the TV.
+
For a small value, it behaves like the :math:`\ell_1`-Frobenius norm of the Hessian.
+
+
The problem is solved with an over-relaxed Chambolle-Pock algorithm (see L. Condat, "A primal-dual splitting method
+
for convex optimization involving Lipschitzian, proximable and linear composite terms", J. Optimization Theory and
+
Applications, vol. 158, no. 2, pp. 460-479, 2013.
+
+
Code (and description) adapted from Laurent Condat's matlab version (https://lcondat.github.io/software.html) and
+
Daniil Smolyakov's `code <https://github.com/RoundedGlint585/TGVDenoising/blob/master/TGV%20WithoutHist.ipynb>`_.
+
+
Attributes
+
----------
+
ndim : int
+
Number of spatial dimensions, can be either ``2`` or ``3``.
+
ths : float, optional
+
Denoise threshold. The default is ``0.1``.
+
trainable : bool, optional
+
If ``True``, threshold value is trainable, otherwise it is not.
+
The default is ``False``.
+
device : str, optional
+
Device on which the wavelet transform is computed.
+
The default is ``None`` (infer from input).
+
verbose : bool, optional
+
Whether to print computation details or not. The default is ``False``.
+
niter : int, optional,
+
Maximum number of iterations. The default is ``1000``.
+
crit : float, optional
+
Convergence criterion. The default is 1e-5.
+
x2 : torch.Tensor, optional
+
Primary variable for warm restart. The default is ``None``.
+
u2 : torch.Tensor, optional
+
Dual variable for warm restart. The default is ``None``.
+
r2 : torch.Tensor, optional
+
Auxiliary variable for warm restart. The default is ``None``.
+
+
Notes
+
-----
+
The regularization term :math:`\|r\|_{1,2} + \|J(Dx-r)\|_{1,F}` is implicitly normalized by its Lipschitz
+
constant, i.e. :math:`\sqrt{72}`, see e.g. K. Bredies et al., "Total generalized variation," SIAM J. Imaging
+
Sci., 3(3), 492-526, 2010.
+
+
"""
+
+
[docs] def __init__(
+
self,
+
ndim,
+
ths=0.1,
+
trainable=False,
+
device=None,
+
verbose=False,
+
niter=100,
+
crit=1e-5,
+
x2=None,
+
u2=None,
+
r2=None,
+
):
+
super().__init__()
+
+
if trainable:
+
self.ths = nn.Parameter(ths)
+
else:
+
self.ths = ths
+
+
self.denoiser = _TGVDenoiser(
+
ndim=ndim,
+
device=device,
+
verbose=verbose,
+
n_it_max=niter,
+
crit=crit,
+
x2=x2,
+
u2=u2,
+
r2=r2,
+
)
+
self.denoiser.device = device
+
+
def forward(self, input):
+
# get complex
+
if torch.is_complex(input):
+
iscomplex = True
+
else:
+
iscomplex = False
+
+
# default device
+
idevice = input.device
+
if self.denoiser.device is None:
+
device = idevice
+
else:
+
device = self.denoiser.device
+
+
# get input shape
+
ndim = self.denoiser.ndim
+
ishape = input.shape
+
+
# reshape for computation
+
input = input.reshape(-1, *ishape[-ndim:])
+
if iscomplex:
+
input = torch.stack((input.real, input.imag), axis=1)
+
input = input.reshape(-1, *ishape[-ndim:])
+
+
# apply denoising
+
output = self.denoiser(input[:, None, ...].to(device), self.ths).to(
+
idevice
+
) # perform the denoising on the real-valued tensor
+
+
# reshape back
+
if iscomplex:
+
output = (
+
output[::2, ...] + 1j * output[1::2, ...]
+
) # build the denoised complex data
+
output = output.reshape(ishape)
+
+
return output.to(idevice)
+
+
+[docs]def tgv_denoise(
+
input,
+
ndim,
+
ths=0.1,
+
device=None,
+
verbose=False,
+
niter=100,
+
crit=1e-5,
+
x2=None,
+
u2=None,
+
):
+
r"""
+
Apply Total Generalized Variation denoising.
+
+
(see K. Bredies, K. Kunisch, and T. Pock, "Total generalized variation," SIAM J. Imaging Sci., 3(3), 492-526, 2010.)
+
+
This algorithm converges to the unique image :math:`x` (and the auxiliary vector field :math:`r`) minimizing
+
+
.. math::
+
+
\underset{x, r}{\arg\min} \; \frac{1}{2}\|x-y\|_2^2 + \lambda_1 \|r\|_{1,2} + \lambda_2 \|J(Dx-r)\|_{1,F}
+
+
where :math:`D` maps an image to its gradient field and :math:`J` maps a vector field to its Jacobian.
+
For a large value of :math:`\lambda_2`, the TGV behaves like the TV.
+
For a small value, it behaves like the :math:`\ell_1`-Frobenius norm of the Hessian.
+
+
The problem is solved with an over-relaxed Chambolle-Pock algorithm (see L. Condat, "A primal-dual splitting method
+
for convex optimization involving Lipschitzian, proximable and linear composite terms", J. Optimization Theory and
+
Applications, vol. 158, no. 2, pp. 460-479, 2013.
+
+
Code (and description) adapted from Laurent Condat's matlab version (https://lcondat.github.io/software.html) and
+
Daniil Smolyakov's `code <https://github.com/RoundedGlint585/TGVDenoising/blob/master/TGV%20WithoutHist.ipynb>`_.
+
+
Arguments
+
---------
+
input : np.ndarray | torch.Tensor
+
Input image of shape (..., n_ndim, ..., n_0).
+
ndim : int
+
Number of spatial dimensions, can be either ``2`` or ``3``.
+
ths : float, optional
+
Denoise threshold. Default is ``0.1``.
+
ndim : int
+
Number of spatial dimensions, can be either ``2`` or ``3``.
+
ths : float, optional
+
Denoise threshold. The default is ``0.1``.
+
trainable : bool, optional
+
If ``True``, threshold value is trainable, otherwise it is not.
+
The default is ``False``.
+
device : str, optional
+
Device on which the wavelet transform is computed.
+
The default is ``None`` (infer from input).
+
verbose : bool, optional
+
Whether to print computation details or not. The default is ``False``.
+
niter : int, optional,
+
Maximum number of iterations. The default is ``1000``.
+
crit : float, optional
+
Convergence criterion. The default is 1e-5.
+
x2 : torch.Tensor, optional
+
Primary variable for warm restart. The default is ``None``.
+
u2 : torch.Tensor, optional
+
Dual variable for warm restart. The default is ``None``.
+
r2 : torch.Tensor, optional
+
Auxiliary variable for warm restart. The default is ``None``.
+
+
Returns
+
-------
+
output : np.ndarray | torch.Tensor
+
Denoised image of shape (..., n_ndim, ..., n_0).
+
+
"""
+
# cast to numpy if required
+
if isinstance(input, np.ndarray):
+
isnumpy = True
+
input = torch.as_tensor(input)
+
else:
+
isnumpy = False
+
+
# initialize denoiser
+
TV = TGVDenoiser(ndim, ths, False, device, verbose, niter, crit, x2, u2)
+
output = TV(input)
+
+
# cast back to numpy if requried
+
if isnumpy:
+
output = output.numpy(force=True)
+
+
return output
+
+
+# %% local utils
+class _TGVDenoiser(nn.Module):
+ def __init__(
+ self,
+ ndim,
+ device,
+ verbose=False,
+ n_it_max=1000,
+ crit=1e-5,
+ x2=None,
+ u2=None,
+ r2=None,
+ ):
+ super().__init__()
+ self.device = device
+ self.ndim = ndim
+
+ if ndim == 2:
+ self.nabla = self.nabla2
+ self.nabla_adjoint = self.nabla2_adjoint
+ self.epsilon = self.epsilon2
+ self.epsilon_adjoint = self.epsilon2_adjoint
+ elif ndim == 3:
+ self.nabla = self.nabla3
+ self.nabla_adjoint = self.nabla3_adjoint
+ self.epsilon = self.epsilon3
+ self.epsilon_adjoint = self.epsilon3_adjoint
+
+ self.verbose = verbose
+ self.n_it_max = n_it_max
+ self.crit = crit
+ self.restart = True
+
+ self.tau = 0.01 # > 0
+
+ self.rho = 1.99 # in 1,2
+ self.sigma = 1 / self.tau / 72
+
+ self.x2 = x2
+ self.r2 = r2
+ self.u2 = u2
+
+ self.has_converged = False
+
+ def prox_tau_fx(self, x, y):
+ return (x + self.tau * y) / (1 + self.tau)
+
+ def prox_tau_fr(self, r, lambda1):
+ left = torch.sqrt(torch.sum(r**2, axis=-1)) / (self.tau * lambda1)
+ tmp = r - r / (
+ torch.maximum(
+ left, torch.tensor([1], device=left.device).type(left.dtype)
+ ).unsqueeze(-1)
+ )
+ return tmp
+
+ def prox_sigma_g_conj(self, u, lambda2):
+ return u / (
+ torch.maximum(
+ torch.sqrt(torch.sum(u**2, axis=-1)) / lambda2,
+ torch.tensor([1], device=u.device).type(u.dtype),
+ ).unsqueeze(-1)
+ )
+
+ def forward(self, y, ths=None):
+ restart = (
+ True
+ if (self.restart or self.x2 is None or self.x2.shape != y.shape)
+ else False
+ )
+
+ if restart:
+ self.x2 = y.clone()
+ self.r2 = torch.zeros((*self.x2.shape, 2), device=self.x2.device).type(
+ self.x2.dtype
+ )
+ self.u2 = torch.zeros((*self.x2.shape, 4), device=self.x2.device).type(
+ self.x2.dtype
+ )
+ self.restart = False
+
+ if ths is not None:
+ lambda1 = ths * 0.1
+ lambda2 = ths * 0.15
+
+ cy = (y**2).sum() / 2
+ primalcostlowerbound = 0
+
+ for _ in range(self.n_it_max):
+ x_prev = self.x2.clone()
+ tmp = self.tau * self.epsilon_adjoint(self.u2)
+ x = self.prox_tau_fx(self.x2 - self.nabla_adjoint(tmp), y)
+ r = self.prox_tau_fr(self.r2 + tmp, lambda1)
+ u = self.prox_sigma_g_conj(
+ self.u2
+ + self.sigma
+ * self.epsilon(self.nabla(2 * x - self.x2) - (2 * r - self.r2)),
+ lambda2,
+ )
+ self.x2 = self.x2 + self.rho * (x - self.x2)
+ self.r2 = self.r2 + self.rho * (r - self.r2)
+ self.u2 = self.u2 + self.rho * (u - self.u2)
+
+ rel_err = torch.linalg.norm(
+ x_prev.flatten() - self.x2.flatten()
+ ) / torch.linalg.norm(self.x2.flatten() + 1e-12)
+
+ if _ > 1 and rel_err < self.crit:
+ self.has_converged = True
+ if self.verbose:
+ print("TGV prox reached convergence")
+ break
+
+ if self.verbose and _ % 100 == 0:
+ primalcost = (
+ torch.linalg.norm(x.flatten() - y.flatten()) ** 2
+ + lambda1 * torch.sum(torch.sqrt(torch.sum(r**2, axis=-1)))
+ + lambda2
+ * torch.sum(
+ torch.sqrt(
+ torch.sum(self.epsilon(self.nabla(x) - r) ** 2, axis=-1)
+ )
+ )
+ )
+ # dualcost = cy - ((y - nablaT(epsilonT(u))) ** 2).sum() / 2.0
+ tmp = torch.max(
+ torch.sqrt(torch.sum(self.epsilon_adjoint(u) ** 2, axis=-1))
+ ) # to check feasibility: the value will be <= lambda1 only at convergence. Since u is not feasible, the dual cost is not reliable: the gap=primalcost-dualcost can be <0 and cannot be used as stopping criterion.
+ u3 = u / torch.maximum(
+ tmp / lambda1, torch.tensor([1], device=tmp.device).type(tmp.dtype)
+ ) # u3 is a scaled version of u, which is feasible. so, its dual cost is a valid, but very rough lower bound of the primal cost.
+ dualcost2 = (
+ cy
+ - torch.sum((y - self.nabla_adjoint(self.epsilon_adjoint(u3))) ** 2)
+ / 2.0
+ ) # we display the best value of dualcost2 computed so far.
+ primalcostlowerbound = max(primalcostlowerbound, dualcost2.item())
+ if self.verbose:
+ print(
+ "Iter: ",
+ _,
+ " Primal cost: ",
+ primalcost.item(),
+ " Rel err:",
+ rel_err,
+ )
+
+ if _ == self.n_it_max - 1:
+ if self.verbose:
+ print(
+ "The algorithm did not converge, stopped after "
+ + str(_ + 1)
+ + " iterations."
+ )
+
+ return self.x2
+
+ @staticmethod
+ def nabla2(x):
+ r"""
+ Applies the finite differences operator associated with tensors of the same shape as x.
+ """
+ b, c, h, w = x.shape
+ u = torch.zeros((b, c, h, w, 2), device=x.device).type(x.dtype)
+ u[:, :, :-1, :, 0] = u[:, :, :-1, :, 0] - x[:, :, :-1]
+ u[:, :, :-1, :, 0] = u[:, :, :-1, :, 0] + x[:, :, 1:]
+ u[:, :, :, :-1, 1] = u[:, :, :, :-1, 1] - x[..., :-1]
+ u[:, :, :, :-1, 1] = u[:, :, :, :-1, 1] + x[..., 1:]
+ return u
+
+ @staticmethod
+ def nabla2_adjoint(x):
+ r"""
+ Applies the adjoint of the finite difference operator.
+ """
+ b, c, h, w = x.shape[:-1]
+ u = torch.zeros((b, c, h, w), device=x.device).type(
+ x.dtype
+ ) # note that we just reversed left and right sides of each line to obtain the transposed operator
+ u[:, :, :-1] = u[:, :, :-1] - x[:, :, :-1, :, 0]
+ u[:, :, 1:] = u[:, :, 1:] + x[:, :, :-1, :, 0]
+ u[..., :-1] = u[..., :-1] - x[..., :-1, 1]
+ u[..., 1:] = u[..., 1:] + x[..., :-1, 1]
+ return u
+
+ @staticmethod
+ def nabla3(x):
+ r"""
+ Applies the finite differences operator associated with tensors of the same shape as x.
+ """
+ b, c, d, h, w = x.shape
+ u = torch.zeros((b, c, d, h, w, 3), device=x.device).type(x.dtype)
+ u[:, :, :-1, :, :, 0] = u[:, :, :-1, :, :, 0] - x[:, :, :-1]
+ u[:, :, :-1, :, :, 0] = u[:, :, :-1, :, :, 0] + x[:, :, 1:]
+ u[:, :, :, :-1, :, 1] = u[:, :, :, :-1, :, 1] - x[:, :, :, :-1]
+ u[:, :, :, :-1, :, 1] = u[:, :, :, :-1, :, 1] + x[:, :, :, 1:]
+ u[:, :, :, :, :-1, 2] = u[:, :, :, :, :-1, 2] - x[:, :, :, :, :-1]
+ u[:, :, :, :, :-1, 2] = u[:, :, :, :, :-1, 2] + x[:, :, :, :, 1:]
+
+ return u
+
+ @staticmethod
+ def nabla3_adjoint(x):
+ r"""
+ Applies the adjoint of the finite difference operator.
+ """
+ b, c, d, h, w = x.shape
+ u = torch.zeros((b, c, d, h, w), device=x.device).type(x.dtype)
+ u[:, :, :-1, :, 0] = u[:, :, :-1, :, 0] - x[:, :, :-1]
+ u[:, :, 1:, :, 0] = u[:, :, 1:, :, 0] + x[:, :, :-1]
+ u[:, :, :, :-1, 1] = u[:, :, :, :-1, 1] - x[:, :, :, :-1]
+ u[:, :, :, 1:, 1] = u[:, :, :, 1:, 1] + x[:, :, :, :-1]
+ u[:, :, :, :, :-1, 2] = u[:, :, :, :, :-1, 2] - x[:, :, :, :, :-1]
+ u[:, :, :, :, 1:, 2] = u[:, :, :, :, 1:, 2] + x[:, :, :, :, :-1]
+
+ return u
+
+ @staticmethod
+ def epsilon2(I): # Simplified
+ r"""
+ Applies the jacobian of a vector field.
+ """
+ b, c, h, w, _ = I.shape
+ G = torch.zeros((b, c, h, w, 4), device=I.device).type(I.dtype)
+ G[:, :, 1:, :, 0] = G[:, :, 1:, :, 0] - I[:, :, :-1, :, 0] # xdy
+ G[..., 0] = G[..., 0] + I[..., 0]
+ G[..., 1:, 1] = G[..., 1:, 1] - I[..., :-1, 0] # xdx
+ G[..., 1:, 1] = G[..., 1:, 1] + I[..., 1:, 0]
+ G[..., 1:, 2] = G[..., 1:, 2] - I[..., :-1, 1] # xdx
+ G[..., 2] = G[..., 2] + I[..., 1]
+ G[:, :, :-1, :, 3] = G[:, :, :-1, :, 3] - I[:, :, :-1, :, 1] # xdy
+ G[:, :, :-1, :, 3] = G[:, :, :-1, :, 3] + I[:, :, 1:, :, 1]
+ return G
+
+ @staticmethod
+ def epsilon2_adjoint(G):
+ r"""
+ Applies the adjoint of the jacobian of a vector field.
+ """
+ b, c, h, w, _ = G.shape
+ I = torch.zeros((b, c, h, w, 2), device=G.device).type(G.dtype)
+ I[:, :, :-1, :, 0] = I[:, :, :-1, :, 0] - G[:, :, 1:, :, 0]
+ I[..., 0] = I[..., 0] + G[..., 0]
+ I[..., :-1, 0] = I[..., :-1, 0] - G[..., 1:, 1]
+ I[..., 1:, 0] = I[..., 1:, 0] + G[..., 1:, 1]
+ I[..., :-1, 1] = I[..., :-1, 1] - G[..., 1:, 2]
+ I[..., 1] = I[..., 1] + G[..., 2]
+ I[:, :, :-1, :, 1] = I[:, :, :-1, :, 1] - G[:, :, :-1, :, 3]
+ I[:, :, 1:, :, 1] = I[:, :, 1:, :, 1] + G[:, :, :-1, :, 3]
+ return I
+
+ @staticmethod
+ def epsilon3(I): # Adapted for 3D matrices
+ r"""
+ Applies the jacobian of a vector field.
+ """
+ b, c, d, h, w = I.shape
+ G = torch.zeros((b, c, d, h, w, 6), device=I.device).type(I.dtype)
+ G[:, :, :, 1:, :, 0] = G[:, :, :, 1:, :, 0] - I[:, :, :, :-1, :, 0] # xdy
+ G[..., 0] = G[..., 0] + I[..., 0]
+ G[..., 1:, :, 1] = G[..., 1:, :, 1] - I[..., :, :-1, 0] # xdx
+ G[..., 1:, :, 1] = G[..., 1:, :, 1] + I[..., :, 1:, 0]
+ G[..., 1:, :, 2] = G[..., 1:, :, 2] - I[..., :, :-1, 1] # xdz
+ G[..., 2] = G[..., 2] + I[..., :, 1, 0]
+ G[:, :, :, :-1, :, 3] = G[:, :, :, :-1, :, 3] - I[:, :, :, :-1, :, 1] # xdy
+ G[:, :, :, :-1, :, 3] = G[:, :, :, :-1, :, 3] + I[:, :, :, 1:, :, 1]
+ G[..., 3] = G[..., 3] + I[..., 0]
+ G[..., 1:, :, 4] = G[..., 1:, :, 4] - I[..., 1:, :, :-1, 2] # xdz
+ G[..., 4] = G[..., 4] + I[..., 1, :, :, 0]
+ G[:, :, :, :, :-1, 5] = G[:, :, :, :, :-1, 5] - I[:, :, :, :, :-1, 2] # xdy
+ G[:, :, :, 1:, :, 5] = G[:, :, :, 1:, :, 5] + I[:, :, :, :, :-1, 2]
+ return G
+
+ @staticmethod
+ def epsilon3_adjoint(G): # Adapted for 3D matrices
+ r"""
+ Applies the adjoint of the jacobian of a vector field.
+ """
+ b, c, d, h, w, _ = G.shape
+ I = torch.zeros((b, c, d, h, w, 3), device=G.device).type(G.dtype)
+ I[:, :, :, :-1, :, 0] = I[:, :, :, :-1, :, 0] - G[:, :, :, 1:, :, 0]
+ I[..., 0] = I[..., 0] + G[..., 0]
+ I[..., :-1, :, 0] = I[..., :-1, :, 0] - G[..., 1:, :, 1]
+ I[..., 1:, :, 0] = I[..., 1:, :, 0] + G[..., 1:, :, 1]
+ I[..., :-1, :, 1] = I[..., :-1, :, 1] - G[..., 1:, :, 2]
+ I[..., 0] = I[..., 0] + G[..., 2]
+ I[:, :, :, :-1, :, 1] = I[:, :, :, :-1, :, 1] - G[:, :, :, :-1, :, 3]
+ I[:, :, :, 1:, :, 1] = I[:, :, :, 1:, :, 1] + G[:, :, :, :-1, :, 3]
+ I[..., 1] = I[..., 1] + G[..., 3]
+ I[..., :-1, :, 2] = I[..., :-1, :, 2] - G[..., 1:, :, 4]
+ I[..., 0] = I[..., 0] + G[..., 4]
+ I[:, :, :, :, :-1, 2] = I[:, :, :, :, :-1, 2] - G[:, :, :, :, :-1, 5]
+ I[:, :, :, 1:, :, 2] = I[:, :, :, 1:, :, 2] + G[:, :, :, :, :-1, 5]
+ return I
+
+
+
+
+