File size: 3,116 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 95 96 97 98 99 100 101 102 103 104 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)
|