video_blend.py

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)