Run Face Search Locally

run_face_search.py

import argparse
import sys
import numpy as np
from pathlib import Path
from PIL import Image

from insightface.app import FaceAnalysis


def run_face_search(query_face_image: Path, target_search_dir: Path,
                    similarity_threshold=0.8, min_resolution=50):
    """
    Run face search on target images and return matches that are above the similarity threshold.
    :param query_face_image: path to the query face image
    :param target_search_dir: directory with target images to search from
    :param similarity_threshold: threshold for face similarity
    :param min_resolution: minimum resolution for a face to be considered
    """
    # init insightface model
    face_model = FaceAnalysis(allowed_modules=['detection', 'recognition'])
    face_model.prepare(ctx_id=0)  # ctx_id=-1 to use CPU

    # get embedding for query face
    query_image = Image.open(query_face_image).convert("RGB")
    query_faces = face_model.get(np.array(query_image))
    if len(query_faces) == 0:
        raise Exception("No face detected in the query image")
    elif len(query_faces) > 1:
        print("Multiple faces detected in the query image. Picking the first one.")
    query_emb = query_faces[0].embedding

    # run face similarity on target images
    # currently computing similarity individually for each pair. We will optimize this process in future iterations
    matches = []
    for extension in ['*.jpg', '*.png']:
        for img_path in target_search_dir.glob(extension):
            faces = face_model.get(np.array(Image.open(img_path).convert("RGB")))  # get all faces from current image
            for face in faces:
                # ignore when resolution low
                if any(map(lambda x: x < min_resolution, [face.bbox[2] - face.bbox[0], face.bbox[3] - face.bbox[1]])):
                    continue
                # otherwise compute similarity and add match if above threshold
                similarity = cosine_distance(query_emb, face.embedding)
                if similarity > similarity_threshold:
                    matches.append((img_path, similarity, face.bbox))
    return matches


def cosine_distance(v1, v2):
    """Returns cosine distance between the two given vectors"""
    return np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2))


def main(_=None):
    parser = argparse.ArgumentParser(description='Local Face Search')
    parser.add_argument('-q', '--query', required=True, help="path to the image containing the query face")
    parser.add_argument('-t', '--target', required=True, help="path to the directory with target images to search from")
    parser.add_argument('--sim-threshold', default=0.5, help="similarity threshold for face search")
    parser.add_argument('--min-res', default=20, help="minimum resolution for a face to be considered")

    args = parser.parse_args()

    print('Running face search...')
    matches = run_face_search(Path(args.query), Path(args.target), float(args.sim_threshold), int(args.min_res))
    print('Face search completed. Matches found:')
    print(matches)


if __name__ == "__main__":
    main(sys.argv[1:])