| 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)) |
|
|
