From 893d93749116aa7605191b7ae03bd4cf11a4e323 Mon Sep 17 00:00:00 2001 From: Egor Panfilov Date: Fri, 15 Jan 2016 13:59:32 +0300 Subject: [PATCH] ENH: Moved nested functions, refactor code --- skimage/restoration/inpaint.py | 164 +++++++++++++++++---------------- 1 file changed, 83 insertions(+), 81 deletions(-) diff --git a/skimage/restoration/inpaint.py b/skimage/restoration/inpaint.py index f0c417b3..1b318b97 100644 --- a/skimage/restoration/inpaint.py +++ b/skimage/restoration/inpaint.py @@ -1,4 +1,4 @@ -from __future__ import print_function, division +from __future__ import division import numpy as np import skimage @@ -7,23 +7,89 @@ from scipy.sparse.linalg import spsolve from scipy.ndimage.filters import laplace +def _get_neighborhood(nd_idx, radius, nd_shape): + bounds_lo = (nd_idx - radius).clip(min=0) + bounds_hi = (nd_idx + radius + 1).clip(max=nd_shape) + return bounds_lo, bounds_hi + + +def _inpaint_biharmonic_single_channel(img, mask, out, limits): + # Initialize sparse matrices + matrix_unknown = sparse.lil_matrix((np.sum(mask), out.size)) + matrix_known = sparse.lil_matrix((np.sum(mask), out.size)) + + # Find indexes of masked points in flatten array + mask_i = np.ravel_multi_index(np.where(mask), mask.shape) + + # Find masked points and prepare them to be easily enumerate over + mask_pts = np.array(np.where(mask)).T + + # Iterate over masked points + for mask_pt_n, mask_pt_idx in enumerate(mask_pts): + # Get bounded neighborhood of selected radius + b_lo, b_hi = _get_neighborhood(mask_pt_idx, 2, out.shape) + + # Create biharmonic coefficients ndarray + neigh_coef = np.zeros(b_hi - b_lo) + neigh_coef[tuple(mask_pt_idx - b_lo)] = 1 + neigh_coef = laplace(laplace(neigh_coef)) + + # Iterate over masked point's neighborhood + it_inner = np.nditer(neigh_coef, flags=['multi_index']) + for coef in it_inner: + if coef == 0: + continue + tmp_pt_idx = np.add(b_lo, it_inner.multi_index) + tmp_pt_i = np.ravel_multi_index(tmp_pt_idx, mask.shape) + + if mask[tuple(tmp_pt_idx)]: + matrix_unknown[mask_pt_n, tmp_pt_i] = coef + else: + matrix_known[mask_pt_n, tmp_pt_i] = coef + + # Prepare diagonal matrix + flat_diag_image = sparse.dia_matrix((out.flatten(), np.array([0])), + shape=(out.size, out.size)) + + # Calculate right hand side as a sum of known matrix's columns + matrix_known = matrix_known.tocsr() + rhs = -(matrix_known * flat_diag_image).sum(axis=1) + + # Solve linear system for masked points + matrix_unknown = matrix_unknown[:, mask_i] + matrix_unknown = sparse.csr_matrix(matrix_unknown) + result = spsolve(matrix_unknown, rhs) + + # Handle enormous values + result = np.clip(result, *limits) + + result = result.ravel() + + # Substitute masked points with inpainted versions + for mask_pt_n, mask_pt_idx in enumerate(mask_pts): + out[tuple(mask_pt_idx)] = result[mask_pt_n] + + return out + + def inpaint_biharmonic(img, mask, multichannel=False): """Inpaint masked points in image with biharmonic equations. Parameters ---------- - img : nD{+color channel} np.ndarray + img : (M, N[, ..., P][, C]) ndarray Input image. - mask : nD np.ndarray - Array of pixels to be inpainted. Have to be the same size as one - of the 'img' channels. Unknown pixels has to be represented with 1, + mask : (M, N[, ..., P]) ndarray + Array of pixels to be inpainted. Have to be the same shape as one + of the 'img' channels. Unknown pixels have to be represented with 1, known pixels - with 0. multichannel : boolean, optional - If True, the last `img` dimension is considered as a color channel. + If True, the last `img` dimension is considered as a color channel, + otherwise as spatial. Returns ------- - out : nD{+color channel} np.array + out : (M, N[, ..., P][, C] ndarray Input image with masked pixels inpainted. Example @@ -43,73 +109,6 @@ def inpaint_biharmonic(img, mask, multichannel=False): http://www.ima.umn.edu/~damelin/biharmonic """ - def _inpaint(img, mask): - out = np.copy(img) - - # Initialize sparse matrices - matrix_unknown = sparse.lil_matrix((np.sum(mask), out.size)) - matrix_known = sparse.lil_matrix((np.sum(mask), out.size)) - - def _get_neighborhood(idx, radii): - bounds_lo = (idx - radii).clip(min=0) - bounds_hi = (idx + np.add(radii, 1)).clip(max=out.shape) - return bounds_lo, bounds_hi - - # Find indexes of masked points in flatten array - mask_i = np.ravel_multi_index(np.where(mask), mask.shape) - - # Find masked points and prepare them to be easily enumerate over - mask_pts = np.array(np.where(mask)).T - - # Iterate over masked points - for mask_pt_n, mask_pt_idx in enumerate(mask_pts): - # Get bounded neighborhood of selected radii - b_lo, b_hi = _get_neighborhood(mask_pt_idx, radii=np.array([2])) - - # Create biharmonic coefficients ndarray - neigh_coef = np.zeros(b_hi - b_lo) - neigh_coef[tuple(mask_pt_idx - b_lo)] = 1 - neigh_coef = laplace(laplace(neigh_coef)) - - # Iterate over masked point's neighborhood - it_inner = np.nditer(neigh_coef, flags=['multi_index']) - for coef in it_inner: - if coef == 0: - continue - tmp_pt_idx = np.add(b_lo, it_inner.multi_index) - tmp_pt_i = np.ravel_multi_index(tmp_pt_idx, mask.shape) - - if mask[tuple(tmp_pt_idx)]: - matrix_unknown[mask_pt_n, tmp_pt_i] = coef - else: - matrix_known[mask_pt_n, tmp_pt_i] = coef - - # Prepare diagonal matrix - flat_diag_image = sparse.dia_matrix((out.flatten(), np.array([0])), - shape=(out.size, out.size)) - - # Calculate right hand side as a sum of known matrix's columns - matrix_known = matrix_known.tocsr() - rhs = -(matrix_known * flat_diag_image).sum(axis=1) - - # Solve linear system for masked points - matrix_unknown = matrix_unknown[:, mask_i] - matrix_unknown = sparse.csr_matrix(matrix_unknown) - result = spsolve(matrix_unknown, rhs) - - # Handle enormous values - # TODO: consider images in [-1:1] scale - result[np.where(result < 0)] = 0 - result[np.where(result > 1)] = 1 - - result = result.ravel() - - # Substitute masked points with inpainted versions - for mask_pt_n, mask_pt_idx in enumerate(mask_pts): - out[tuple(mask_pt_idx)] = result[mask_pt_n] - - return out - img_baseshape = img.shape[:-1] if multichannel else img.shape if img_baseshape != mask.shape: raise ValueError('Input arrays have to be the same shape') @@ -119,16 +118,19 @@ def inpaint_biharmonic(img, mask, multichannel=False): img = skimage.img_as_float(img) mask = mask.astype(np.bool) - - if not multichannel: - img = img.reshape(img.shape + (1,)) - out = np.zeros_like(img) - + if not multichannel: + img = img[..., np.newaxis] + + out = np.copy(img) + for i in range(img.shape[-1]): - out[..., i] = _inpaint(img[..., i], mask) + known_points = img[..., i][~mask] + limits = (np.min(known_points), np.max(known_points)) + _inpaint_biharmonic_single_channel(img[..., i], mask, + out[..., i], limits) if not multichannel: - out = out.reshape(out.shape[:-1]) + out = out[..., 0] return out