99991
commented
4 years ago Hi, Thanks for your work! It's very nice!
Thank you! :D
And then I want to know estimate_fb_ml()function is base on which paper and "linear equation system" is what, and do you have some idea ?
Maybe the term system of linear equations is more common. The paper is currently being reviewed.
But I find the time it costs is big. So I want to put it to GPU by Cupy or pytorch.
The estimate_fb_cf()-method does not work great on GPU because conjugate gradient descent to solve the system of linear equations converges very slowly without preconditioner. The incomplete thresholded Cholesky preconditioner works best, but it can not be parallelized easily. The Jacobi preconditioner can be used instead, but it is not as good. See the following code.
You can increase the regularization parameter or decrease rtol to get faster results but blurrier images.
from matting import load_image, save_image, stack_images
from matting.util import solve_cg, vec_vec_outer, pixel_coordinates
from matting.util import sparse_conv_matrix
import argparse
import numpy as np
import scipy.sparse
import cupy as cp
import cupyx.scipy.sparse
def estimate_fb_cf_cupy(
image,
alpha,
regularization=1e-5,
neighbors=[(-1, 0), (1, 0), (0, -1), (0, 1)],
max_iter=10000,
atol=0.0,
rtol=1e-5,
print_info=False,
):
h, w = image.shape[:2]
n = w * h
a = alpha.flatten()
# Build sparse linear equation system
U = scipy.sparse.bmat([[
scipy.sparse.diags(a),
scipy.sparse.diags(1 - a)]])
# Directional derivative matrices
Ds = [
sparse_conv_matrix(w, h, [0, dx], [0, dy], [1.0, -1.0])
for dx, dy in neighbors]
S = sum(
D.T.dot(scipy.sparse.diags(regularization + np.abs(D.dot(a)))).dot(D)
for D in Ds)
V = scipy.sparse.bmat([
[S, None],
[None, S]])
A = U.T.dot(U) + V
inv_diag = 1.0 / A.diagonal()
# upload matrix and inverse diagonal GPU
A = cupyx.scipy.sparse.csr_matrix(A)
inv_diag = cp.asarray(inv_diag)
def precondition(x):
return x * inv_diag
foreground = np.zeros((h, w, 3))
background = np.zeros((h, w, 3))
# For each color channel
for channel in range(3):
if print_info:
print("solving channel %d" % (1 + channel))
image_channel = image[:, :, channel].flatten()
b = U.T.dot(image_channel)
b = cp.asarray(b)
norm_b = cp.linalg.norm(b)
x = cp.asarray(np.concatenate([image_channel]*2))
r = b - A.dot(x)
z = precondition(r)
p = z.copy()
rz = np.inner(r, z)
for iteration in range(max_iter):
Ap = A.dot(p)
alpha = rz / cp.inner(p, Ap)
x += alpha * p
r -= alpha * Ap
residual_error = cp.linalg.norm(r)
if print_info:
print("iteration %05d - residual error %e" % (
iteration,
residual_error))
if residual_error < atol or residual_error < rtol * norm_b:
break
z = precondition(r)
beta = 1.0 / rz
rz = cp.inner(r, z)
beta *= rz
p *= beta
p += z
fb = cp.asnumpy(x)
foreground[:, :, channel] = fb[:n].reshape(h, w)
background[:, :, channel] = fb[n:].reshape(h, w)
foreground = np.clip(foreground, 0, 1)
background = np.clip(background, 0, 1)
return foreground, background
def main():
parser = argparse.ArgumentParser(description="Specify input image, alpha and output path.")
parser.add_argument("input", help="Path of input image")
parser.add_argument("alpha", help="Path of alpha image")
parser.add_argument("output", help="Path of output image")
args = parser.parse_args()
image = load_image(args.input, "RGB")
alpha = load_image(args.alpha, "GRAY")
foreground, background = estimate_fb_cf_cupy(
image,
alpha,
regularization=1e-3,
rtol=1e-4,
print_info=True)
save_image(args.output, stack_images(foreground, alpha))
if __name__ == "__main__":
main()The estimate_fb_ml()-method could also be adapted to work on the GPU, but CuPy and PyTorch are not a great fit for that because most of the time would be spent copying around memory and doing very little calculations. There will be a fork of this library (possibly in a few weeks/months) which uses Numba instead of NumPy to accelerate the ml method which is much faster. You can try this code for now:
from matting import load_image, save_image, stack_images
from numba import njit
import numpy as np
import argparse
@njit("void(f4[:, :, :], f4[:, :, :])", cache=True)
def _resize_nearest_3d(dst, src):
h_src, w_src, depth = src.shape
h_dst, w_dst, depth = dst.shape
for y_dst in range(h_dst):
for x_dst in range(w_dst):
x_src = max(0, min(w_src - 1, x_dst * w_src // w_dst))
y_src = max(0, min(h_src - 1, y_dst * h_src // h_dst))
for c in range(depth):
dst[y_dst, x_dst, c] = src[y_src, x_src, c]
@njit("void(f4[:, :], f4[:, :])", cache=True)
def _resize_nearest_2d(dst, src):
h_src, w_src = src.shape
h_dst, w_dst = dst.shape
for y_dst in range(h_dst):
for x_dst in range(w_dst):
x_src = max(0, min(w_src - 1, x_dst * w_src // w_dst))
y_src = max(0, min(h_src - 1, y_dst * h_src // h_dst))
dst[y_dst, x_dst] = src[y_src, x_src]
@njit("Tuple((f4[:, :, :], f4[:, :, :]))(f4[:, :, :], f4[:, :], f4, i4, i4, i4)", cache=True)
def _estimate_fb_ml(
input_image,
input_alpha,
regularization,
n_small_iterations,
n_big_iterations,
small_size,
):
h0, w0, depth = input_image.shape
dtype = np.float32
w_prev = 1
h_prev = 1
F_prev = np.empty((h_prev, w_prev, depth), dtype=dtype)
B_prev = np.empty((h_prev, w_prev, depth), dtype=dtype)
n_levels = int(np.ceil(np.log2(max(w0, h0))))
for i_level in range(n_levels + 1):
w = round(w0 ** (i_level / n_levels))
h = round(h0 ** (i_level / n_levels))
image = np.empty((h, w, depth), dtype=dtype)
alpha = np.empty((h, w), dtype=dtype)
_resize_nearest_3d(image, input_image)
_resize_nearest_2d(alpha, input_alpha)
F = np.empty((h, w, depth), dtype=dtype)
B = np.empty((h, w, depth), dtype=dtype)
_resize_nearest_3d(F, F_prev)
_resize_nearest_3d(B, B_prev)
if w <= small_size and h <= small_size:
n_iter = n_small_iterations
else:
n_iter = n_big_iterations
b = np.zeros((2, depth), dtype=dtype)
dx = [-1, 1, 0, 0]
dy = [0, 0, -1, 1]
for i_iter in range(n_iter):
for y in range(h):
for x in range(w):
a0 = alpha[y, x];
a1 = 1.0 - a0
a00 = a0 * a0
a01 = a0 * a1
# a10 = a01 due to symmetry, can be omitted
a11 = a1 * a1
for c in range(depth):
b[0, c] = a0 * image[y, x, c]
b[1, c] = a1 * image[y, x, c]
for d in range(4):
x2 = max(0, min(w - 1, x + dx[d]))
y2 = max(0, min(h - 1, y + dy[d]))
da = regularization + abs(a0 - alpha[y2, x2])
a00 += da
a11 += da
for c in range(depth):
b[0, c] += da * F[y2, x2, c]
b[1, c] += da * B[y2, x2, c]
determinant = a00 * a11 - a01 * a01
inv_det = 1.0 / determinant
b00 = inv_det * a11
b01 = inv_det * -a01
b11 = inv_det * a00
for c in range(depth):
F_c = b00 * b[0, c] + b01 * b[1, c];
B_c = b01 * b[0, c] + b11 * b[1, c];
F_c = max(0.0, min(1.0, F_c))
B_c = max(0.0, min(1.0, B_c))
F[y, x, c] = F_c
B[y, x, c] = B_c
F_prev = F
B_prev = B
w_prev = w
h_prev = h
return F, B
def estimate_fb_ml_numba(
image,
alpha,
regularization=1e-5,
n_small_iterations=10,
n_big_iterations=2,
small_size=32,
):
return _estimate_fb_ml(
image.astype(np.float32),
alpha.astype(np.float32),
regularization,
n_small_iterations,
n_big_iterations,
small_size)
def main():
parser = argparse.ArgumentParser(description="Specify input image, alpha and output path.")
parser.add_argument("input", help="Path of input image")
parser.add_argument("alpha", help="Path of alpha image")
parser.add_argument("output", help="Path of output image")
args = parser.parse_args()
image = load_image(args.input, "RGB")
alpha = load_image(args.alpha, "GRAY")
foreground, background = estimate_fb_ml_numba(image, alpha)
save_image(args.output, stack_images(foreground, alpha))
if __name__ == "__main__":
main()
onefish51
Hi, Thanks for your work! It's very nice!
I used your approach to estimate foreground and background. But I find the time it costs is big. So I want to put it to GPU by Cupy or pytorch. And then I want to know
estimate_fb_ml()function is base on which paper and "linear equation system" is what, and do you have some idea ?Thank you very much !