File size: 3,112 Bytes
ff715ca |
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 |
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)
|