Source code for darsia.restoration.tvd

"""Module wrapping TV denoising algoriths from skimage into the realm of DarSIA. These
can be directly integrated in darsia.ConcentrationAnalysis, in particular as part of the
definition of a restoration object.

"""

from typing import Union

import numpy as np
import skimage

import darsia


[docs] class TVD: """Total variation denoising interface. Connects to skimage.restoration.denoise_tv_* methods as well as to the DarSIA implementation of the split Bregman method (:mod:'split_bregman_tvd.py'). """ def __init__(self, key: str = "", **kwargs) -> None: """Constructor. Args: key (str): Prefix for kwargs arguments. """ # Determine method self.method = kwargs.pop(key + "method", "chambolle").lower() # Method-specific parameters if self.method == "heterogeneous bregman": # DarSIA type methods self.omega = kwargs.pop("omega", 1) self.regularization = kwargs.get("regularization", 1.0) # General parameters self.weight = kwargs.pop(key + "weight", 0.1) self.max_num_iter = kwargs.pop(key + "max_num_iter", 200) self.eps = kwargs.pop(key + "eps", 2e-4) self.kwargs = kwargs def __call__( self, img: Union[np.ndarray, darsia.Image] ) -> Union[np.ndarray, darsia.Image]: """Application of TV denoising. Args: img (np.ndarray or Image): image Returns: np.ndarray or Image: upscaled image (same type as input) """ if isinstance(img, np.ndarray): return self._tvd_array(img) elif isinstance(img, darsia.Image): return self._tvd_image(img) else: raise TypeError(f"Input type {type(img)} not supported.") def _tvd_array(self, img: np.ndarray) -> np.ndarray: """Application of anisotropic resizing and tv denoising to numpy array. Args: img (np.ndarray): image Returns: np.ndarray: upscaled image """ # Apply TVD if self.method == "chambolle": return skimage.restoration.denoise_tv_chambolle( img, weight=self.weight, max_num_iter=self.max_num_iter, eps=self.eps, ) elif self.method == "anisotropic bregman": return skimage.restoration.denoise_tv_bregman( img, weight=self.weight, max_num_iter=self.max_num_iter, eps=self.eps, isotropic=False, ) elif self.method == "isotropic bregman": return skimage.restoration.denoise_tv_bregman( img, weight=self.weight, max_num_iter=self.max_num_iter, eps=self.eps, isotropic=True, ) elif self.method == "heterogeneous bregman": return darsia.split_bregman_tvd( img, mu=self.weight, omega=self.omega, ell=self.regularization, max_num_iter=self.max_num_iter, eps=self.eps, **self.kwargs, ) else: raise ValueError(f"Method {self.method} not supported.") def _tvd_image(self, img: darsia.Image) -> darsia.Image: """Application of anisotropic resizing and tv denoising to darsia.Image. Args: img (darsia.Image): image Returns: darsia.Image: upscaled image """ img_copy = img.copy() img_copy.img = self._tvd_array(img.img) return img_copy
[docs] def tvd( img: Union[np.ndarray, darsia.Image], method: str = "chambolle", weight: Union[float, np.ndarray] = 1.0, max_num_iter: int = 200, eps: float = 2e-4, **kwargs, ) -> Union[np.ndarray, darsia.Image]: """Inline application of TVD. Args: img (np.ndarray or Image): image method (str): TVD method weight (float or np.ndarray): weight max_num_iter (int): maximum number of iterations eps (float): tolerance **kwargs: additional arguments - omega (array or float): data fidelity weight for heterogeneous bregman - regularization (float): regularization parameter for heterogeneous bregman Returns: np.ndarray or Image: upscaled image (same type as input) """ tvd = TVD( method=method, weight=weight, max_num_iter=max_num_iter, eps=eps, **kwargs, ) return tvd(img)