-
Notifications
You must be signed in to change notification settings - Fork 0
/
preprocess.py
95 lines (70 loc) · 3.16 KB
/
preprocess.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
import argparse
import glob
import logging
import multiprocessing as mp
import os
import time
import cv2
from align_dlib import AlignDlib
logger = logging.getLogger(__name__)
align_dlib = AlignDlib(os.path.join(os.path.dirname(__file__), 'shape_predictor_68_face_landmarks.dat'))
def main(input_dir, output_dir, crop_dim):
start_time = time.time()
pool = mp.Pool(processes=mp.cpu_count())
if not os.path.exists(output_dir):
os.makedirs(output_dir)
# for image_dir in os.listdir(input_dir):
# image_output_dir = os.path.join(output_dir, os.path.basename(os.path.basename(image_dir)))
# if not os.path.exists(image_output_dir):
# os.makedirs(image_output_dir)
image_paths = glob.glob(os.path.join(input_dir, '*.jpg'))
for index, image_path in enumerate(image_paths):
# image_output_dir = os.path.join(output_dir, os.path.basename(os.path.dirname(image_path)))
# output_path = os.path.join(image_output_dir, os.path.basename(image_path))
pool.apply_async(preprocess_image, (image_path, output_dir, crop_dim))
pool.close()
pool.join()
logger.info('Completed in {} seconds'.format(time.time() - start_time))
def preprocess_image(input_path, output_path, crop_dim):
"""
Detect face, align and crop :param input_path. Write output to :param output_path
:param input_path: Path to input image
:param output_path: Path to write processed image
:param crop_dim: dimensions to crop image to
"""
image = _process_image(input_path, crop_dim)
output_file_path = output_path.rstrip("/") + os.sep+os.path.basename(input_path)
if image is not None:
logger.debug('Writing processed file: {}'.format(output_file_path))
cv2.imwrite(output_file_path, image)
else:
logger.warning("Skipping filename: {}".format(input_path))
def _process_image(filename, crop_dim):
image = None
aligned_image = None
image = _buffer_image(filename)
if image is not None:
aligned_image = _align_image(image, crop_dim)
else:
raise IOError('Error buffering image: {}'.format(filename))
return aligned_image
def _buffer_image(filename):
logger.debug('Reading image: {}'.format(filename))
image = cv2.imread(filename, )
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
return image
def _align_image(image, crop_dim):
bb = align_dlib.getLargestFaceBoundingBox(image)
aligned = align_dlib.align(crop_dim, image, bb, landmarkIndices=AlignDlib.INNER_EYES_AND_BOTTOM_LIP)
if aligned is not None:
aligned = cv2.cvtColor(aligned, cv2.COLOR_BGR2RGB)
return aligned
if __name__ == '__main__':
logging.basicConfig(level=logging.INFO)
parser = argparse.ArgumentParser(add_help=True)
parser.add_argument('--input-dir', type=str, action='store', default='data', dest='input_dir')
parser.add_argument('--output-dir', type=str, action='store', default='output', dest='output_dir')
parser.add_argument('--crop-dim', type=int, action='store', default=180, dest='crop_dim',
help='Size to crop images to')
args = parser.parse_args()
main(args.input_dir, args.output_dir, args.crop_dim)