parent
25d1bea718
commit
d3c51749b0
@ -0,0 +1,72 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
from __future__ import print_function
|
||||
import click
|
||||
import os
|
||||
import re
|
||||
import face_recognition.api as face_recognition
|
||||
import multiprocessing
|
||||
import sys
|
||||
import itertools
|
||||
|
||||
|
||||
def print_result(filename, location):
|
||||
top, right, bottom, left = location
|
||||
print("{},{},{},{},{}".format(filename, top, right, bottom, left))
|
||||
|
||||
|
||||
def test_image(image_to_check, model, upsample):
|
||||
unknown_image = face_recognition.load_image_file(image_to_check)
|
||||
face_locations = face_recognition.face_locations(unknown_image, number_of_times_to_upsample=upsample, model=model)
|
||||
|
||||
for face_location in face_locations:
|
||||
print_result(image_to_check, face_location)
|
||||
|
||||
|
||||
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, number_of_cpus, model, upsample):
|
||||
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(model),
|
||||
itertools.repeat(upsample),
|
||||
)
|
||||
|
||||
pool.starmap(test_image, function_parameters)
|
||||
|
||||
|
||||
@click.command()
|
||||
@click.argument('image_to_check')
|
||||
@click.option('--cpus', default=1, help='number of CPU cores to use in parallel. -1 means "use all in system"')
|
||||
@click.option('--model', default="hog", help='Which face detection model to use. Options are "hog" or "cnn".')
|
||||
@click.option('--upsample', default=0, help='How many times to upsample the image looking for faces. Higher numbers find smaller faces.')
|
||||
def main(image_to_check, cpus, model, upsample):
|
||||
# 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):
|
||||
if cpus == 1:
|
||||
[test_image(image_file, model, upsample) for image_file in image_files_in_folder(image_to_check)]
|
||||
else:
|
||||
process_images_in_process_pool(image_files_in_folder(image_to_check), cpus, model, upsample)
|
||||
else:
|
||||
test_image(image_to_check, model, upsample)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
Loading…
Reference in new issue