mirror of
https://github.com/wassname/scikit-image.git
synced 2026-07-11 19:50:43 +08:00
ENH: Moved nested functions, refactor code
This commit is contained in:
Egor Panfilov
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user