mirror of
https://github.com/storytold/storyteller-ml.git
synced 2026-10-17 00:23:14 +00:00
Justin John
94 lines
3.0 KiB
Python
94 lines
3.0 KiB
Python
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)
|