Spaces:
Running
on
Zero
Running
on
Zero
| import cv2 | |
| import numpy as np | |
| import scipy | |
| As = None | |
| prev_states = None | |
| def construct_A(h, w, grad_weight): | |
| indgx_x = [] | |
| indgx_y = [] | |
| indgy_x = [] | |
| indgy_y = [] | |
| vdx = [] | |
| vdy = [] | |
| for i in range(h): | |
| for j in range(w): | |
| if i < h - 1: | |
| indgx_x += [i * w + j] | |
| indgx_y += [i * w + j] | |
| vdx += [1] | |
| indgx_x += [i * w + j] | |
| indgx_y += [(i + 1) * w + j] | |
| vdx += [-1] | |
| if j < w - 1: | |
| indgy_x += [i * w + j] | |
| indgy_y += [i * w + j] | |
| vdy += [1] | |
| indgy_x += [i * w + j] | |
| indgy_y += [i * w + j + 1] | |
| vdy += [-1] | |
| Ix = scipy.sparse.coo_array( | |
| (np.ones(h * w), (np.arange(h * w), np.arange(h * w))), | |
| shape=(h * w, h * w)).tocsc() | |
| Gx = scipy.sparse.coo_array( | |
| (np.array(vdx), (np.array(indgx_x), np.array(indgx_y))), | |
| shape=(h * w, h * w)).tocsc() | |
| Gy = scipy.sparse.coo_array( | |
| (np.array(vdy), (np.array(indgy_x), np.array(indgy_y))), | |
| shape=(h * w, h * w)).tocsc() | |
| As = [] | |
| for i in range(3): | |
| As += [ | |
| scipy.sparse.vstack([Gx * grad_weight[i], Gy * grad_weight[i], Ix]) | |
| ] | |
| return As | |
| # blendI, I1, I2, mask should be RGB unit8 type | |
| # return poissson fusion result (RGB unit8 type) | |
| # I1 and I2: propagated results from previous and subsequent key frames | |
| # mask: pixel selection mask | |
| # blendI: contrastive-preserving blending results of I1 and I2 | |
| def poisson_fusion(blendI, I1, I2, mask, grad_weight=[2.5, 0.5, 0.5]): | |
| global As | |
| global prev_states | |
| Iab = cv2.cvtColor(blendI, cv2.COLOR_BGR2LAB).astype(float) | |
| Ia = cv2.cvtColor(I1, cv2.COLOR_BGR2LAB).astype(float) | |
| Ib = cv2.cvtColor(I2, cv2.COLOR_BGR2LAB).astype(float) | |
| m = (mask > 0).astype(float)[:, :, np.newaxis] | |
| h, w, c = Iab.shape | |
| # fuse the gradient of I1 and I2 with mask | |
| gx = np.zeros_like(Ia) | |
| gy = np.zeros_like(Ia) | |
| gx[:-1, :, :] = (Ia[:-1, :, :] - Ia[1:, :, :]) * (1 - m[:-1, :, :]) + ( | |
| Ib[:-1, :, :] - Ib[1:, :, :]) * m[:-1, :, :] | |
| gy[:, :-1, :] = (Ia[:, :-1, :] - Ia[:, 1:, :]) * (1 - m[:, :-1, :]) + ( | |
| Ib[:, :-1, :] - Ib[:, 1:, :]) * m[:, :-1, :] | |
| # construct A for solving Ax=b | |
| crt_states = (h, w, grad_weight) | |
| if As is None or crt_states != prev_states: | |
| As = construct_A(*crt_states) | |
| prev_states = crt_states | |
| final = [] | |
| for i in range(3): | |
| weight = grad_weight[i] | |
| im_dx = np.clip(gx[:, :, i].reshape(h * w, 1), -100, 100) | |
| im_dy = np.clip(gy[:, :, i].reshape(h * w, 1), -100, 100) | |
| im = Iab[:, :, i].reshape(h * w, 1) | |
| im_mean = im.mean() | |
| im = im - im_mean | |
| A = As[i] | |
| b = np.vstack([im_dx * weight, im_dy * weight, im]) | |
| out = scipy.sparse.linalg.lsqr(A, b) | |
| out_im = (out[0] + im_mean).reshape(h, w, 1) | |
| final += [out_im] | |
| final = np.clip(np.concatenate(final, axis=2), 0, 255) | |
| return cv2.cvtColor(final.astype(np.uint8), cv2.COLOR_LAB2BGR) | |