mirror of
https://github.com/storytold/storyteller-ml.git
synced 2026-10-29 09:18:21 +00:00
add rerender lora
This commit is contained in:
salt
@@ -0,0 +1,105 @@
|
||||
import os
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from flow.flow_utils import flow_calc, read_flow, read_mask
|
||||
|
||||
|
||||
class BaseGuide:
|
||||
|
||||
def __init__(self):
|
||||
...
|
||||
|
||||
def get_cmd(self, i, weight) -> str:
|
||||
return (f'-guide {os.path.abspath(self.imgs[0])} '
|
||||
f'{os.path.abspath(self.imgs[i])} -weight {weight}')
|
||||
|
||||
|
||||
class ColorGuide(BaseGuide):
|
||||
|
||||
def __init__(self, imgs):
|
||||
super().__init__()
|
||||
self.imgs = imgs
|
||||
|
||||
|
||||
class PositionalGuide(BaseGuide):
|
||||
|
||||
def __init__(self, flow_paths, save_paths):
|
||||
super().__init__()
|
||||
flows = [read_flow(f) for f in flow_paths]
|
||||
masks = [read_mask(f) for f in flow_paths]
|
||||
# TODO: modify the format of flow to numpy
|
||||
H, W = flows[0].shape[2:]
|
||||
first_img = PositionalGuide.__generate_first_img(H, W)
|
||||
prev_img = first_img
|
||||
imgs = [first_img]
|
||||
cid = 0
|
||||
for flow, mask in zip(flows, masks):
|
||||
cur_img = flow_calc.warp(prev_img, flow,
|
||||
'nearest').astype(np.uint8)
|
||||
cur_img = cv2.inpaint(cur_img, mask, 30, cv2.INPAINT_TELEA)
|
||||
prev_img = cur_img
|
||||
imgs.append(cur_img)
|
||||
cid += 1
|
||||
cv2.imwrite(f'guide/{cid}.jpg', mask)
|
||||
|
||||
for path, img in zip(save_paths, imgs):
|
||||
cv2.imwrite(path, img)
|
||||
self.imgs = save_paths
|
||||
|
||||
@staticmethod
|
||||
def __generate_first_img(H, W):
|
||||
Hs = np.linspace(0, 1, H)
|
||||
Ws = np.linspace(0, 1, W)
|
||||
i, j = np.meshgrid(Hs, Ws, indexing='ij')
|
||||
r = (i * 255).astype(np.uint8)
|
||||
g = (j * 255).astype(np.uint8)
|
||||
b = np.zeros(r.shape)
|
||||
res = np.stack((b, g, r), 2)
|
||||
return res
|
||||
|
||||
|
||||
class EdgeGuide(BaseGuide):
|
||||
|
||||
def __init__(self, imgs, save_paths):
|
||||
super().__init__()
|
||||
edges = [EdgeGuide.__generate_edge(cv2.imread(img)) for img in imgs]
|
||||
for path, img in zip(save_paths, edges):
|
||||
cv2.imwrite(path, img)
|
||||
self.imgs = save_paths
|
||||
|
||||
@staticmethod
|
||||
def __generate_edge(img):
|
||||
filter = np.array([[0, -1, 0], [-1, 4, -1], [0, -1, 0]])
|
||||
res = cv2.filter2D(img, -1, filter)
|
||||
return res
|
||||
|
||||
|
||||
class TemporalGuide(BaseGuide):
|
||||
|
||||
def __init__(self, key_img, stylized_imgs, flow_paths, save_paths):
|
||||
super().__init__()
|
||||
self.flows = [read_flow(f) for f in flow_paths]
|
||||
self.masks = [read_mask(f) for f in flow_paths]
|
||||
self.stylized_imgs = stylized_imgs
|
||||
self.imgs = save_paths
|
||||
|
||||
first_img = cv2.imread(key_img)
|
||||
cv2.imwrite(self.imgs[0], first_img)
|
||||
|
||||
def get_cmd(self, i, weight) -> str:
|
||||
if i == 0:
|
||||
warped_img = self.stylized_imgs[0]
|
||||
else:
|
||||
prev_img = cv2.imread(self.stylized_imgs[i - 1])
|
||||
warped_img = flow_calc.warp(prev_img, self.flows[i - 1],
|
||||
'nearest').astype(np.uint8)
|
||||
|
||||
warped_img = cv2.inpaint(warped_img, self.masks[i - 1], 30,
|
||||
cv2.INPAINT_TELEA)
|
||||
|
||||
cv2.imwrite(self.imgs[i], warped_img)
|
||||
|
||||
return super().get_cmd(i, weight)
|
||||
|
||||
@@ -0,0 +1,50 @@
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
|
||||
def histogram_transform(img: np.ndarray, means: np.ndarray, stds: np.ndarray,
|
||||
target_means: np.ndarray, target_stds: np.ndarray):
|
||||
means = means.reshape((1, 1, 3))
|
||||
stds = stds.reshape((1, 1, 3))
|
||||
target_means = target_means.reshape((1, 1, 3))
|
||||
target_stds = target_stds.reshape((1, 1, 3))
|
||||
x = img.astype(np.float32)
|
||||
x = (x - means) * target_stds / stds + target_means
|
||||
# x = np.round(x)
|
||||
# x = np.clip(x, 0, 255)
|
||||
# x = x.astype(np.uint8)
|
||||
return x
|
||||
|
||||
|
||||
def blend(a: np.ndarray,
|
||||
b: np.ndarray,
|
||||
min_error: np.ndarray,
|
||||
weight1=0.5,
|
||||
weight2=0.5):
|
||||
a = cv2.cvtColor(a, cv2.COLOR_BGR2Lab)
|
||||
b = cv2.cvtColor(b, cv2.COLOR_BGR2Lab)
|
||||
min_error = cv2.cvtColor(min_error, cv2.COLOR_BGR2Lab)
|
||||
a_mean = np.mean(a, axis=(0, 1))
|
||||
a_std = np.std(a, axis=(0, 1))
|
||||
b_mean = np.mean(b, axis=(0, 1))
|
||||
b_std = np.std(b, axis=(0, 1))
|
||||
min_error_mean = np.mean(min_error, axis=(0, 1))
|
||||
min_error_std = np.std(min_error, axis=(0, 1))
|
||||
|
||||
t_mean_val = 0.5 * 256
|
||||
t_std_val = (1 / 36) * 256
|
||||
t_mean = np.ones([3], dtype=np.float32) * t_mean_val
|
||||
t_std = np.ones([3], dtype=np.float32) * t_std_val
|
||||
a = histogram_transform(a, a_mean, a_std, t_mean, t_std)
|
||||
|
||||
b = histogram_transform(b, b_mean, b_std, t_mean, t_std)
|
||||
ab = (a * weight1 + b * weight2 - t_mean_val) / 0.5 + t_mean_val
|
||||
ab_mean = np.mean(ab, axis=(0, 1))
|
||||
ab_std = np.std(ab, axis=(0, 1))
|
||||
ab = histogram_transform(ab, ab_mean, ab_std, min_error_mean,
|
||||
min_error_std)
|
||||
ab = np.round(ab)
|
||||
ab = np.clip(ab, 0, 255)
|
||||
ab = ab.astype(np.uint8)
|
||||
ab = cv2.cvtColor(ab, cv2.COLOR_Lab2BGR)
|
||||
return ab
|
||||
@@ -0,0 +1,93 @@
|
||||
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)
|
||||
@@ -0,0 +1,189 @@
|
||||
import os
|
||||
import shutil
|
||||
|
||||
|
||||
class VideoSequence:
|
||||
|
||||
def __init__(self,
|
||||
base_dir,
|
||||
beg_frame,
|
||||
end_frame,
|
||||
interval,
|
||||
input_subdir='videos',
|
||||
key_subdir='keys0',
|
||||
tmp_subdir='tmp',
|
||||
input_format='frame%04d.jpg',
|
||||
key_format='%04d.jpg',
|
||||
out_subdir_format='out_%d',
|
||||
blending_out_subdir='blend',
|
||||
output_format='%04d.jpg'):
|
||||
if (end_frame - beg_frame) % interval != 0:
|
||||
end_frame -= (end_frame - beg_frame) % interval
|
||||
|
||||
self.__base_dir = base_dir
|
||||
self.__input_dir = os.path.join(base_dir, input_subdir)
|
||||
self.__key_dir = os.path.join(base_dir, key_subdir)
|
||||
self.__tmp_dir = os.path.join(base_dir, tmp_subdir)
|
||||
self.__input_format = input_format
|
||||
self.__blending_out_dir = os.path.join(base_dir, blending_out_subdir)
|
||||
self.__key_format = key_format
|
||||
self.__out_subdir_format = out_subdir_format
|
||||
self.__output_format = output_format
|
||||
self.__beg_frame = beg_frame
|
||||
self.__end_frame = end_frame
|
||||
self.__interval = interval
|
||||
self.__n_seq = (end_frame - beg_frame) // interval
|
||||
self.__make_out_dirs()
|
||||
os.makedirs(self.__tmp_dir, exist_ok=True)
|
||||
|
||||
@property
|
||||
def beg_frame(self):
|
||||
return self.__beg_frame
|
||||
|
||||
@property
|
||||
def end_frame(self):
|
||||
return self.__end_frame
|
||||
|
||||
@property
|
||||
def n_seq(self):
|
||||
return self.__n_seq
|
||||
|
||||
@property
|
||||
def interval(self):
|
||||
return self.__interval
|
||||
|
||||
@property
|
||||
def blending_dir(self):
|
||||
return os.path.abspath(self.__blending_out_dir)
|
||||
|
||||
def remove_out_and_tmp(self):
|
||||
for i in range(self.n_seq + 1):
|
||||
out_dir = self.__get_out_subdir(i)
|
||||
shutil.rmtree(out_dir)
|
||||
shutil.rmtree(self.__tmp_dir)
|
||||
|
||||
def get_input_sequence(self, i, is_forward=True):
|
||||
beg_id = self.get_sequence_beg_id(i)
|
||||
end_id = self.get_sequence_beg_id(i + 1)
|
||||
if is_forward:
|
||||
id_list = list(range(beg_id, end_id))
|
||||
else:
|
||||
id_list = list(range(end_id, beg_id, -1))
|
||||
path_dir = [
|
||||
os.path.join(self.__input_dir, self.__input_format % id)
|
||||
for id in id_list
|
||||
]
|
||||
return path_dir
|
||||
|
||||
def get_output_sequence(self, i, is_forward=True):
|
||||
beg_id = self.get_sequence_beg_id(i)
|
||||
end_id = self.get_sequence_beg_id(i + 1)
|
||||
if is_forward:
|
||||
id_list = list(range(beg_id, end_id))
|
||||
else:
|
||||
i += 1
|
||||
id_list = list(range(end_id, beg_id, -1))
|
||||
out_subdir = self.__get_out_subdir(i)
|
||||
path_dir = [
|
||||
os.path.join(out_subdir, self.__output_format % id)
|
||||
for id in id_list
|
||||
]
|
||||
return path_dir
|
||||
|
||||
def get_temporal_sequence(self, i, is_forward=True):
|
||||
beg_id = self.get_sequence_beg_id(i)
|
||||
end_id = self.get_sequence_beg_id(i + 1)
|
||||
if is_forward:
|
||||
id_list = list(range(beg_id, end_id))
|
||||
else:
|
||||
i += 1
|
||||
id_list = list(range(end_id, beg_id, -1))
|
||||
tmp_dir = self.__get_tmp_out_subdir(i)
|
||||
path_dir = [
|
||||
os.path.join(tmp_dir, 'temporal_' + self.__output_format % id)
|
||||
for id in id_list
|
||||
]
|
||||
return path_dir
|
||||
|
||||
def get_edge_sequence(self, i, is_forward=True):
|
||||
beg_id = self.get_sequence_beg_id(i)
|
||||
end_id = self.get_sequence_beg_id(i + 1)
|
||||
if is_forward:
|
||||
id_list = list(range(beg_id, end_id))
|
||||
else:
|
||||
i += 1
|
||||
id_list = list(range(end_id, beg_id, -1))
|
||||
tmp_dir = self.__get_tmp_out_subdir(i)
|
||||
path_dir = [
|
||||
os.path.join(tmp_dir, 'edge_' + self.__output_format % id)
|
||||
for id in id_list
|
||||
]
|
||||
return path_dir
|
||||
|
||||
def get_pos_sequence(self, i, is_forward=True):
|
||||
beg_id = self.get_sequence_beg_id(i)
|
||||
end_id = self.get_sequence_beg_id(i + 1)
|
||||
if is_forward:
|
||||
id_list = list(range(beg_id, end_id))
|
||||
else:
|
||||
i += 1
|
||||
id_list = list(range(end_id, beg_id, -1))
|
||||
tmp_dir = self.__get_tmp_out_subdir(i)
|
||||
path_dir = [
|
||||
os.path.join(tmp_dir, 'pos_' + self.__output_format % id)
|
||||
for id in id_list
|
||||
]
|
||||
return path_dir
|
||||
|
||||
def get_flow_sequence(self, i, is_forward=True):
|
||||
beg_id = self.get_sequence_beg_id(i)
|
||||
end_id = self.get_sequence_beg_id(i + 1)
|
||||
if is_forward:
|
||||
id_list = list(range(beg_id, end_id - 1))
|
||||
path_dir = [
|
||||
os.path.join(self.__tmp_dir, 'flow_f_%04d.npy' % id)
|
||||
for id in id_list
|
||||
]
|
||||
else:
|
||||
id_list = list(range(end_id, beg_id + 1, -1))
|
||||
path_dir = [
|
||||
os.path.join(self.__tmp_dir, 'flow_b_%04d.npy' % id)
|
||||
for id in id_list
|
||||
]
|
||||
|
||||
return path_dir
|
||||
|
||||
def get_input_img(self, i):
|
||||
return os.path.join(self.__input_dir, self.__input_format % i)
|
||||
|
||||
def get_key_img(self, i):
|
||||
sequence_beg_id = self.get_sequence_beg_id(i)
|
||||
return os.path.join(self.__key_dir,
|
||||
self.__key_format % sequence_beg_id)
|
||||
|
||||
def get_blending_img(self, i):
|
||||
return os.path.join(self.__blending_out_dir, self.__output_format % i)
|
||||
|
||||
def get_sequence_beg_id(self, i):
|
||||
return i * self.__interval + self.__beg_frame
|
||||
|
||||
def __get_out_subdir(self, i):
|
||||
dir_id = self.get_sequence_beg_id(i)
|
||||
out_subdir = os.path.join(self.__base_dir,
|
||||
self.__out_subdir_format % dir_id)
|
||||
return out_subdir
|
||||
|
||||
def __get_tmp_out_subdir(self, i):
|
||||
dir_id = self.get_sequence_beg_id(i)
|
||||
tmp_out_subdir = os.path.join(self.__tmp_dir,
|
||||
self.__out_subdir_format % dir_id)
|
||||
return tmp_out_subdir
|
||||
|
||||
def __make_out_dirs(self):
|
||||
os.makedirs(self.__base_dir, exist_ok=True)
|
||||
os.makedirs(self.__blending_out_dir, exist_ok=True)
|
||||
for i in range(self.__n_seq + 1):
|
||||
out_subdir = self.__get_out_subdir(i)
|
||||
tmp_subdir = self.__get_tmp_out_subdir(i)
|
||||
os.makedirs(out_subdir, exist_ok=True)
|
||||
os.makedirs(tmp_subdir, exist_ok=True)
|
||||
Reference in New Issue
Block a user