|
| 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:])) |
0 commit comments