crop frames

1 files changed, 111 insertions, 0 deletions

diff --git a/crop_frames.py b/crop_frames.py
new file mode 100644
index 0000000..4e83fd4
--- /dev/null
+++ b/crop_frames.py
@@ -0,0 +1,111 @@
+import os
+import argparse
+import glob
+import operator
+from shutil import rmtree
+from PIL import Image
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--step', default=256, type=int)
+parser.add_argument('--overlap', default=4, type=int)
+parser.add_argument('--scale', default=1, type=int)
+parser.add_argument('cmd', metavar='crop', help='Command')
+parser.add_argument('dir', metavar='dir', help='Directory to process')
+args = parser.parse_args()
+
+def crop_dir():
+    step = args.step - args.overlap
+    files = glob.glob('{}/*.png'.format(args.dir))
+    img = Image.open(files[0])
+    width, height = img.size
+    print("{}x{} {}".format(width, height, step))
+    for x in range(0, width, step):
+        for y in range(0, height, step):
+            w = min(args.step, width - x)
+            h = min(args.step, height - y)
+            crop_dir = "{}/crop_{}_{}_{}_{}".format(args.dir, x, y, w, h)
+            if os.path.exists(crop_dir):
+                rmtree(crop_dir)
+            os.makedirs(crop_dir)
+
+    for raw_fn in files:
+        img = Image.open(raw_fn)
+        fn = os.path.basename(raw_fn)
+        for x in range(0, width, step):
+            for y in range(0, height, step):
+                w = min(args.step, width - x)
+                h = min(args.step, height - y)
+                crop_dir = "{}/crop_{}_{}_{}_{}".format(args.dir, x, y, w, h)
+                print("{}x{} {}x{}".format(x, y, w, h))
+                crop = img.crop((x, y, x + w, y + h))
+                crop.save("{}/{}".format(crop_dir, fn))
+
+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 create_mask(overlap, x, y, w, h):
+    mask_img = Image.new('L', (w, h,), "white")
+    mask = mask_img.load()
+    if x > 0:
+        for i in range(0, overlap, 1):
+            for j in range(0, h, 1):
+                mask[i, j] = (int((i / overlap) * mask[i, j]),)
+    if y > 0:
+        for i in range(0, w, 1):
+            for j in range(0, overlap, 1):
+                mask[i, j] = (int((i / overlap) * mask[i, j]),)
+    return mask_img
+
+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)
+        print("{} {} {} {}".format(x, y, w, h))
+        if x not in widths:
+            overlap = args.overlap
+            if x == 0:
+                overlap = 0
+            ww += (w - overlap) * args.scale
+            widths.append(x)
+        if y not in heights:
+            overlap = args.overlap
+            if x == 0:
+                overlap = 0
+            hh += (h - overlap) * args.scale
+            heights.append(y)
+        # print(w, h)
+        mask = create_mask(overlap, x, y, w * args.scale, h * args.scale)
+        masks[os.path.basename(crop_dir)] = mask
+        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])))
+    for file in files:
+        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))
+            mask = masks[crop_dir]
+            canvas.paste(image, (x * args.scale, y * args.scale,), mask)
+        canvas.save("{}/{}".format(merge_dir, fn))
+
+if __name__ == "__main__":
+    if args.cmd == 'crop':
+        crop_dir()
+    else:
+        merge_dir()
\ No newline at end of file