start/end

1 files changed, 24 insertions, 3 deletions

diff --git a/merge_frames.py b/merge_frames.py
index e265fea..32e0341 100644
--- a/merge_frames.py
+++ b/merge_frames.py
@@ -11,13 +11,29 @@ parser = argparse.ArgumentParser()
 parser.add_argument('--step', default=192, type=int)
 parser.add_argument('--overlap', default=2, type=int)
 parser.add_argument('--scale', default=4, type=int)
-parser.add_argument('dir', metavar='dir', help='Directory to process')
+parser.add_argument('--in_dir', help='Directory to process')
+parser.add_argument('--render_dir', help='Directory to output render to')
+parser.add_argument('--mov', action='store_true', help='whether to create an mp4')
 args = parser.parse_args()
 
 def split_crop_dir(fn):
   crop, x, y, w, h = os.path.basename(fn).split("_")
   return int(x), int(y), int(w), int(h)
 
+def make_movie(merge_dir, dataset):
+  subprocess.call([
+    'ffmpeg',
+    '-hide_banner',
+    '-i', os.path.join(merge_dir, 'frame_%05d.png'),
+    '-y',
+    '-c:v', 'libx264',
+    '-preset', 'slow',
+    '-crf', '19',
+    '-vf', 'fps=25',
+    '-pix_fmt', 'yuv420p',
+    os.path.join(opt.render_dir, dataset + '.mp4')
+  ])
+
 def merge_files(file, crop_dir_list, ww, hh, overlap, merge_dir):
     fn = os.path.basename(file)
     print(fn)
@@ -32,8 +48,10 @@ def merge_files(file, crop_dir_list, ww, hh, overlap, merge_dir):
 
 def merge_dir():
   overlap = args.overlap * args.scale
-  crop_dirs = glob.glob("{}/crop_*".format(args.dir))
-  merge_dir = "{}/merged".format(args.dir)
+  partz = args.dir.trim('/').split('/')
+  dataset = partz[-1]
+  crop_dirs = glob.glob("{}/crop_*".format(args.in_dir))
+  merge_dir = "{}/merged".format(args.in_dir)
   os.makedirs(merge_dir, exist_ok=True)
   ww = 0
   hh = 0
@@ -69,5 +87,8 @@ def merge_dir():
   with Pool(processes=cpu_count()) as pool:
     pool.starmap(merge_files, dataset, chunksize)
 
+  if args.mov:
+    make_movie(merge_dir, dataset)
+
 if __name__ == "__main__":
   merge_dir()