Skip to content

Commit e49528b

Browse files
Add files via upload
1 parent 1f5acf0 commit e49528b

11 files changed

+2888
-0
lines changed

align_dataset_mtcnn.py

Lines changed: 159 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,159 @@
1+
"""Performs face alignment and stores face thumbnails in the output directory."""
2+
# MIT License
3+
#
4+
# Copyright (c) 2016 David Sandberg
5+
#
6+
# Permission is hereby granted, free of charge, to any person obtaining a copy
7+
# of this software and associated documentation files (the "Software"), to deal
8+
# in the Software without restriction, including without limitation the rights
9+
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
10+
# copies of the Software, and to permit persons to whom the Software is
11+
# furnished to do so, subject to the following conditions:
12+
#
13+
# The above copyright notice and this permission notice shall be included in all
14+
# copies or substantial portions of the Software.
15+
#
16+
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
17+
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
18+
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
19+
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
20+
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
21+
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
22+
# SOFTWARE.
23+
24+
from __future__ import absolute_import
25+
from __future__ import division
26+
from __future__ import print_function
27+
28+
from scipy import misc
29+
import sys
30+
import os
31+
import argparse
32+
import tensorflow as tf
33+
import numpy as np
34+
import facenet
35+
import align.detect_face
36+
import random
37+
from time import sleep
38+
39+
def main(args):
40+
sleep(random.random())
41+
output_dir = os.path.expanduser(args.output_dir)
42+
if not os.path.exists(output_dir):
43+
os.makedirs(output_dir)
44+
# Store some git revision info in a text file in the log directory
45+
src_path,_ = os.path.split(os.path.realpath(__file__))
46+
facenet.store_revision_info(src_path, output_dir, ' '.join(sys.argv))
47+
dataset = facenet.get_dataset(args.input_dir)
48+
49+
print('Creating networks and loading parameters')
50+
51+
with tf.Graph().as_default():
52+
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=args.gpu_memory_fraction)
53+
sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options, log_device_placement=False))
54+
with sess.as_default():
55+
pnet, rnet, onet = align.detect_face.create_mtcnn(sess, None)
56+
57+
minsize = 20 # minimum size of face
58+
threshold = [ 0.6, 0.7, 0.7 ] # three steps's threshold
59+
factor = 0.709 # scale factor
60+
61+
# Add a random key to the filename to allow alignment using multiple processes
62+
random_key = np.random.randint(0, high=99999)
63+
bounding_boxes_filename = os.path.join(output_dir, 'bounding_boxes_%05d.txt' % random_key)
64+
65+
with open(bounding_boxes_filename, "w") as text_file:
66+
nrof_images_total = 0
67+
nrof_successfully_aligned = 0
68+
if args.random_order:
69+
random.shuffle(dataset)
70+
for cls in dataset:
71+
output_class_dir = os.path.join(output_dir, cls.name)
72+
if not os.path.exists(output_class_dir):
73+
os.makedirs(output_class_dir)
74+
if args.random_order:
75+
random.shuffle(cls.image_paths)
76+
for image_path in cls.image_paths:
77+
nrof_images_total += 1
78+
filename = os.path.splitext(os.path.split(image_path)[1])[0]
79+
output_filename = os.path.join(output_class_dir, filename+'.png')
80+
print(image_path)
81+
if not os.path.exists(output_filename):
82+
try:
83+
img = misc.imread(image_path)
84+
except (IOError, ValueError, IndexError) as e:
85+
errorMessage = '{}: {}'.format(image_path, e)
86+
print(errorMessage)
87+
else:
88+
if img.ndim<2:
89+
print('Unable to align "%s"' % image_path)
90+
text_file.write('%s\n' % (output_filename))
91+
continue
92+
if img.ndim == 2:
93+
img = facenet.to_rgb(img)
94+
img = img[:,:,0:3]
95+
96+
bounding_boxes, _ = align.detect_face.detect_face(img, minsize, pnet, rnet, onet, threshold, factor)
97+
nrof_faces = bounding_boxes.shape[0]
98+
if nrof_faces>0:
99+
det = bounding_boxes[:,0:4]
100+
det_arr = []
101+
img_size = np.asarray(img.shape)[0:2]
102+
if nrof_faces>1:
103+
if args.detect_multiple_faces:
104+
for i in range(nrof_faces):
105+
det_arr.append(np.squeeze(det[i]))
106+
else:
107+
bounding_box_size = (det[:,2]-det[:,0])*(det[:,3]-det[:,1])
108+
img_center = img_size / 2
109+
offsets = np.vstack([ (det[:,0]+det[:,2])/2-img_center[1], (det[:,1]+det[:,3])/2-img_center[0] ])
110+
offset_dist_squared = np.sum(np.power(offsets,2.0),0)
111+
index = np.argmax(bounding_box_size-offset_dist_squared*2.0) # some extra weight on the centering
112+
det_arr.append(det[index,:])
113+
else:
114+
det_arr.append(np.squeeze(det))
115+
116+
for i, det in enumerate(det_arr):
117+
det = np.squeeze(det)
118+
bb = np.zeros(4, dtype=np.int32)
119+
bb[0] = np.maximum(det[0]-args.margin/2, 0)
120+
bb[1] = np.maximum(det[1]-args.margin/2, 0)
121+
bb[2] = np.minimum(det[2]+args.margin/2, img_size[1])
122+
bb[3] = np.minimum(det[3]+args.margin/2, img_size[0])
123+
cropped = img[bb[1]:bb[3],bb[0]:bb[2],:]
124+
scaled = misc.imresize(cropped, (args.image_size, args.image_size), interp='bilinear')
125+
nrof_successfully_aligned += 1
126+
filename_base, file_extension = os.path.splitext(output_filename)
127+
if args.detect_multiple_faces:
128+
output_filename_n = "{}_{}{}".format(filename_base, i, file_extension)
129+
else:
130+
output_filename_n = "{}{}".format(filename_base, file_extension)
131+
misc.imsave(output_filename_n, scaled)
132+
text_file.write('%s %d %d %d %d\n' % (output_filename_n, bb[0], bb[1], bb[2], bb[3]))
133+
else:
134+
print('Unable to align "%s"' % image_path)
135+
text_file.write('%s\n' % (output_filename))
136+
137+
print('Total number of images: %d' % nrof_images_total)
138+
print('Number of successfully aligned images: %d' % nrof_successfully_aligned)
139+
140+
141+
def parse_arguments(argv):
142+
parser = argparse.ArgumentParser()
143+
144+
parser.add_argument('input_dir', type=str, help='Directory with unaligned images.')
145+
parser.add_argument('output_dir', type=str, help='Directory with aligned face thumbnails.')
146+
parser.add_argument('--image_size', type=int,
147+
help='Image size (height, width) in pixels.', default=182)
148+
parser.add_argument('--margin', type=int,
149+
help='Margin for the crop around the bounding box (height, width) in pixels.', default=44)
150+
parser.add_argument('--random_order',
151+
help='Shuffles the order of images to enable alignment using multiple processes.', action='store_true')
152+
parser.add_argument('--gpu_memory_fraction', type=float,
153+
help='Upper bound on the amount of GPU memory that will be used by the process.', default=1.0)
154+
parser.add_argument('--detect_multiple_faces', type=bool,
155+
help='Detect and align multiple faces per image.', default=False)
156+
return parser.parse_args(argv)
157+
158+
if __name__ == '__main__':
159+
main(parse_arguments(sys.argv[1:]))

det1.npy

26.7 KB
Binary file not shown.

det2.npy

392 KB
Binary file not shown.

det3.npy

1.49 MB
Binary file not shown.

0 commit comments

Comments
 (0)