|
from ..patch_match import PyramidPatchMatcher |
|
import functools, os |
|
import numpy as np |
|
from PIL import Image |
|
from tqdm import tqdm |
|
|
|
|
|
class TableManager: |
|
def __init__(self): |
|
pass |
|
|
|
def task_list(self, n): |
|
tasks = [] |
|
max_level = 1 |
|
while (1<<max_level)<=n: |
|
max_level += 1 |
|
for i in range(n): |
|
j = i |
|
for level in range(max_level): |
|
if i&(1<<level): |
|
continue |
|
j |= 1<<level |
|
if j>=n: |
|
break |
|
meta_data = { |
|
"source": i, |
|
"target": j, |
|
"level": level + 1 |
|
} |
|
tasks.append(meta_data) |
|
tasks.sort(key=functools.cmp_to_key(lambda u, v: u["level"]-v["level"])) |
|
return tasks |
|
|
|
def build_remapping_table(self, frames_guide, frames_style, patch_match_engine, batch_size, desc=""): |
|
n = len(frames_guide) |
|
tasks = self.task_list(n) |
|
remapping_table = [[(frames_style[i], 1)] for i in range(n)] |
|
for batch_id in tqdm(range(0, len(tasks), batch_size), desc=desc): |
|
tasks_batch = tasks[batch_id: min(batch_id+batch_size, len(tasks))] |
|
source_guide = np.stack([frames_guide[task["source"]] for task in tasks_batch]) |
|
target_guide = np.stack([frames_guide[task["target"]] for task in tasks_batch]) |
|
source_style = np.stack([frames_style[task["source"]] for task in tasks_batch]) |
|
_, target_style = patch_match_engine.estimate_nnf(source_guide, target_guide, source_style) |
|
for task, result in zip(tasks_batch, target_style): |
|
target, level = task["target"], task["level"] |
|
if len(remapping_table[target])==level: |
|
remapping_table[target].append((result, 1)) |
|
else: |
|
frame, weight = remapping_table[target][level] |
|
remapping_table[target][level] = ( |
|
frame * (weight / (weight + 1)) + result / (weight + 1), |
|
weight + 1 |
|
) |
|
return remapping_table |
|
|
|
def remapping_table_to_blending_table(self, table): |
|
for i in range(len(table)): |
|
for j in range(1, len(table[i])): |
|
frame_1, weight_1 = table[i][j-1] |
|
frame_2, weight_2 = table[i][j] |
|
frame = (frame_1 + frame_2) / 2 |
|
weight = weight_1 + weight_2 |
|
table[i][j] = (frame, weight) |
|
return table |
|
|
|
def tree_query(self, leftbound, rightbound): |
|
node_list = [] |
|
node_index = rightbound |
|
while node_index>=leftbound: |
|
node_level = 0 |
|
while (1<<node_level)&node_index and node_index-(1<<node_level+1)+1>=leftbound: |
|
node_level += 1 |
|
node_list.append((node_index, node_level)) |
|
node_index -= 1<<node_level |
|
return node_list |
|
|
|
def process_window_sum(self, frames_guide, blending_table, patch_match_engine, window_size, batch_size, desc=""): |
|
n = len(blending_table) |
|
tasks = [] |
|
frames_result = [] |
|
for target in range(n): |
|
node_list = self.tree_query(max(target-window_size, 0), target) |
|
for source, level in node_list: |
|
if source!=target: |
|
meta_data = { |
|
"source": source, |
|
"target": target, |
|
"level": level |
|
} |
|
tasks.append(meta_data) |
|
else: |
|
frames_result.append(blending_table[target][level]) |
|
for batch_id in tqdm(range(0, len(tasks), batch_size), desc=desc): |
|
tasks_batch = tasks[batch_id: min(batch_id+batch_size, len(tasks))] |
|
source_guide = np.stack([frames_guide[task["source"]] for task in tasks_batch]) |
|
target_guide = np.stack([frames_guide[task["target"]] for task in tasks_batch]) |
|
source_style = np.stack([blending_table[task["source"]][task["level"]][0] for task in tasks_batch]) |
|
_, target_style = patch_match_engine.estimate_nnf(source_guide, target_guide, source_style) |
|
for task, frame_2 in zip(tasks_batch, target_style): |
|
source, target, level = task["source"], task["target"], task["level"] |
|
frame_1, weight_1 = frames_result[target] |
|
weight_2 = blending_table[source][level][1] |
|
weight = weight_1 + weight_2 |
|
frame = frame_1 * (weight_1 / weight) + frame_2 * (weight_2 / weight) |
|
frames_result[target] = (frame, weight) |
|
return frames_result |
|
|
|
|
|
class FastModeRunner: |
|
def __init__(self): |
|
pass |
|
|
|
def run(self, frames_guide, frames_style, batch_size, window_size, ebsynth_config, save_path=None): |
|
frames_guide = frames_guide.raw_data() |
|
frames_style = frames_style.raw_data() |
|
table_manager = TableManager() |
|
patch_match_engine = PyramidPatchMatcher( |
|
image_height=frames_style[0].shape[0], |
|
image_width=frames_style[0].shape[1], |
|
channel=3, |
|
**ebsynth_config |
|
) |
|
|
|
table_l = table_manager.build_remapping_table(frames_guide, frames_style, patch_match_engine, batch_size, desc="Fast Mode Step 1/4") |
|
table_l = table_manager.remapping_table_to_blending_table(table_l) |
|
table_l = table_manager.process_window_sum(frames_guide, table_l, patch_match_engine, window_size, batch_size, desc="Fast Mode Step 2/4") |
|
|
|
table_r = table_manager.build_remapping_table(frames_guide[::-1], frames_style[::-1], patch_match_engine, batch_size, desc="Fast Mode Step 3/4") |
|
table_r = table_manager.remapping_table_to_blending_table(table_r) |
|
table_r = table_manager.process_window_sum(frames_guide[::-1], table_r, patch_match_engine, window_size, batch_size, desc="Fast Mode Step 4/4")[::-1] |
|
|
|
frames = [] |
|
for (frame_l, weight_l), frame_m, (frame_r, weight_r) in zip(table_l, frames_style, table_r): |
|
weight_m = -1 |
|
weight = weight_l + weight_m + weight_r |
|
frame = frame_l * (weight_l / weight) + frame_m * (weight_m / weight) + frame_r * (weight_r / weight) |
|
frames.append(frame) |
|
frames = [frame.clip(0, 255).astype("uint8") for frame in frames] |
|
if save_path is not None: |
|
for target, frame in enumerate(frames): |
|
Image.fromarray(frame).save(os.path.join(save_path, "%05d.png" % target)) |
|
|
