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