1# -*- coding: utf-8 -*- 2from __future__ import print_function 3import click 4import os 5import re 6import face_recognition.api as face_recognition 7import multiprocessing 8import sys 9import itertools 10 11 12def print_result(filename, location): 13 top, right, bottom, left = location 14 print("{},{},{},{},{}".format(filename, top, right, bottom, left)) 15 16 17def test_image(image_to_check, model): 18 unknown_image = face_recognition.load_image_file(image_to_check) 19 face_locations = face_recognition.face_locations(unknown_image, number_of_times_to_upsample=0, model=model) 20 21 for face_location in face_locations: 22 print_result(image_to_check, face_location) 23 24 25def image_files_in_folder(folder): 26 return [os.path.join(folder, f) for f in os.listdir(folder) if re.match(r'.*\.(jpg|jpeg|png)', f, flags=re.I)] 27 28 29def process_images_in_process_pool(images_to_check, number_of_cpus, model): 30 if number_of_cpus == -1: 31 processes = None 32 else: 33 processes = number_of_cpus 34 35 # macOS will crash due to a bug in libdispatch if you don't use 'forkserver' 36 context = multiprocessing 37 if "forkserver" in multiprocessing.get_all_start_methods(): 38 context = multiprocessing.get_context("forkserver") 39 40 pool = context.Pool(processes=processes) 41 42 function_parameters = zip( 43 images_to_check, 44 itertools.repeat(model), 45 ) 46 47 pool.starmap(test_image, function_parameters) 48 49 50@click.command() 51@click.argument('image_to_check') 52@click.option('--cpus', default=1, help='number of CPU cores to use in parallel. -1 means "use all in system"') 53@click.option('--model', default="hog", help='Which face detection model to use. Options are "hog" or "cnn".') 54def main(image_to_check, cpus, model): 55 # Multi-core processing only supported on Python 3.4 or greater 56 if (sys.version_info < (3, 4)) and cpus != 1: 57 click.echo("WARNING: Multi-processing support requires Python 3.4 or greater. Falling back to single-threaded processing!") 58 cpus = 1 59 60 if os.path.isdir(image_to_check): 61 if cpus == 1: 62 [test_image(image_file, model) for image_file in image_files_in_folder(image_to_check)] 63 else: 64 process_images_in_process_pool(image_files_in_folder(image_to_check), cpus, model) 65 else: 66 test_image(image_to_check, model) 67 68 69if __name__ == "__main__": 70 main() 71