1 files changed, 63 insertions, 0 deletions
diff --git a/merge_frames.py b/merge_frames.py
new file mode 100644
index 0000000..bae78b6
--- /dev/null
+++ b/merge_frames.py
@@ -0,0 +1,63 @@
+import os
+import argparse
+import glob
+import operator
+from multiprocessing import Pool, cpu_count
+from shutil import rmtree
+from PIL import Image
+
+parser = argparse.ArgumentParser()
+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')
+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 merge_files(file, crop_dir_list, ww, hh):
+    fn = os.path.basename(file)
+    print(fn)
+    canvas = Image.new('RGB', (ww, hh,))
+    for crop_dir_tuple in crop_dir_list:
+      x, y, w, h = crop_dir_tuple
+      crop_dir = "crop_{}_{}_{}_{}".format(x, y, w, h)
+      image = Image.open("{}/{}/{}".format(args.dir, crop_dir, fn))
+      crop = image.crop((overlap, overlap, image.size[0] - overlap * 2, [1] - overlap * 2))
+      canvas.paste(image, (x * args.scale, y * args.scale,), mask)
+    canvas.save("{}/{}".format(merge_dir, fn))
+
+def merge_dir():
+  overlap = args.overlap * args.scale
+  crop_dirs = glob.glob("{}/crop_*".format(args.dir))
+  merge_dir = "{}/merged".format(args.dir)
+  os.makedirs(merge_dir, exist_ok=True)
+  ww = 0
+  hh = 0
+  masks = {}
+  crop_dir_list = []
+  widths = []
+  heights = []
+  for crop_dir in crop_dirs:
+    x, y, w, h = split_crop_dir(crop_dir)
+    if x == 0:
+      hh = max(y + h, hh)
+    if y == 0:
+      ww = max(x + w, ww)
+    crop_dir_list.append((x, y, w, h,))
+  crop_dir_list = sorted(crop_dir_list, key=operator.itemgetter(0, 1))
+
+  print("{}x{}".format(ww, hh))
+
+  files = sorted(glob.glob("{}/*.png".format(crop_dirs[0])))
+  dataset = []
+  for file in files:
+    dataset.append((file, crop_dir_list, ww, hh,))
+
+  chunksize = 3
+  with Pool(processes=cpu_count()) as pool:
+    pool.starmap(merge_files, dataset, chunksize)
+
+if __name__ == "__main__":
+  merge_dir()