Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- from __future__ import print_function
- from argparse import ArgumentParser
- import sys
- sys.path.insert(0, 'src')
- import os, random, subprocess, evaluate, shutil
- from utils import exists, list_files
- import pdb
- TMP_DIR = '.fns_frames_%s/' % random.randint(0,99999)
- DEVICE = '/gpu:0'
- BATCH_SIZE = 4
- def build_parser():
- parser = ArgumentParser()
- parser.add_argument('--checkpoint', type=str,
- dest='checkpoint', help='checkpoint directory or .ckpt file',
- metavar='CHECKPOINT', required=True)
- parser.add_argument('--in-path', type=str,
- dest='in_path', help='in video path',
- metavar='IN_PATH', required=True)
- parser.add_argument('--out-path', type=str,
- dest='out', help='path to save processed video to',
- metavar='OUT', required=True)
- parser.add_argument('--tmp-dir', type=str, dest='tmp_dir',
- help='tmp dir for processing', metavar='TMP_DIR',
- default=TMP_DIR)
- parser.add_argument('--device', type=str, dest='device',
- help='device for eval. CPU discouraged. ex: \'/gpu:0\'',
- metavar='DEVICE', default=DEVICE)
- parser.add_argument('--batch-size', type=int,
- dest='batch_size',help='batch size for eval. default 4.',
- metavar='BATCH_SIZE', default=BATCH_SIZE)
- parser.add_argument('--no-disk', type=bool, dest='no_disk',
- help='Don\'t save intermediate files to disk. Default False',
- metavar='NO_DISK', default=False)
- return parser
- def check_opts(opts):
- exists(opts.checkpoint)
- exists(opts.out)
- def main():
- parser = build_parser()
- opts = parser.parse_args()
- if opts.no_disk:
- evaluate.from_pipe(opts)
- else:
- in_dir = os.path.join(opts.tmp_dir, 'in')
- out_dir = os.path.join(opts.tmp_dir, 'out')
- if not os.path.exists(in_dir):
- os.makedirs(in_dir)
- if not os.path.exists(out_dir):
- os.makedirs(out_dir)
- in_args = [
- 'ffmpeg',
- '-i', opts.in_path,
- '%s/frame_%%d.png' % in_dir
- ]
- subprocess.call(" ".join(in_args), shell=True)
- base_names = list_files(in_dir)
- in_files = [os.path.join(in_dir, x) for x in base_names]
- out_files = [os.path.join(out_dir, x) for x in base_names]
- evaluate.ffwd(in_files, out_files, opts.checkpoint, device_t=opts.device,
- batch_size=opts.batch_size)
- fr = 30 # wtf
- out_args = [
- 'ffmpeg',
- '-i', '%s/frame_%%d.png' % out_dir,
- '-f', 'mp4',
- '-q:v', '0',
- '-vcodec', 'mpeg4',
- '-r', str(fr),
- opts.out
- ]
- subprocess.call(" ".join(out_args), shell=True)
- print('Video at: %s' % opts.out)
- shutil.rmtree(opts.tmp_dir)
- if __name__ == '__main__':
- main()
Add Comment
Please, Sign In to add comment