from distutils.spawn import find_executable | |
import glob | |
import os | |
import random | |
import shutil | |
import subprocess | |
import sys | |
STYLE_IMAGES_PATH = 'paint-by-numbers/test' | |
CONTENT_IMAGE_PATH = 'boston.jpg' | |
WORKSPACE_PATH = 'workspace' | |
CODEC = 'libx264' # switch to 'huffyuv' for lossless (with .avi extension) | |
OUTPUT = 'medley.mov' | |
INIT_REPEATS = 24 # Number of times to repeat initial frame | |
FINAL_REPEATS = 36 # Number of times to repeat final frame | |
NUM_STEPS = 2000 # Number of steps for each style transfer | |
MAX_ITERS = 275 # Maximum number of stylizations | |
assert find_executable('pastiche') and find_executable('ffmpeg') | |
if os.path.exists(WORKSPACE_PATH): | |
sys.stderr.write('Workspace already exists. This may cause problems.\n') | |
style_images = glob.glob(os.path.join(STYLE_IMAGES_PATH, '*.jpg')) | |
random.shuffle(style_images) | |
num_iters = min(len(style_images), MAX_ITERS) | |
frames_workspace = os.path.join(WORKSPACE_PATH, 'frames_workspace') | |
os.makedirs(frames_workspace, exist_ok=True) | |
count = 0 | |
for style_image in style_images: | |
if count >= num_iters: | |
break | |
print(f'Style Image: {style_image}') | |
pastiche_workspace = os.path.join(WORKSPACE_PATH, 'pastiche_workspace') | |
pastiche_argv = [ | |
'pastiche', | |
'--num-steps', str(NUM_STEPS), | |
'--workspace', pastiche_workspace, | |
'--workspace-step', '1', | |
'--content', CONTENT_IMAGE_PATH, | |
'--styles', style_image, | |
'--output', '/dev/null' | |
] | |
code = subprocess.call(pastiche_argv) | |
if code != 0: | |
print('Failure') | |
continue | |
group_name_padding = len(str(num_iters - 1)) | |
group_name = f'{count:0{group_name_padding}d}' | |
group_path = os.path.join(frames_workspace, group_name) | |
os.makedirs(group_path, exist_ok=True) | |
frames = [] | |
for frame in sorted(glob.glob(os.path.join(pastiche_workspace, '*.png'))): | |
frame_idx = int(os.path.splitext(os.path.basename(frame))[0]) | |
divisor = (frame_idx // 50) + 1 | |
if frame_idx % divisor == 0: | |
frames.append(frame) | |
frames = frames + [frames[-1]] * FINAL_REPEATS + list(reversed(frames)) | |
frames = [frames[0]] * INIT_REPEATS + frames | |
num_frames = len(frames) | |
frame_name_padding = len(str((num_frames - 1))) | |
for frame_idx, frame in enumerate(frames): | |
frame_name = f'{frame_idx:0{frame_name_padding}d}' | |
frame_dst = os.path.join(group_path, f'{frame_name}.png') | |
shutil.copy(frame, frame_dst) | |
shutil.rmtree(pastiche_workspace) | |
print('Success') | |
print(f'Count: {count}') | |
count = count + 1 | |
ffmpeg_glob = os.path.join('') | |
ffmpeg_argv = [ | |
'ffmpeg', | |
'-y', | |
'-pattern_type', 'glob', | |
'-i', os.path.join(frames_workspace, '*', '*.png'), | |
'-c:v', CODEC, | |
OUTPUT | |
] | |
subprocess.call(ffmpeg_argv) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment