diff options
| -rw-r--r-- | crop_frames.py | 134 | ||||
| -rw-r--r-- | merge_frames.py | 63 | ||||
| -rw-r--r-- | run.sh | 2 |
3 files changed, 106 insertions, 93 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() 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() @@ -11,7 +11,7 @@ mv "$name/crop*" inputs/L/ python test.py 28 L "" python crop_frames.py merge results/28L/L/ ffmpeg -i ./results/28L/L/merged/frame_%05d.png -y -c:v libx264 -vf fps=25 -pix_fmt yuv420p "./results/${name}_uprez.mp4" -scp "./results/${name}_uprez.mp4" jules@lmno:asdf/neural/uprez/ +# scp "./results/${name}_uprez.mp4" jules@lmno:asdf/neural/uprez/ endd=`date` echo |