-
Notifications
You must be signed in to change notification settings - Fork 0
/
process_images.py
78 lines (64 loc) · 2.46 KB
/
process_images.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
import os
from PIL import Image
import PIL
import glob
from pathlib import Path
from tqdm import tqdm
import zipfile
## Parameters
BOX = (1,10,23,34) # cropping box parameters
BACKGROUND = (255,255,255) # white border
OLD_FOLDER = 'kanjis'
NEW_FOLDER = 'kanjis_edited'
def zipdir(path, ziph):
# ziph is zipfile handle
for root, dirs, files in os.walk(path):
for file in files:
ziph.write(os.path.join(root, file))
def add_margin(pil_img, top, right, bottom, left, color):
width, height = pil_img.size
new_width = width + right + left
new_height = height + top + bottom
result = Image.new(pil_img.mode, (new_width, new_height), color)
result.paste(pil_img, (left, top))
return result
def expand2square(pil_img, background_color):
width, height = pil_img.size
if width == height:
return pil_img
elif width > height:
result = Image.new(pil_img.mode, (width, width), background_color)
result.paste(pil_img, (0, (width - height) // 2))
return result
else:
result = Image.new(pil_img.mode, (height, height), background_color)
result.paste(pil_img, ((height - width) // 2, 0))
return result
if __name__ == '__main__':
print(f"Converting files, cropping, removing alpha channel and pasting to {NEW_FOLDER}")
Path(NEW_FOLDER).mkdir(parents=True, exist_ok=True)
print("Clearing kanjis folder...")
for f in tqdm(glob.glob(os.path.join(NEW_FOLDER, '*.jpg'))):
os.remove(f)
file_list = glob.glob(os.path.join(OLD_FOLDER, '*.jpg'))
print(len(file_list))
i = 0
for old_path in tqdm(file_list):
file_name = Path(old_path).parts[-1]
try:
new_image = Image.\
open(old_path).\
crop(box=BOX).\
convert('RGB')
new_image = add_margin(
expand2square(new_image, BACKGROUND),
2,2,2,2,BACKGROUND)
assert new_image.size == (28,28), "Image has wrong dimensions"
new_image.save(os.path.join(NEW_FOLDER, file_name))
except PIL.UnidentifiedImageError:
print(f"Could not read {file_name}")
i += 1
print(f"Out of {len(file_list)} files, could not read {i}")
print("Zipping and saving to processed_images.zip")
with zipfile.ZipFile('processed_images.zip', 'w', zipfile.ZIP_DEFLATED) as zipf:
zipdir(os.path.join('.', NEW_FOLDER), zipf)