1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
|
import os
import argparse
import glob
import operator
from multiprocessing import Pool, cpu_count
from shutil import rmtree
from PIL import Image
from math import floor
parser = argparse.ArgumentParser()
parser.add_argument('--step', default=192, type=int)
parser.add_argument('--overlap', default=2, type=int)
parser.add_argument('--scale', default=4, type=int)
parser.add_argument('--in_dir', help='Directory to process')
parser.add_argument('--render_dir', help='Directory to output render to')
parser.add_argument('--mov', action='store_true', help='whether to create an mp4')
parser.add_argument('--out_fn', help='mp4 basename')
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 make_movie(merge_dir, out_fn):
subprocess.call([
'ffmpeg',
'-hide_banner',
'-i', os.path.join(merge_dir, 'frame_%05d.png'),
'-y',
'-c:v', 'libx264',
'-preset', 'slow',
'-crf', '19',
'-vf', 'fps=25',
'-pix_fmt', 'yuv420p',
os.path.join(args.render_dir, out_fn + '.mp4')
])
def merge_files(file, crop_dir_list, ww, hh, overlap, merge_dir):
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, ix, iy = crop_dir_tuple
crop_dir = "crop_{}_{}_{}_{}".format(x, y, w, h)
image = Image.open("{}/{}/{}".format(args.in_dir, crop_dir, fn))
crop = image.crop((overlap, overlap, image.size[0] - overlap, image.size[1] - overlap))
canvas.paste(crop, ((x - (ix * args.overlap)) * args.scale, (y - (iy * args.overlap)) * args.scale,))
canvas.save("{}/{}".format(merge_dir, fn))
def merge_dir():
overlap = args.overlap * args.scale
crop_dirs = glob.glob("{}/crop_*".format(args.in_dir))
merge_dir = "{}/merged".format(args.in_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 - args.overlap * 2, hh)
if y == 0:
ww = max(x + w - args.overlap * 2, ww)
ix = 0 if x == 0 else 1 # floor(x / args.step)
iy = floor(y / args.step)
crop_dir_list.append((x, y, w, h, ix, iy))
crop_dir_list = sorted(crop_dir_list, key=operator.itemgetter(0, 1))
ww *= args.scale
hh *= args.scale
ww -= overlap
hh -= overlap
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, overlap, merge_dir,))
chunksize = 3
with Pool(processes=cpu_count()) as pool:
pool.starmap(merge_files, dataset, chunksize)
if args.mov:
make_movie(merge_dir, args.out_fn)
if __name__ == "__main__":
merge_dir()
|
