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
|
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=192, type=int, help='Width of frames')
parser.add_argument('--overlap', default=2, type=int, help='Expand all frames by this margin')
parser.add_argument('--ratio', default=2.0, type=float, help='Pre-crop frames to this ratio')
parser.add_argument('--in_dir', help='Directory to process')
parser.add_argument('--out_dir', default='inputs', help='Directory to output crops to')
parser.add_argument('--start', default=0, type=int, help='Minimum frame number')
parser.add_argument('--end', default=99999, type=int, help='Maximum frame number')
args = parser.parse_args()
def crop_dir():
step = args.step
overlap = args.overlap
side = step + overlap * 2
files = glob.glob('{}/*.png'.format(args.in_dir))
# print(files)
img = Image.open(files[0])
width, old_height = img.size
if args.ratio == (width / old_height):
height = old_height
top_offset = 0
else:
height = round(width / args.ratio)
top_offset = round((old_height - height) / 2)
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, x, y, w, h)
if os.path.exists(crop_dir):
rmtree(crop_dir)
os.makedirs(crop_dir)
dataset = []
for raw_fn in files:
fn, ext = os.path.splitext(os.path.basename(raw_fn))
partz = fn.split('_')
try:
num = int(partz[-1].replace('.png', ''))
if num >= opt.start and num <= opt.end:
dataset.append((raw_fn, width, height, step, side, overlap, top_offset,))
except:
pass
chunksize = 3
with Pool(processes=cpu_count()) as pool:
pool.starmap(crop_image, dataset, chunksize)
def crop_image(raw_fn, width, height, step, side, overlap, top_offset):
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, x, y, w, h)
# print("{}x{} {}x{}".format(x, y, w, h))
crop = img.crop((xx, yy + top_offset, xx + w, yy + h + top_offset))
crop.save("{}/{}".format(crop_dir, fn))
if __name__ == "__main__":
crop_dir()
|
