Face recognition script using face_recognition Python library

faces.py

# Requirements:
# sudo apt update && sudo apt install -y cmake
# pip install face-recognition numpy opencv-python Pillow tqdm
import cv2
import face_recognition
import numpy as np
from PIL import Image
from tqdm import tqdm


def detect_faces(image):
    """
    Detecta rostos em uma imagem usando face_recognition.

    :param image: Imagem PIL
    :return: Lista de coordenadas de rostos (top, right, bottom, left)
    """
    img_array = np.array(image)
    face_locations = face_recognition.face_locations(img_array, model="cnn")
    return face_locations


def normalize_face(image, face_location):
    """
    Normaliza um rosto detectado na imagem.

    :param image: Imagem PIL
    :param face_location: Coordenadas do rosto (top, right, bottom, left)
    :return: Imagem PIL normalizada do rosto
    """
    top, right, bottom, left = face_location
    face_image = image.crop((left, top, right, bottom))
    face_image = face_image.resize((160, 160))
    face_array = np.array(face_image)

    # Equalização de histograma adaptativa
    lab = cv2.cvtColor(face_array, cv2.COLOR_RGB2LAB)
    l, a, b = cv2.split(lab)
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
    cl = clahe.apply(l)
    limg = cv2.merge((cl,a,b))
    face_array = cv2.cvtColor(limg, cv2.COLOR_LAB2RGB)

    # Normalização de intensidade
    face_array = cv2.normalize(face_array, None, 0, 255, cv2.NORM_MINMAX)

    return Image.fromarray(face_array)


def generate_embedding(face_image):
    """
    Gera um embedding para um rosto normalizado.

    :param face_image: Imagem PIL normalizada do rosto
    :return: Array numpy com o embedding ou None se nenhum rosto for detectado
    """
    face_array = np.array(face_image)
    encodings = face_recognition.face_encodings(face_array)
    return encodings[0] if encodings else None


def find_matching_faces(reference_embedding, embeddings_list, threshold=0.6):
    """
    Encontra rostos correspondentes com base em uma lista de embeddings.

    :param reference_embedding: Embedding de referência
    :param embeddings_list: Lista de embeddings para comparação
    :param threshold: Limiar de distância para considerar uma correspondência
    :return: Lista de índices dos rostos correspondentes
    """
    if reference_embedding is None or not embeddings_list:
        return []
    valid_embeddings = [emb for emb in embeddings_list if emb is not None]
    if not valid_embeddings:
        return []
    distances = face_recognition.face_distance(valid_embeddings, reference_embedding)
    return [i for i, distance in enumerate(distances) if distance <= threshold]


def create_average_embedding(image_files):
    """
    Cria um embedding médio a partir de várias imagens de uma pessoa.
    Pode melhorar o embedding caso possua mais de uma imagem de referência da mesma pessoa.

    :param image_files: Lista de caminhos para as imagens
    :return: Embedding médio
    """
    embeddings = []
    for file in image_files:
        image = Image.open(file).convert('RGB')
        face_locations = detect_faces(image)
        if len(face_locations) > 0:
            normalized_face = normalize_face(image, face_locations[0])
            embedding = generate_embedding(normalized_face)
            if embedding is not None:
                embeddings.append(embedding)
    if embeddings:
        return np.mean(embeddings, axis=0)
    else:
        return None


if __name__ == "__main__":
    import argparse
    import sys
    from pathlib import Path

    parser = argparse.ArgumentParser()
    parser.add_argument("photo_dir", type=Path, help="Pasta contendo todo o conjunto de fotos a serem analisadas")
    parser.add_argument("reference_picture", type=Path, help="Foto com a face de referência, que será buscada nas outras fotos")
    args = parser.parse_args()

    photo_dir = args.photo_dir
    reference_filename = args.reference_picture
    ref_image = Image.open(reference_filename).convert('RGB')
    ref_face_locations = detect_faces(ref_image)

    if not ref_face_locations:
        print("ERRO: não foi possível encontrar faces na imagem de referência {reference_filename}", file=sys.stderr)
        exit(1)
    elif len(ref_face_locations) > 1:
        print(f"Encontradas {len(ref_face_locations)} faces na imagem de referência. Usando a primeira delas.")

    ref_normalized_face = normalize_face(ref_image, ref_face_locations[0])
    ref_embedding = generate_embedding(ref_normalized_face)
    if ref_embedding is None:
        print("ERRO: não foi possível gerar o embedding para a imagem de referência.", file=sys.stderr)
        exit(2)
    print(f"Embedding da face de referência:\n{ref_embedding}\n")

    other_embeddings = []
    other_files = []
    for file_path in tqdm(photo_dir.glob("foto-*.jpg"), desc="Processando fotos"):
        image = Image.open(file_path).convert('RGB')
        face_locations = detect_faces(image)
        if face_locations:
            # TODO: precisamos fazer testes para verificar se normalização é realmente necessária (salvar as faces
            # originais e as versões normalizadas para checagem manual e também verificar acurácia com/sem normalização
            # para um grande dataset)
            normalized_face = normalize_face(image, face_locations[0])
            embedding = generate_embedding(normalized_face)
            if embedding is not None:
                other_embeddings.append(embedding)
                other_files.append(file_path)
    matches = find_matching_faces(ref_embedding, other_embeddings)
    print(f"Encontradas {len(matches)} correspondências:")
    for idx in matches:
        print(f"- {other_files[idx]}")