-
Notifications
You must be signed in to change notification settings - Fork 1
/
convert.py
144 lines (111 loc) · 4.35 KB
/
convert.py
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
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
"""Resize and crop images to square, save as tiff."""
import os
from multiprocessing.pool import Pool
import click
import numpy as np
from PIL import Image, ImageFilter
import data
N_PROC = 20
def convert(fname, crop_size):
img = Image.open(fname)
blurred = img.filter(ImageFilter.BLUR)
ba = np.array(blurred)
h, w, _ = ba.shape
if w > 1.2 * h:
left_max = ba[:, : w // 32, :].max(axis=(0, 1)).astype(int)
right_max = ba[:, - w // 32:, :].max(axis=(0, 1)).astype(int)
max_bg = np.maximum(left_max, right_max)
foreground = (ba > max_bg + 10).astype(np.uint8)
bbox = Image.fromarray(foreground).getbbox()
if bbox is None:
print('bbox none for {} (???)'.format(fname))
else:
left, upper, right, lower = bbox
# if we selected less than 80% of the original
# height, just crop the square
if right - left < 0.8 * h or lower - upper < 0.8 * h:
print('bbox too small for {}'.format(fname))
bbox = None
else:
bbox = None
if bbox is None:
bbox = square_bbox(img)
cropped = img.crop(bbox)
resized = cropped.resize([crop_size, crop_size])
return resized
def square_bbox(img):
w, h = img.size
left = max((w - h) // 2, 0)
upper = 0
right = min(w - (w - h) // 2, w)
lower = h
return (left, upper, right, lower)
def convert_square(fname, crop_size):
img = Image.open(fname)
bbox = square_bbox(img)
cropped = img.crop(bbox)
resized = cropped.resize([crop_size, crop_size])
return resized
def get_convert_fname(fname, extension, directory, convert_directory):
return fname.replace('jpeg', extension).replace(directory,
convert_directory)
def process(args):
fun, arg = args
directory, convert_directory, fname, crop_size, extension = arg
convert_fname = get_convert_fname(fname, extension, directory,
convert_directory)
if not os.path.exists(convert_fname):
img = fun(fname, crop_size)
save(img, convert_fname)
def save(img, fname):
img.save(fname, quality=97)
@click.command()
@click.option('--directory', default='data/train', show_default=True,
help="Directory with original images.")
@click.option('--convert_directory', default='data/train_res', show_default=True,
help="Where to save converted images.")
@click.option('--test', is_flag=True, default=False, show_default=True,
help="Convert images one by one and examine them on screen.")
@click.option('--crop_size', default=256, show_default=True,
help="Size of converted images.")
@click.option('--extension', default='jpg', show_default=True,
help="Filetype of converted images.")
def main(directory, convert_directory, test, crop_size, extension):
try:
os.mkdir(convert_directory)
except OSError:
pass
filenames = [os.path.join(dp, f) for dp, dn, fn in os.walk(directory)
for f in fn if f.endswith('jpeg') or f.endswith('jpg') or f.endswith('png') or f.endswith('tiff')]
filenames = sorted(filenames)
if test:
names = data.get_names(filenames)
y = data.get_labels(names)
for f, level in zip(filenames, y):
if level == 1:
try:
img = convert(f, crop_size)
img.show()
Image.open(f).show()
real_raw_input = vars(__builtins__).get('raw_input',input)
real_raw_input('enter for next')
except KeyboardInterrupt:
exit(0)
print("Resizing images in {} to {}, this takes a while."
"".format(directory, convert_directory))
n = len(filenames)
# process in batches, sometimes weird things happen with Pool on my machine
batchsize = 500
batches = n // batchsize + 1
pool = Pool(N_PROC)
args = []
for f in filenames:
args.append((convert, (directory, convert_directory, f, crop_size,
extension)))
for i in range(batches):
print("batch {:>2} / {}".format(i + 1, batches))
pool.map(process, args[i * batchsize: (i + 1) * batchsize])
pool.close()
print('done')
if __name__ == '__main__':
main()