merge

1 files changed, 5 insertions, 6 deletions

diff --git a/merge_frames.py b/merge_frames.py
index 650d51a..5d42950 100644
--- a/merge_frames.py
+++ b/merge_frames.py
@@ -14,13 +14,14 @@ parser.add_argument('--scale', default=4, type=int)
 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')
+parser.add_argument('--out_fn', help='mp4 basename')
 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):
+def make_movie(merge_dir, out_fn):
   subprocess.call([
     'ffmpeg',
     '-hide_banner',
@@ -31,7 +32,7 @@ def make_movie(merge_dir, dataset):
     '-crf', '19',
     '-vf', 'fps=25',
     '-pix_fmt', 'yuv420p',
-    os.path.join(args.render_dir, dataset + '.mp4')
+    os.path.join(args.render_dir, out_fn + '.mp4')
   ])
 
 def merge_files(file, crop_dir_list, ww, hh, overlap, merge_dir):
@@ -41,15 +42,13 @@ def merge_files(file, crop_dir_list, ww, hh, overlap, merge_dir):
     for crop_dir_tuple in crop_dir_list:
       x, y, w, h, ix, iy = crop_dir_tuple
       crop_dir = "crop_{}_{}_{}_{}".format(x, y, w, h)
-      image = Image.open("{}/{}/{}".format(args.dir, crop_dir, fn))
+      image = Image.open("{}/{}/{}".format(args.in_dir, crop_dir, fn))
       crop = image.crop((overlap, overlap, image.size[0] - overlap, image.size[1] - overlap))
       canvas.paste(crop, ((x - (ix * args.overlap)) * args.scale, (y - (iy * args.overlap)) * args.scale,))
     canvas.save("{}/{}".format(merge_dir, fn))
 
 def merge_dir():
   overlap = args.overlap * args.scale
-  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)
@@ -88,7 +87,7 @@ def merge_dir():
     pool.starmap(merge_files, dataset, chunksize)
 
   if args.mov:
-    make_movie(merge_dir, dataset)
+    make_movie(merge_dir, args.out_fn)
 
 if __name__ == "__main__":
   merge_dir()