diff --git a/examples/regularization.py b/examples/regularization.py index ae8513d4..b125e473 100644 --- a/examples/regularization.py +++ b/examples/regularization.py @@ -8,17 +8,35 @@ # Load 3D image image_folder = f"{os.path.dirname(__file__)}/images/" img_path = image_folder + "dicom_3d.npy" -img = darsia.imread(img_path, dimensions=[0.1, 0.1, 0.1]) +img = darsia.imread(img_path, dimensions=[0.1, 0.1, 0.1], dim=3) -# Make regularization parameter heterogenous (for illustration purposes) +print("Image shape: ", img.img.shape) + +# Make regularization parameter heterogenous and anisotropic (for illustration purposes) mu = np.ones_like(img.img) mu[:, 0:100, :] = 0.5 +mu_anisotropic = [10 * mu, mu, mu] + # Regularize image using anisotropic tvd -img_regularized_tvd = darsia.tvd( +img_regularized_anisotropic_tvd = darsia.tvd( img=img, method="heterogeneous bregman", isotropic=False, + weight=mu_anisotropic, + omega=0.1, + max_num_iter=30, + eps=1e-6, + dim=3, + verbose=True, + solver=darsia.Jacobi(maxiter=20), + # solver=darsia.CG(maxiter=20, tol=1e-3), +) + +img_regularized_tvd = darsia.tvd( + img=img.img, + method="heterogeneous bregman", + isotropic=False, weight=mu, omega=0.1, max_num_iter=30, @@ -26,7 +44,7 @@ dim=3, verbose=True, solver=darsia.Jacobi(maxiter=20), - # solver = darsia.CG(maxiter = 20, tol = 1e-3) + # solver=darsia.CG(maxiter=20, tol=1e-3), ) # Regularize image using H1 regularization @@ -35,14 +53,16 @@ mu=1, omega=1, dim=3, - solver=darsia.Jacobi(maxiter=30, tol=1e-6, verbose=True), + solver=darsia.Jacobi(maxiter=30, tol=1e-4, verbose=True), # solver = darsia.MG(3) ) plt.figure("img slice") plt.imshow(img.img[100, :, :]) +plt.figure("regularized anisotropic weight tvd slice") +plt.imshow(img_regularized_anisotropic_tvd.img[100, :, :]) plt.figure("regularized tvd slice") -plt.imshow(img_regularized_tvd.img[100, :, :]) +plt.imshow(img_regularized_tvd[100, :, :]) plt.figure("regularized H1 slice") -plt.imshow(img_regularized_H1.img[100, :, :]) +plt.imshow(img_regularized_h1.img[100, :, :]) plt.show() diff --git a/src/darsia/restoration/split_bregman_tvd.py b/src/darsia/restoration/split_bregman_tvd.py index 24286373..023f0f49 100644 --- a/src/darsia/restoration/split_bregman_tvd.py +++ b/src/darsia/restoration/split_bregman_tvd.py @@ -11,8 +11,8 @@ def split_bregman_tvd( - img: np.ndarray, - mu: Union[float, np.ndarray] = 1.0, + img: Union[np.ndarray, da.Image], + mu: Union[float, np.ndarray, list] = 1.0, omega: Union[float, np.ndarray] = 1.0, ell: Optional[Union[float, np.ndarray]] = None, dim: int = 2, @@ -31,8 +31,10 @@ def split_bregman_tvd( shrinkage step. The regularization term is weighted by the parameter ell. Args: - img (array): image array - mu (float or array): TV penalization parameter + img (array, da.Image): image + mu (float or array or list): inverse TV penalization parameter, if list + it must have the lenght of the dimension and each entry corresponds + to the penalization in the respective direction. omega (float or array): mass penalization parameter ell (float or array): regularization parameter dim (int): spatial dimension of the image @@ -48,16 +50,26 @@ def split_bregman_tvd( array: denoised image """ + # Keep track of input type and convert input image to float for further calculations img_dtype = img.dtype # Store input image and its norm for convergence check img_nrm = np.linalg.norm(img) + # Controll that mu has correct length if its a list + if isinstance(mu, list): + assert len(mu) == dim, "mu must be a list of length dim" + # Feed the solver with parameters, follow the suggestion to use double the weight - # for the diffusion coefficient if no value is provided. + # for the mass coefficient if no value is provided. if ell is None: - ell = 2 * mu + if isinstance(omega, float): + ell = 2 * omega + elif isinstance(omega, np.ndarray): + ell = 2 * omega.copy() + ell[ell == 0] = 1 + solver.update_params( mass_coeff=omega, diffusion_coeff=ell, @@ -66,12 +78,20 @@ def split_bregman_tvd( # Define energy functional for verbose def _functional(x: np.ndarray) -> float: - return 0.5 * np.linalg.norm(omega * (x - img)) ** 2 + sum( - [ - np.sum(np.abs(mu * da.backward_diff(img=x, axis=j, dim=dim))) - for j in range(dim) - ] - ) + if isinstance(mu, list): + return 0.5 * np.linalg.norm(omega * (x - img)) ** 2 + sum( + [ + np.sum(np.abs(mu[j] * da.backward_diff(img=x, axis=j, dim=dim))) + for j in range(dim) + ] + ) + else: + return 0.5 * np.linalg.norm(omega * (x - img)) ** 2 + sum( + [ + np.sum(np.abs(mu * da.backward_diff(img=x, axis=j, dim=dim))) + for j in range(dim) + ] + ) # Define shrinkage operator (shrinks element-wise by k) def _shrink(x: np.ndarray, k: Union[float, np.ndarray]) -> np.ndarray: @@ -86,7 +106,7 @@ def _rhs_function(dt: np.ndarray, bt: np.ndarray) -> np.ndarray: ] ) - # Define initial guess if provided, otherwise start with input image and allovate + # Define initial guess if provided, otherwise start with input image and allocate # zero arrays for the split Bregman variables. if x0 is not None: img0, d0, b0 = x0 @@ -115,6 +135,11 @@ def _rhs_function(dt: np.ndarray, bt: np.ndarray) -> np.ndarray: shrinkage_factor = np.maximum(s - mu / ell, 0) / (s + 1e-18) d = dub * shrinkage_factor[..., None] b = dub - d + elif isinstance(mu, list): + for j in range(dim): + dub = da.backward_diff(img=img_new, axis=j, dim=dim) + b[..., j] + d[..., j] = _shrink(dub, mu[j] / ell) + b[..., j] = dub - d[..., j] else: for j in range(dim): dub = da.backward_diff(img=img_new, axis=j, dim=dim) + b[..., j] @@ -123,16 +148,17 @@ def _rhs_function(dt: np.ndarray, bt: np.ndarray) -> np.ndarray: # Monitor performance relative_increment = np.linalg.norm(img_new - img_iter) / img_nrm + + # Update of result + img_iter = img_new.copy() + if verbose if isinstance(verbose, bool) else iter % verbose == 0: print( f"""Split Bregman iteration {iter} - """ - f"""relative increment: {relative_increment}, """ + f"""relative increment: {round(relative_increment,5)}, """ f"""energy functional: {_functional(img_iter)}""" ) - # Update of result - img_iter = img_new.copy() - # Convergence check if eps is not None: if relative_increment < eps: diff --git a/src/darsia/utils/derivatives.py b/src/darsia/utils/derivatives.py index 66a2feb7..334f398e 100644 --- a/src/darsia/utils/derivatives.py +++ b/src/darsia/utils/derivatives.py @@ -13,7 +13,7 @@ def backward_diff( - img: np.ndarray, axis: int, dim: int = 2, h: Optional[float] = None + img: np.ndarray, axis: int, dim: int = 2, h: Optional[Union[float, list]] = None ) -> np.ndarray: """Backward difference of image matrix in direction of axis. @@ -31,13 +31,16 @@ def backward_diff( if h is None: return np.diff(img, axis=axis, append=da.array_slice(img, axis, -1, None, 1)) else: + if isinstance(h, list): + assert len(h) == dim, "h must be a list of length dim" + h = h[axis] return ( np.diff(img, axis=axis, append=da.array_slice(img, axis, -1, None, 1)) / h ) def forward_diff( - img: np.ndarray, axis: int, dim: int = 2, h: Optional[float] = None + img: np.ndarray, axis: int, dim: int = 2, h: Optional[Union[float, list]] = None ) -> np.ndarray: """Forward difference of image matrix in direction of axis. @@ -45,7 +48,7 @@ def forward_diff( img (np.ndarray): image array axis (int): axis along which the difference is taken dim (int): dimension of image array - h (Optional[float]): grid spacing + h (Optional[float, list]): grid spacing Returns: np.ndarray: forward difference image matrix @@ -55,6 +58,9 @@ def forward_diff( if h is None: return np.diff(img, axis=axis, prepend=da.array_slice(img, axis, 0, 1, 1)) else: + if isinstance(h, list): + assert len(h) == dim, "h must be a list of length dim" + h = h[axis] return np.diff(img, axis=axis, prepend=da.array_slice(img, axis, 0, 1, 1)) / h @@ -62,7 +68,7 @@ def laplace( img: np.ndarray, axis: Optional[int] = None, dim: int = 2, - h: Optional[float] = None, + h: Optional[Union[float, list]] = None, diffusion_coeff: Union[np.ndarray, float] = 1, ) -> np.ndarray: """Laplace operator with homogeneous boundary conditions. @@ -73,7 +79,7 @@ def laplace( img (np.ndarray): image array axis (int): axis along which the difference is taken dim (int): dimension of image array - h (Optional[float]): grid spacing + h (Optional[float, list]): grid spacing diffision_coeff (Optional[np.ndarray]): diffusion coefficient Returns: diff --git a/src/darsia/utils/linear_solvers/cg.py b/src/darsia/utils/linear_solvers/cg.py index faa27609..e6079552 100644 --- a/src/darsia/utils/linear_solvers/cg.py +++ b/src/darsia/utils/linear_solvers/cg.py @@ -1,4 +1,4 @@ -from typing import Optional +from typing import Optional, Union import numpy as np import scipy.sparse as sps @@ -14,7 +14,7 @@ class CG(da.Solver): """ def __call__( - self, x0: np.ndarray, rhs: np.ndarray, h: Optional[float] = None + self, x0: np.ndarray, rhs: np.ndarray, h: Optional[Union[float, list]] = None ) -> np.ndarray: """Solve the problem. diff --git a/src/darsia/utils/linear_solvers/jacobi.py b/src/darsia/utils/linear_solvers/jacobi.py index d8bcdcef..4df1cf12 100644 --- a/src/darsia/utils/linear_solvers/jacobi.py +++ b/src/darsia/utils/linear_solvers/jacobi.py @@ -3,7 +3,7 @@ """ from __future__ import annotations -from typing import Union +from typing import Optional, Union import numpy as np @@ -21,7 +21,9 @@ class Jacobi(da.Solver): """ - def _neighbor_accumulation(self, im: np.ndarray) -> np.ndarray: + def _neighbor_accumulation( + self, im: np.ndarray, h: Optional[Union[float, list]] = None + ) -> np.ndarray: """Accumulation of neighbor pixels. Accumulates for each entry, the entries of neighbors, regardless of dimension. @@ -37,44 +39,72 @@ def _neighbor_accumulation(self, im: np.ndarray) -> np.ndarray: """ im_av: np.ndarray = np.zeros_like(im) - for ax in range(im.ndim): - im_av += np.concatenate( - (da.array_slice(im, ax, 0, 1), da.array_slice(im, ax, 0, -1)), axis=ax - ) + np.concatenate( - (da.array_slice(im, ax, 1, None), da.array_slice(im, ax, -1, None)), - axis=ax, - ) + if h is None: + for ax in range(self.dim): + im_av += np.concatenate( + (da.array_slice(im, ax, 0, 1), da.array_slice(im, ax, 0, -1)), + axis=ax, + ) + np.concatenate( + (da.array_slice(im, ax, 1, None), da.array_slice(im, ax, -1, None)), + axis=ax, + ) + else: + if isinstance(h, float): + h = [h] * self.dim + for ax in range(self.dim): + im_av += ( + np.concatenate( + (da.array_slice(im, ax, 0, 1), da.array_slice(im, ax, 0, -1)), + axis=ax, + ) + + np.concatenate( + ( + da.array_slice(im, ax, 1, None), + da.array_slice(im, ax, -1, None), + ), + axis=ax, + ) + ) / h[ax] ** 2 return im_av - def _diag(self, h: float = 1) -> Union[float, np.ndarray]: + def _diag(self, h: Optional[Union[float, list]] = None) -> Union[float, np.ndarray]: """ Compute diagonal of the stiffness matrix. The stiffness matrix is understood in finite difference sense, with the - possibility to identify pixel sizes with phyiscal mesh sizes. This is a helper + possibility to identify pixel sizes with phyiscal voxel sizes. This is a helper function for the main Jacobi solver. Args: - h (float): mesh diameter + h (float): voxel size Returns: Union[float, np.ndarray]: diagonal of the stiffness matrix """ - return self.mass_coeff + self.diffusion_coeff * 2 * self.dim / h**2 + if h is None: + voxel_weights = self.dim + else: + if isinstance(h, float): + voxel_weights = self.dim / (h**2) + else: + assert len(h) == self.dim, "h must be a float or a list of length dim" + voxel_weights = sum([1 / h[i] ** 2 for i in range(self.dim)]) + + return self.mass_coeff + 2 * self.diffusion_coeff * voxel_weights def __call__( self, x0: np.ndarray, rhs: np.ndarray, - h: float = 1.0, + h: Optional[Union[float, list]] = None, ) -> np.ndarray: """One iteration of a Jacobi solver for linear systems. Args: x0 (np.ndarray): initial guess rhs (np.ndarray): right hand side of the linear system - h (float): mesh diameter + h (optional[float, list]): voxel size Returns: np.ndarray: solution to the linear system @@ -88,14 +118,14 @@ def __call__( if self.tol is None: for _ in range(self.maxiter): x = ( - rhs + self.diffusion_coeff * self._neighbor_accumulation(x) / h**2 + rhs + self.diffusion_coeff * self._neighbor_accumulation(x, h) ) / diag if self.verbose: print(f"Jacobi iteration {_} of {self.maxiter} completed.") else: for _ in range(self.maxiter): x_new = ( - rhs + self.diffusion_coeff * self._neighbor_accumulation(x) / h**2 + rhs + self.diffusion_coeff * self._neighbor_accumulation(x, h) ) / diag err = np.linalg.norm(x_new - x) / np.linalg.norm(x0) if err < self.tol: diff --git a/src/darsia/utils/linear_solvers/solver.py b/src/darsia/utils/linear_solvers/solver.py index 2c6b5531..dfcc4e55 100644 --- a/src/darsia/utils/linear_solvers/solver.py +++ b/src/darsia/utils/linear_solvers/solver.py @@ -58,12 +58,15 @@ def update_params( ) @abc.abstractmethod - def __call__(self, x0: np.ndarray, rhs: np.ndarray) -> np.ndarray: + def __call__( + self, x0: np.ndarray, rhs: np.ndarray, h: Optional[Union[float, list]] = None + ) -> np.ndarray: """Main method of the solver - run the solver. Args: x0 (np.ndarray): initial guess rhs (np.ndarray): right hand side + h (float or list, optional): voxel size Returns: np.ndarray: solution