medley.py

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)