| from ..patch_match import PyramidPatchMatcher |
| import os |
| import numpy as np |
| from PIL import Image |
| from tqdm import tqdm |
|
|
|
|
| class AccurateModeRunner: |
| def __init__(self): |
| pass |
|
|
| def run(self, frames_guide, frames_style, batch_size, window_size, ebsynth_config, desc="Accurate Mode", save_path=None): |
| patch_match_engine = PyramidPatchMatcher( |
| image_height=frames_style[0].shape[0], |
| image_width=frames_style[0].shape[1], |
| channel=3, |
| use_mean_target_style=True, |
| **ebsynth_config |
| ) |
| |
| n = len(frames_style) |
| for target in tqdm(range(n), desc=desc): |
| l, r = max(target - window_size, 0), min(target + window_size + 1, n) |
| remapped_frames = [] |
| for i in range(l, r, batch_size): |
| j = min(i + batch_size, r) |
| source_guide = np.stack([frames_guide[source] for source in range(i, j)]) |
| target_guide = np.stack([frames_guide[target]] * (j - i)) |
| source_style = np.stack([frames_style[source] for source in range(i, j)]) |
| _, target_style = patch_match_engine.estimate_nnf(source_guide, target_guide, source_style) |
| remapped_frames.append(target_style) |
| frame = np.concatenate(remapped_frames, axis=0).mean(axis=0) |
| frame = frame.clip(0, 255).astype("uint8") |
| if save_path is not None: |
| Image.fromarray(frame).save(os.path.join(save_path, "%05d.png" % target)) |
