From 89784631e79ff3358c267b3e807ef34ea4e9abb8 Mon Sep 17 00:00:00 2001 From: Egor Panfilov Date: Wed, 30 Dec 2015 22:45:27 +0300 Subject: [PATCH] ENH: Massively reworked code to enable nD-support --- doc/examples/filters/plot_inpaint.py | 10 +- skimage/restoration/inpaint.py | 127 +++++++--------------- skimage/restoration/tests/test_inpaint.py | 26 +---- 3 files changed, 51 insertions(+), 112 deletions(-) diff --git a/doc/examples/filters/plot_inpaint.py b/doc/examples/filters/plot_inpaint.py index ba863295..ce661843 100644 --- a/doc/examples/filters/plot_inpaint.py +++ b/doc/examples/filters/plot_inpaint.py @@ -26,18 +26,20 @@ import matplotlib.pyplot as plt from skimage import data, color from skimage.restoration import inpaint -image_orig = color.rgb2gray(data.astronaut()) +image_orig = data.astronaut() # Create mask with three defect regions: left, middle, right respectively -mask = np.zeros_like(image_orig) +mask = np.zeros(image_orig.shape[:-1]) mask[20:60, 0:20] = 1 mask[200:300, 150:170] = 1 mask[50:100, 400:430] = 1 +# Defect image over the same region in each color channel image_defect = image_orig.copy() -image_defect[np.where(mask)] = 0 +for layer in range(image_defect.shape[-1]): + image_defect[np.where(mask)] = 0 -image_result = inpaint.inpaint_biharmonic(image_defect, mask) +image_result = inpaint.inpaint_biharmonic(image_defect, mask, multichannel=True) fig, axes = plt.subplots(ncols=3, nrows=1) diff --git a/skimage/restoration/inpaint.py b/skimage/restoration/inpaint.py index b71ad324..f0c417b3 100644 --- a/skimage/restoration/inpaint.py +++ b/skimage/restoration/inpaint.py @@ -2,9 +2,9 @@ from __future__ import print_function, division import numpy as np import skimage -from .._shared.utils import assert_nD from scipy import sparse from scipy.sparse.linalg import spsolve +from scipy.ndimage.filters import laplace def inpaint_biharmonic(img, mask, multichannel=False): @@ -12,18 +12,18 @@ def inpaint_biharmonic(img, mask, multichannel=False): Parameters ---------- - img : 2D{+color} np.array + img : nD{+color channel} np.ndarray Input image. - mask : 2D np.array + 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, known pixels - with 0. multichannel : boolean, optional - Defines if the last `img` dimension is a color dimension. + If True, the last `img` dimension is considered as a color channel. Returns ------- - out : 2D{+color} np.array + out : nD{+color channel} np.array Input image with masked pixels inpainted. Example @@ -41,102 +41,58 @@ def inpaint_biharmonic(img, mask, multichannel=False): .. [1] N.S.Hoang, S.B.Damelin, "On surface completion and image inpainting by biharmonic functions: numerical aspects", http://www.ima.umn.edu/~damelin/biharmonic - - Realization is based on: - .. [2] John D'Errico, - http://www.mathworks.com/matlabcentral/fileexchange/4551-inpaint-nans, - method 3 """ def _inpaint(img, mask): out = np.copy(img) - out_h, out_w = out.shape - out_l = out.size - def _in_bounds(idx): - if len(idx) == 1: - return 0 <= idx <= out_l - 1 - else: - return (0 <= idx[0] <= out_h - 1) and (0 <= idx[1] <= out_w - 1) - + # 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_mn = np.array(np.where(mask)).T mask_i = np.ravel_multi_index(np.where(mask), mask.shape) - # Initialize sparse matrix - # TODO: only points required for computation might be considered - matrix_unknown = sparse.lil_matrix((np.sum(mask), out.size), dtype=np.int32) - matrix_known = sparse.lil_matrix((np.sum(mask), out.size), dtype=np.int32) + # Find masked points and prepare them to be easily enumerate over + mask_pts = np.array(np.where(mask)).T - # INFO: kernels can be reworked using scipy.signal.convolve2d - # and np.array([0, 1, 0], [1, -4, 1], [0, 1, 0]) + # 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])) - # 1 stage. Find points 2 or more pixels far from edges - kernel = [1, 2, -8, 2, 1, -8, 20, -8, 1, 2, -8, 2, 1] - offset = [-2 * out_w, -out_w - 1, -out_w, -out_w + 1, - -2, -1, 0, 1, 2, out_w - 1, out_w, out_w + 1, 2 * out_w] + # 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)) - for idx, (i, (m, n)) in enumerate(zip(mask_i, mask_mn)): - if 2 <= m <= out_h - 3 and 2 <= n <= out_w - 3: - for k, o in zip(kernel, offset): - if i + o in mask_i: - matrix_unknown[idx, i + o] = k - else: - matrix_known[idx, i + o] = k + # 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) - # 2 stage. Find points 1 pixel far from edges - kernel = [1, 1, -4, 1, 1] - offset = [-out_w, -1, 0, 1, out_w] - - for idx, (i, (m, n)) in enumerate(zip(mask_i, mask_mn)): - if (m in [1, out_h - 2] and 1 <= n <= out_h - 2) or \ - (n in [1, out_w - 2] and 1 <= m <= out_w - 2): - for k, o in zip(kernel, offset): - if i + o in mask_i: - matrix_unknown[idx, i + o] = k - else: - matrix_known[idx, i + o] = k - - # 3 stage. Find points on the edges - kernel = [1, 1, -3, 1, 1] - offset = [-out_w, -1, 0, 1, out_w] - offset_mn = [(-1, 0), (0, -1), (0, 0), (0, 1), (1, 0)] - - for idx, (i, (m, n)) in enumerate(zip(mask_i, mask_mn)): - if (m in [0, out_h - 1] and 1 <= n <= out_w - 2) or \ - (n in [0, out_w - 1] and 1 <= m <= out_h - 2): - for k, o_mn in zip(kernel, offset_mn): - if _in_bounds((m + o_mn[0], n + o_mn[1])): - o = offset[offset_mn.index(o_mn)] - if i + o in mask_i: - matrix_unknown[idx, i + o] = k - else: - matrix_known[idx, i + o] = k - - # 4 stage. Find corner points - kernel = [1, 1, -2, 1, 1] - offset = [-out_w, -1, 0, 1, out_w] - offset_mn = [(-1, 0), (0, -1), (0, 0), (0, 1), (1, 0)] - - for idx, (i, (m, n)) in enumerate(zip(mask_i, mask_mn)): - if m in [0, out_h - 1] and n in [0, out_w - 1]: - for k, o_mn in zip(kernel, offset_mn): - if _in_bounds((m + o_mn[0], n + o_mn[1])): - o = offset[offset_mn.index(o_mn)] - if i + o in mask_i: - matrix_unknown[idx, i + o] = k - else: - matrix_known[idx, i + o] = k + 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 columns + # 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 over defect points + # Solve linear system for masked points matrix_unknown = matrix_unknown[:, mask_i] matrix_unknown = sparse.csr_matrix(matrix_unknown) result = spsolve(matrix_unknown, rhs) @@ -148,15 +104,12 @@ def inpaint_biharmonic(img, mask, multichannel=False): result = result.ravel() - # Put calculated points into the image - for idx, (m, n) in enumerate(mask_mn): - out[m, n] = result[idx] + # 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 - assert_nD(img, 2 + multichannel) - assert_nD(mask, 2) - 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') diff --git a/skimage/restoration/tests/test_inpaint.py b/skimage/restoration/tests/test_inpaint.py index 1fe3e766..abe89df3 100644 --- a/skimage/restoration/tests/test_inpaint.py +++ b/skimage/restoration/tests/test_inpaint.py @@ -14,27 +14,11 @@ def test_inpaint_biharmonic(): mask[0, 4:] = 1 out = inpaint.inpaint_biharmonic(img, mask) ref = np.array( - [[0., 0.0625, 0.25, 0.5625, 0.56947314], - [0., 0.0625, 0.25, 0.47029959, 0.57644628], - [0., 0.0625, 0.24664256, 0.49225207, 0.68956612], - [0., 0.0625, 0.25, 0.5625, 1.], - [0., 0.0625, 0.25, 0.5625, 1.]] - ) - assert_allclose(ref, out) - - -def test_inpaint_edges(): - img = np.tile(np.square(np.linspace(0, 1, 5)), (5, 1)) - mask = np.zeros_like(img) - mask[[0, -1], :] = 1 - mask[:, [0, -1]] = 1 - out = inpaint.inpaint_biharmonic(img, mask) - ref = np.array( - [[0.12199519, 0.15599245, 0.28348214, 0.44445398, 0.48737981], - [0.08799794, 0.0625, 0.25, 0.5625, 0.53030563], - [0.07949863, 0.0625, 0.25, 0.5625, 0.54103709], - [0.08799794, 0.0625, 0.25, 0.5625, 0.53030563], - [0.12199519, 0.15599245, 0.28348214, 0.44445398, 0.48737981]] + [[0., 0.0625, 0.25000000, 0.5625000, 0.73925058], + [0., 0.0625, 0.25000000, 0.5478048, 0.76557821], + [0., 0.0625, 0.25842878, 0.5623079, 0.85927796], + [0., 0.0625, 0.25000000, 0.5625000, 1.00000000], + [0., 0.0625, 0.25000000, 0.5625000, 1.00000000]] ) assert_allclose(ref, out)