Skip to content

Instantly share code, notes, and snippets.

@dstein64
Last active January 21, 2021 03:57
Show Gist options
  • Save dstein64/5dcc67fa43cc0d13d6d4d544095a1382 to your computer and use it in GitHub Desktop.
Save dstein64/5dcc67fa43cc0d13d6d4d544095a1382 to your computer and use it in GitHub Desktop.
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