-
Notifications
You must be signed in to change notification settings - Fork 0
/
face_gen.py
104 lines (97 loc) · 4.2 KB
/
face_gen.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
import manga109api
from pprint import pprint
import cv2
import time
import os
from dotenv import load_dotenv
import numpy as np
load_dotenv()
TEST_RATIO = 0.1
manga109_root_dir = os.getenv('MANGA109_ROOT_DIR')
p = manga109api.Parser(root_dir=manga109_root_dir)
book_cnt = 0
os.system('rm -r face body')
os.mkdir('face/')
os.mkdir('body/')
face_train = open("face_training_set.txt", "w")
face_test = open("face_test_set.txt", "w")
for book in p.books:
book_cnt += 1
face_cnt = 0
body_cnt = 0
os.mkdir('face/' + str(book_cnt))
os.mkdir('body/' + str(book_cnt))
for page in p.annotations[book]["book"]["pages"]["page"]:
img = cv2.imread(p.img_path(book=book, index=page["@index"]))
if "face" in page.keys() and type(page["face"]) is list:
for face in page["face"]:
#pprint(face)
x, y = face["@xmin"], face["@ymin"]
w = face["@xmax"] - x
h = face["@ymax"] - y
if w > h:
offset = int((w - h) / 2)
face_img = img[y:y + h, x + offset:x + h + offset]
else:
offset = int((h - w) / 2)
face_img = img[y + offset:y + w + offset, x:x + w]
if face_img.shape[0] >= 50:
ok_img = cv2.resize(
face_img, (128, 128), interpolation=cv2.INTER_LANCZOS4)
face_cnt += 1
cv2.imwrite(
'face/' + str(book_cnt) + '/' + str(face_cnt) + '.jpg',
ok_img)
if np.random.rand() < TEST_RATIO:
face_test.write('face/' + str(book_cnt) + '/' + str(face_cnt) + '.jpg ' + str(book_cnt) + '\n')
else:
face_train.write('face/' + str(book_cnt) + '/' + str(face_cnt) + '.jpg ' + str(book_cnt) + '\n')
if "face" in page.keys() and type(page["face"]) is dict:
face = page["face"]
x, y = face["@xmin"], face["@ymin"]
w = face["@xmax"] - x
h = face["@ymax"] - y
if w > h:
offset = int((w - h) / 2)
face_img = img[y:y + h, x + offset:x + h + offset]
else:
offset = int((h - w) / 2)
face_img = img[y + offset:y + w + offset, x:x + w]
if face_img.shape[0] >= 50:
ok_img = cv2.resize(
face_img, (128, 128), interpolation=cv2.INTER_LANCZOS4)
face_cnt += 1
cv2.imwrite(
'face/' + str(book_cnt) + '/' + str(face_cnt) + '.jpg',
ok_img)
if np.random.rand() < TEST_RATIO:
face_test.write('face/' + str(book_cnt) + '/' + str(face_cnt) + '.jpg ' + str(book_cnt) + '\n')
else:
face_train.write('face/' + str(book_cnt) + '/' + str(face_cnt) + '.jpg ' + str(book_cnt) + '\n')
if "body" in page.keys() and type(page["body"]) is list:
for body in page["body"]:
#pprint(body)
x, y = body["@xmin"], body["@ymin"]
w = body["@xmax"] - x
h = body["@ymax"] - y
body_img = img[y:y + h, x:x + w]
if body_img.shape[0] >= 50 and body_img.shape[1] >= 50:
ok_img = cv2.resize(
body_img, (128, 128), interpolation=cv2.INTER_LANCZOS4)
body_cnt += 1
cv2.imwrite(
'body/' + str(book_cnt) + '/' + str(body_cnt) + '.jpg',
ok_img)
if "body" in page.keys() and type(page["body"]) is dict:
body = page["body"]
x, y = body["@xmin"], body["@ymin"]
w = body["@xmax"] - x
h = body["@ymax"] - y
body_img = img[y:y + h, x:x + w]
if body_img.shape[0] >= 50 and body_img.shape[1] >= 50:
ok_img = cv2.resize(
body_img, (128, 128), interpolation=cv2.INTER_LANCZOS4)
body_cnt += 1
cv2.imwrite(
'body/' + str(book_cnt) + '/' + str(body_cnt) + '.jpg',
ok_img)