diff options
| -rw-r--r-- | crop_frames.py | 111 |
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 |