new crop/merge

1 files changed, 42 insertions, 92 deletions

diff --git a/crop_frames.py b/crop_frames.py
index 516ad37..2462e7b 100644
--- a/crop_frames.py
+++ b/crop_frames.py
@@ -2,110 +2,60 @@ 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('--step', default=256, type=int)
-parser.add_argument('--overlap', default=4, type=int)
+parser.add_argument('--overlap', default=2, type=int)
 parser.add_argument('--scale', default=4, type=int)
-parser.add_argument('cmd', metavar='crop', help='Command')
-parser.add_argument('dir', metavar='dir', help='Directory to process')
+parser.add_argument('--in_dir', help='Directory to process')
+parser.add_argument('--out_dir', default='crops', help='Directory to output to')
 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)
+  step = args.step
+  overlap = args.overlap
+  side = step + overlap * 2
+  files = glob.glob('{}/*.png'.format(args.in_dir))
+  print(files)
 
-    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))
+  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):
+      xx = max(0, x - overlap)
+      yy = max(0, y - overlap)
+      w = min(side, width - xx)
+      h = min(side, height - yy)
+      crop_dir = "{}/crop_{}_{}_{}_{}".format(args.out_dir, xx, yy, w, h)
+      if os.path.exists(crop_dir):
+        rmtree(crop_dir)
+      os.makedirs(crop_dir)
 
-def split_crop_dir(fn):
-    crop, x, y, w, h = os.path.basename(fn).split("_")
-    return int(x), int(y), int(w), int(h)
+  dataset = []
+  for raw_fn in files:
+    dataset.append((raw_fn, width, height, step, side, overlap,))
 
-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
+  chunksize = 3 
+  with Pool(processes=cpu_count()) as pool:
+    pool.starmap(crop_image, dataset, chunksize)
 
-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))
+def crop_image(raw_fn, width, height, step, side, overlap):
+  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):
+      xx = max(0, x - overlap)
+      yy = max(0, y - overlap)
+      w = min(side, width - xx)
+      h = min(side, height - yy)
+      crop_dir = "{}/crop_{}_{}_{}_{}".format(args.out_dir, xx, yy, w, h)
+      # print("{}x{} {}x{}".format(x, y, w, h))
+      crop = img.crop((xx, yy, xx + w, yy + h))
+      crop.save("{}/{}".format(crop_dir, fn))
 
 if __name__ == "__main__":
-    if args.cmd == 'crop':
-        crop_dir()
-    else:
-        merge_dir()
+  crop_dir()