import argparse import os import platform import struct import subprocess import time from typing import List import cv2 import numpy as np import torch.multiprocessing as mp from numba import njit import sys sys.path.append("./src/ebsynth/") import blender.histogram_blend as histogram_blend from blender.guide import (BaseGuide, ColorGuide, EdgeGuide, PositionalGuide, TemporalGuide) from blender.poisson_fusion import poisson_fusion from blender.video_sequence import VideoSequence from flow.flow_utils import flow_calc from src.video_util import frame_to_video OPEN_EBSYNTH_LOG = False MAX_PROCESS = 8 os_str = platform.system() if os_str == 'Windows': ebsynth_bin = '.\\src\\ebsynth\\deps\\ebsynth\\bin\\ebsynth.exe' elif os_str == 'Linux': ebsynth_bin = './src/ebsynth/deps/ebsynth/bin/ebsynth' elif os_str == 'Darwin': ebsynth_bin = './src/ebsynth/deps/ebsynth/bin/ebsynth.app' else: print('Cannot recognize OS. Run Ebsynth failed.') exit(0) @njit def g_error_mask_loop(H, W, dist1, dist2, output, weight1, weight2): for i in range(H): for j in range(W): if weight1 * dist1[i, j] < weight2 * dist2[i, j]: output[i, j] = 0 else: output[i, j] = 1 if weight1 == 0: output[i, j] = 0 elif weight2 == 0: output[i, j] = 1 def g_error_mask(dist1, dist2, weight1=1, weight2=1): H, W = dist1.shape output = np.empty_like(dist1, dtype=np.byte) g_error_mask_loop(H, W, dist1, dist2, output, weight1, weight2) return output def create_sequence(base_dir, key_ind, key_dir): sequence = VideoSequence(base_dir, key_ind, 'video', key_dir, 'tmp', '%04d.png', '%04d.png') return sequence def process_one_sequence(i, video_sequence: VideoSequence): interval = video_sequence.interval(i) for is_forward in [True, False]: input_seq = video_sequence.get_input_sequence(i, is_forward) output_seq = video_sequence.get_output_sequence(i, is_forward) flow_seq = video_sequence.get_flow_sequence(i, is_forward) key_img_id = i if is_forward else i + 1 key_img = video_sequence.get_key_img(key_img_id) for j in range(interval - 1): i1 = cv2.imread(input_seq[j]) i2 = cv2.imread(input_seq[j + 1]) flow_calc.get_flow(i1, i2, flow_seq[j]) guides: List[BaseGuide] = [ ColorGuide(input_seq), EdgeGuide(input_seq, video_sequence.get_edge_sequence(i, is_forward)), TemporalGuide(key_img, output_seq, flow_seq, video_sequence.get_temporal_sequence(i, is_forward)), PositionalGuide(flow_seq, video_sequence.get_pos_sequence(i, is_forward)) ] weights = [6, 0.5, 0.5, 2] for j in range(interval): # key frame if j == 0: img = cv2.imread(key_img) cv2.imwrite(output_seq[0], img) else: cmd = f'{ebsynth_bin} -style {os.path.abspath(key_img)}' for g, w in zip(guides, weights): cmd += ' ' + g.get_cmd(j, w) cmd += (f' -output {os.path.abspath(output_seq[j])}' ' -searchvoteiters 12 -patchmatchiters 6') if OPEN_EBSYNTH_LOG: print(cmd) subprocess.run(cmd, shell=True, capture_output=not OPEN_EBSYNTH_LOG) def process_sequences(i_arr, video_sequence: VideoSequence): for i in i_arr: process_one_sequence(i, video_sequence) def run_ebsynth(video_sequence: VideoSequence): beg = time.time() processes = [] mp.set_start_method('spawn') n_process = min(MAX_PROCESS, video_sequence.n_seq) cnt = video_sequence.n_seq // n_process remainder = video_sequence.n_seq % n_process prev_idx = 0 for i in range(n_process): task_cnt = cnt + 1 if i < remainder else cnt i_arr = list(range(prev_idx, prev_idx + task_cnt)) prev_idx += task_cnt p = mp.Process(target=process_sequences, args=(i_arr, video_sequence)) p.start() processes.append(p) for p in processes: p.join() end = time.time() print(f'ebsynth: {end-beg}') @njit def assemble_min_error_img_loop(H, W, a, b, error_mask, out): for i in range(H): for j in range(W): if error_mask[i, j] == 0: out[i, j] = a[i, j] else: out[i, j] = b[i, j] def assemble_min_error_img(a, b, error_mask): H, W = a.shape[0:2] out = np.empty_like(a) assemble_min_error_img_loop(H, W, a, b, error_mask, out) return out def load_error(bin_path, img_shape): img_size = img_shape[0] * img_shape[1] with open(bin_path, 'rb') as fp: bytes = fp.read() read_size = struct.unpack('q', bytes[:8]) assert read_size[0] == img_size float_res = struct.unpack('f' * img_size, bytes[8:]) res = np.array(float_res, dtype=np.float32).reshape(img_shape[0], img_shape[1]) return res def process_seq(video_sequence: VideoSequence, i, blend_histogram=True, blend_gradient=True): key1_img = cv2.imread(video_sequence.get_key_img(i)) img_shape = key1_img.shape interval = video_sequence.interval(i) beg_id = video_sequence.get_sequence_beg_id(i) oas = video_sequence.get_output_sequence(i) obs = video_sequence.get_output_sequence(i, False) binas = [x.replace('jpg', 'bin') for x in oas] binbs = [x.replace('jpg', 'bin') for x in obs] obs = [obs[0]] + list(reversed(obs[1:])) inputs = video_sequence.get_input_sequence(i) oas = [cv2.imread(x) for x in oas] obs = [cv2.imread(x) for x in obs] inputs = [cv2.imread(x) for x in inputs] flow_seq = video_sequence.get_flow_sequence(i) dist1s = [] dist2s = [] for i in range(interval - 1): bin_a = binas[i + 1] bin_b = binbs[i + 1] dist1s.append(load_error(bin_a, img_shape)) dist2s.append(load_error(bin_b, img_shape)) lb = 0 ub = 1 beg = time.time() p_mask = None # write key img blend_out_path = video_sequence.get_blending_img(beg_id) cv2.imwrite(blend_out_path, key1_img) for i in range(interval - 1): c_id = beg_id + i + 1 blend_out_path = video_sequence.get_blending_img(c_id) dist1 = dist1s[i] dist2 = dist2s[i] oa = oas[i + 1] ob = obs[i + 1] weight1 = i / (interval - 1) * (ub - lb) + lb weight2 = 1 - weight1 mask = g_error_mask(dist1, dist2, weight1, weight2) if p_mask is not None: flow_path = flow_seq[i] flow = flow_calc.get_flow(inputs[i], inputs[i + 1], flow_path) p_mask = flow_calc.warp(p_mask, flow, 'nearest') mask = p_mask | mask p_mask = mask # Save tmp mask # out_mask = np.expand_dims(mask, 2) # cv2.imwrite(f'mask/mask_{c_id:04d}.jpg', out_mask * 255) min_error_img = assemble_min_error_img(oa, ob, mask) if blend_histogram: hb_res = histogram_blend.blend(oa, ob, min_error_img, (1 - weight1), (1 - weight2)) else: # hb_res = min_error_img tmpa = oa.astype(np.float32) tmpb = ob.astype(np.float32) hb_res = (1 - weight1) * tmpa + (1 - weight2) * tmpb # cv2.imwrite(blend_out_path, hb_res) # gradient blend if blend_gradient: res = poisson_fusion(hb_res, oa, ob, mask) else: res = hb_res cv2.imwrite(blend_out_path, res) end = time.time() print('others:', end - beg) def main(args): global MAX_PROCESS MAX_PROCESS = args.n_proc video_sequence = create_sequence(f'{args.name}', args.key_ind, args.key) if not args.ne: run_ebsynth(video_sequence) blend_histogram = True blend_gradient = args.ps for i in range(video_sequence.n_seq): process_seq(video_sequence, i, blend_histogram, blend_gradient) if args.output: frame_to_video(args.output, video_sequence.blending_dir, args.fps, False) if not args.tmp: video_sequence.remove_out_and_tmp() if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('name', type=str, help='Path to input video') parser.add_argument('--output', type=str, default=None, help='Path to output video') parser.add_argument('--fps', type=float, default=30, help='The FPS of output video') parser.add_argument("--key_ind", type=int, nargs='+', default=[1], help="key frame index") parser.add_argument('--key', type=str, default='keys0', help='The subfolder name of stylized key frames') parser.add_argument('--n_proc', type=int, default=8, help='The max process count') parser.add_argument('-ps', action='store_true', help='Use poisson gradient blending') parser.add_argument( '-ne', action='store_true', help='Do not run ebsynth (use previous ebsynth output)') parser.add_argument('-tmp', action='store_true', help='Keep temporary output') args = parser.parse_args() main(args)