All support for parallel processing in the CLI

ageitgey · ageitgey · commit 245160c2f6c7 · 2017-06-26T15:15:53.000-07:00
diff --git a/README.md b/README.md
@@ -140,6 +140,18 @@ Barack Obama
 unknown_person
 ```
 
+Face recognition can also be done in parallel if you have a computer with
+multiple CPU cores. For example if your system has 4 CPU cores, you can
+process about 4 times as many images in the same amount of time by using
+all your CPU cores in parallel.
+
+If you are using Python 3.4 or newer, pass in a `--cpus <number_of_cpu_cores_to_use>` parameter:
+
+```bash
+$ face_recognition -cpus 4 ./pictures_of_people_i_know/ ./unknown_pictures/
+```
+
+You can also pass in `--cpus -1` to use all CPU cores in your system.
 
 #### Python Module
 
diff --git a/face_recognition/cli.py b/face_recognition/cli.py
@@ -6,7 +6,9 @@
 import scipy.misc
 import warnings
 import face_recognition.api as face_recognition
-
+import multiprocessing
+import itertools
+import sys
 
 def scan_known_people(known_people_folder):
     known_names = []
@@ -54,14 +56,40 @@ def image_files_in_folder(folder):
     return [os.path.join(folder, f) for f in os.listdir(folder) if re.match(r'.*\.(jpg|jpeg|png)', f, flags=re.I)]
 
 
+def process_images_in_process_pool(images_to_check, known_names, known_face_encodings, number_of_cpus):
+    if number_of_cpus == -1:
+        processes = None
+    else:
+        processes = number_of_cpus
+
+    # macOS will crash due to a bug in libdispatch if you don't use 'forkserver'
+    context = multiprocessing
+    if "forkserver" in multiprocessing.get_all_start_methods():
+        context = multiprocessing.get_context("forkserver")
+
+    pool = context.Pool(processes=processes)
+    function_parameters = zip(images_to_check, itertools.repeat(known_names), itertools.repeat(known_face_encodings))
+
+    pool.starmap(test_image, function_parameters)
+
+
 @click.command()
 @click.argument('known_people_folder')
 @click.argument('image_to_check')
-def main(known_people_folder, image_to_check):
+@click.option('--cpus', default=1, help='number of CPU cores to use in parallel (can speed up processing lots of images). -1 means "use all in system"')
+def main(known_people_folder, image_to_check, cpus):
     known_names, known_face_encodings = scan_known_people(known_people_folder)
 
+    # Multi-core processing only supported on Python 3.4 or greater
+    if (sys.version_info < (3, 4)) and cpus != 1:
+        click.echo("WARNING: Multi-processing support requires Python 3.4 or greater. Falling back to single-threaded processing!")
+        cpus = 1
+
     if os.path.isdir(image_to_check):
-        [test_image(image_file, known_names, known_face_encodings) for image_file in image_files_in_folder(image_to_check)]
+        if cpus == 1:
+            [test_image(image_file, known_names, known_face_encodings) for image_file in image_files_in_folder(image_to_check)]
+        else:
+            process_images_in_process_pool(image_files_in_folder(image_to_check), known_names, known_face_encodings, cpus)
     else:
         test_image(image_to_check, known_names, known_face_encodings)
 
