microproduct/deformation-sentiral/smallbaselineApp/skimage/restoration/inpaint.py

338 lines
13 KiB
Python
Raw Normal View History

2023-08-28 10:17:29 +00:00
import numpy as np
from scipy import sparse
from scipy.sparse.linalg import spsolve
import scipy.ndimage as ndi
from scipy.ndimage.filters import laplace
import skimage
from .._shared import utils
from ..measure import label
from ._inpaint import _build_matrix_inner
def _get_neighborhood(nd_idx, radius, nd_shape):
bounds_lo = np.maximum(nd_idx - radius, 0)
bounds_hi = np.minimum(nd_idx + radius + 1, nd_shape)
return bounds_lo, bounds_hi
def _get_neigh_coef(shape, center, dtype=float):
# Create biharmonic coefficients ndarray
neigh_coef = np.zeros(shape, dtype=dtype)
neigh_coef[center] = 1
neigh_coef = laplace(laplace(neigh_coef))
# extract non-zero locations and values
coef_idx = np.where(neigh_coef)
coef_vals = neigh_coef[coef_idx]
coef_idx = np.stack(coef_idx, axis=0)
return neigh_coef, coef_idx, coef_vals
def _inpaint_biharmonic_single_region(image, mask, out, neigh_coef_full,
coef_vals, raveled_offsets):
"""Solve a (sparse) linear system corresponding to biharmonic inpainting.
This function creates a linear system of the form:
``A @ u = b``
where ``A`` is a sparse matrix, ``b`` is a vector enforcing smoothness and
boundary constraints and ``u`` is the vector of inpainted values to be
(uniquely) determined by solving the linear system.
``A`` is a sparse matrix of shape (n_mask, n_mask) where `n_mask``
corresponds to the number of non-zero values in ``mask`` (i.e. the number
of pixels to be inpainted). Each row in A will have a number of non-zero
values equal to the number of non-zero values in the biharmonic kernel,
``neigh_coef_full``. In practice, biharmonic kernels with reduced extent
are used at the image borders. This matrix, ``A`` is the same for all
image channels (since the same inpainting mask is currently used for all
channels).
``u`` is a dense matrix of shape ``(n_mask, n_channels)`` and represents
the vector of unknown values for each channel.
``b`` is a dense matrix of shape ``(n_mask, n_channels)`` and represents
the desired output of convolving the solution with the biharmonic kernel.
At mask locations where there is no overlap with known values, ``b`` will
have a value of 0. This enforces the biharmonic smoothness constraint in
the interior of inpainting regions. For regions near the boundary that
overlap with known values, the entries in ``b`` enforce boundary conditions
designed to avoid discontinuity with the known values.
"""
n_channels = out.shape[-1]
radius = neigh_coef_full.shape[0] // 2
edge_mask = np.ones(mask.shape, dtype=bool)
edge_mask[(slice(radius, -radius),) * mask.ndim] = 0
boundary_mask = edge_mask * mask
center_mask = ~edge_mask * mask
boundary_pts = np.where(boundary_mask)
boundary_i = np.flatnonzero(boundary_mask)
center_i = np.flatnonzero(center_mask)
mask_i = np.concatenate((boundary_i, center_i))
center_pts = np.where(center_mask)
mask_pts = tuple(
[np.concatenate((b, c)) for b, c in zip(boundary_pts, center_pts)]
)
# Use convolution to predetermine the number of non-zero entries in the
# sparse system matrix.
structure = neigh_coef_full != 0
tmp = ndi.convolve(mask, structure, output=np.uint8, mode='constant')
nnz_matrix = tmp[mask].sum()
# Need to estimate the number of zeros for the right hand side vector.
# The computation below will slightly overestimate the true number of zeros
# due to edge effects (the kernel itself gets shrunk in size near the
# edges, but that isn't accounted for here). We can trim any excess entries
# later.
n_mask = np.count_nonzero(mask)
n_struct = np.count_nonzero(structure)
nnz_rhs_vector_max = n_mask - np.count_nonzero(tmp == n_struct)
# pre-allocate arrays storing sparse matrix indices and values
row_idx_known = np.empty(nnz_rhs_vector_max, dtype=np.intp)
data_known = np.zeros((nnz_rhs_vector_max, n_channels), dtype=out.dtype)
row_idx_unknown = np.empty(nnz_matrix, dtype=np.intp)
col_idx_unknown = np.empty(nnz_matrix, dtype=np.intp)
data_unknown = np.empty(nnz_matrix, dtype=out.dtype)
# cache the various small, non-square Laplacians used near the boundary
coef_cache = {}
# Iterate over masked points near the boundary
mask_flat = mask.reshape(-1)
out_flat = np.ascontiguousarray(out.reshape((-1, n_channels)))
idx_known = 0
idx_unknown = 0
mask_pt_n = -1
boundary_pts = np.stack(boundary_pts, axis=1)
for mask_pt_n, nd_idx in enumerate(boundary_pts):
# Get bounded neighborhood of selected radius
b_lo, b_hi = _get_neighborhood(nd_idx, radius, mask.shape)
# Create (truncated) biharmonic coefficients ndarray
coef_shape = tuple(b_hi - b_lo)
coef_center = tuple(nd_idx - b_lo)
coef_idx, coefs = coef_cache.get((coef_shape, coef_center),
(None, None))
if coef_idx is None:
_ , coef_idx, coefs = _get_neigh_coef(coef_shape,
coef_center,
dtype=out.dtype)
coef_cache[(coef_shape, coef_center)] = (coef_idx, coefs)
# compute corresponding 1d indices into the mask
coef_idx = coef_idx + b_lo[:, np.newaxis]
index1d = np.ravel_multi_index(coef_idx, mask.shape)
# Iterate over masked point's neighborhood
nvals = 0
for coef, i in zip(coefs, index1d):
if mask_flat[i]:
row_idx_unknown[idx_unknown] = mask_pt_n
col_idx_unknown[idx_unknown] = i
data_unknown[idx_unknown] = coef
idx_unknown += 1
else:
data_known[idx_known, :] -= coef * out_flat[i, :]
nvals += 1
if nvals:
row_idx_known[idx_known] = mask_pt_n
idx_known += 1
# Call an efficient Cython-based implementation for all interior points
row_start = mask_pt_n + 1
known_start_idx = idx_known
unknown_start_idx = idx_unknown
nnz_rhs = _build_matrix_inner(
# starting indices
row_start, known_start_idx, unknown_start_idx,
# input arrays
center_i, raveled_offsets, coef_vals, mask_flat,
out_flat,
# output arrays
row_idx_known, data_known, row_idx_unknown, col_idx_unknown,
data_unknown
)
# trim RHS vector values and indices to the exact length
row_idx_known = row_idx_known[:nnz_rhs]
data_known = data_known[:nnz_rhs, :]
# Form sparse matrix of unknown values
sp_shape = (n_mask, out.size)
matrix_unknown = sparse.coo_matrix(
(data_unknown, (row_idx_unknown, col_idx_unknown)), shape=sp_shape
).tocsr()
# Solve linear system for masked points
matrix_unknown = matrix_unknown[:, mask_i]
# dense vectors representing the right hand side for each channel
rhs = np.zeros((n_mask, n_channels), dtype=out.dtype)
rhs[row_idx_known, :] = data_known
# set use_umfpack to False so float32 data is supported
result = spsolve(matrix_unknown, rhs, use_umfpack=False,
permc_spec='MMD_ATA')
if result.ndim == 1:
result = result[:, np.newaxis]
out[mask_pts] = result
return out
@utils.channel_as_last_axis()
@utils.deprecate_multichannel_kwarg(multichannel_position=2)
def inpaint_biharmonic(image, mask, multichannel=False, *,
split_into_regions=False, channel_axis=None):
"""Inpaint masked points in image with biharmonic equations.
Parameters
----------
image : (M[, N[, ..., P]][, C]) ndarray
Input image.
mask : (M[, N[, ..., P]]) ndarray
Array of pixels to be inpainted. Have to be the same shape as one
of the 'image' channels. Unknown pixels have to be represented with 1,
known pixels - with 0.
multichannel : boolean, optional
If True, the last `image` dimension is considered as a color channel,
otherwise as spatial. This argument is deprecated: specify
`channel_axis` instead.
split_into_regions : boolean, optional
If True, inpainting is performed on a region-by-region basis. This is
likely to be slower, but will have reduced memory requirements.
channel_axis : int or None, optional
If None, the image is assumed to be a grayscale (single channel) image.
Otherwise, this parameter indicates which axis of the array corresponds
to channels.
.. versionadded:: 0.19
``channel_axis`` was added in 0.19.
Returns
-------
out : (M[, N[, ..., P]][, C]) ndarray
Input image with masked pixels inpainted.
References
----------
.. [1] S.B.Damelin and N.S.Hoang. "On Surface Completion and Image
Inpainting by Biharmonic Functions: Numerical Aspects",
International Journal of Mathematics and Mathematical Sciences,
Vol. 2018, Article ID 3950312
:DOI:`10.1155/2018/3950312`
.. [2] C. K. Chui and H. N. Mhaskar, MRA Contextual-Recovery Extension of
Smooth Functions on Manifolds, Appl. and Comp. Harmonic Anal.,
28 (2010), 104-113,
:DOI:`10.1016/j.acha.2009.04.004`
Examples
--------
>>> img = np.tile(np.square(np.linspace(0, 1, 5)), (5, 1))
>>> mask = np.zeros_like(img)
>>> mask[2, 2:] = 1
>>> mask[1, 3:] = 1
>>> mask[0, 4:] = 1
>>> out = inpaint_biharmonic(img, mask)
"""
if image.ndim < 1:
raise ValueError('Input array has to be at least 1D')
multichannel = channel_axis is not None
img_baseshape = image.shape[:-1] if multichannel else image.shape
if img_baseshape != mask.shape:
raise ValueError('Input arrays have to be the same shape')
if np.ma.isMaskedArray(image):
raise TypeError('Masked arrays are not supported')
image = skimage.img_as_float(image)
# float16->float32 and float128->float64
float_dtype = utils._supported_float_type(image.dtype)
image = image.astype(float_dtype, copy=False)
mask = mask.astype(bool, copy=False)
if not multichannel:
image = image[..., np.newaxis]
out = np.copy(image)
# Create biharmonic coefficients ndarray
radius = 2
coef_shape = (2 * radius + 1,) * mask.ndim
coef_center = (radius,) * mask.ndim
neigh_coef_full, coef_idx, coef_vals = _get_neigh_coef(coef_shape,
coef_center,
dtype=out.dtype)
# stride for the last spatial dimension
channel_stride_bytes = out.strides[-2]
# offsets to all neighboring non-zero elements in the footprint
offsets = coef_idx - radius
# determine per-channel intensity limits
known_points = image[~mask]
limits = (known_points.min(axis=0), known_points.max(axis=0))
if split_into_regions:
# Split inpainting mask into independent regions
kernel = ndi.generate_binary_structure(mask.ndim, 1)
mask_dilated = ndi.binary_dilation(mask, structure=kernel)
mask_labeled = label(mask_dilated)
mask_labeled *= mask
bbox_slices = ndi.find_objects(mask_labeled)
for idx_region, bb_slice in enumerate(bbox_slices, 1):
# expand object bounding boxes by the biharmonic kernel radius
roi_sl = tuple(slice(max(sl.start - radius, 0),
min(sl.stop + radius, size))
for sl, size in zip(bb_slice, mask_labeled.shape))
# extract only the region surrounding the label of interest
mask_region = mask_labeled[roi_sl] == idx_region
# add slice for axes
roi_sl += (slice(None), )
# copy for contiguity and to account for possible ROI overlap
otmp = out[roi_sl].copy()
# compute raveled offsets for the ROI
ostrides = np.array([s // channel_stride_bytes
for s in otmp[..., 0].strides])
raveled_offsets = np.sum(offsets * ostrides[..., np.newaxis],
axis=0)
_inpaint_biharmonic_single_region(
image[roi_sl], mask_region, otmp,
neigh_coef_full, coef_vals, raveled_offsets
)
# assign output to the
out[roi_sl] = otmp
else:
# compute raveled offsets for output image
ostrides = np.array([s // channel_stride_bytes
for s in out[..., 0].strides])
raveled_offsets = np.sum(offsets * ostrides[..., np.newaxis], axis=0)
_inpaint_biharmonic_single_region(
image, mask, out, neigh_coef_full, coef_vals, raveled_offsets
)
# Handle enormous values on a per-channel basis
np.clip(out, a_min=limits[0], a_max=limits[1], out=out)
if not multichannel:
out = out[..., 0]
return out