#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#
# Copyright © 2024 StrayWarrior <i@straywarrior.com>


import os
import face_recognition
from PIL import Image

def load_known_faces(known_faces_dir):
    known_face_encodings = []
    known_face_names = []

    for filename in os.listdir(known_faces_dir):
        if filename.endswith(".jpg") or filename.endswith(".png"):
            image_path = os.path.join(known_faces_dir, filename)
            image = face_recognition.load_image_file(image_path)
            face_encoding = face_recognition.face_encodings(image)[0]
            known_face_encodings.append(face_encoding)
            known_face_names.append(os.path.splitext(filename)[0])

    return known_face_encodings, known_face_names

def find_faces_in_images(image_dir, known_face_encodings, known_face_names):
    results = {}

    for filename in os.listdir(image_dir):
        if filename.endswith(".jpg") or filename.endswith(".png"):
            image_path = os.path.join(image_dir, filename)
            image = face_recognition.load_image_file(image_path)
            face_locations = face_recognition.face_locations(image)
            face_encodings = face_recognition.face_encodings(image, face_locations)

            found_faces = []
            for face_encoding in face_encodings:
                matches = face_recognition.compare_faces(known_face_encodings, face_encoding)
                name = "Unknown"

                if True in matches:
                    first_match_index = matches.index(True)
                    name = known_face_names[first_match_index]

                found_faces.append(name)

            if found_faces:
                results[filename] = found_faces

    return results

def main():
    known_faces_dir = "known_faces"  # Directory with images of known faces
    image_dir = "images_to_check"    # Directory with images to check

    known_face_encodings, known_face_names = load_known_faces(known_faces_dir)
    results = find_faces_in_images(image_dir, known_face_encodings, known_face_names)

    for image_name, found_faces in results.items():
        print(f"In {image_name}, found faces: {', '.join(found_faces)}")

if __name__ == "__main__":
    main()