/
inpaint.py
340 lines (280 loc) · 12.4 KB
/
inpaint.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
import numpy as np
from scipy import sparse
from scipy.sparse.linalg import spsolve
import scipy.ndimage as ndi
from scipy.ndimage 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()
def inpaint_biharmonic(image, mask, *, 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.
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, order='C')
# 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