utils/inference.py

from utils.database import get_database
import os
import face_recognition
import cv2
import numpy as np
from PIL import Image
import streamlit as st

PKL_PATH = 'dataset/database.pkl'

def recognize(image,tolerance):
    database = get_database(PKL_PATH)
    known_encoding = [database[id]['encoding'] for id in database.keys()]
    name = 'Unknown'
    face_id = 'Unknown'
    face_locations = face_recognition.face_locations(image)
    face_encodings = face_recognition.face_encodings(image,face_locations)
    for (top,right,bottom,left), face_encoding in zip(face_locations,face_encodings):
        matches = face_recognition.compare_faces(known_encoding,face_encoding,tolerance=tolerance)
        distances = face_recognition.face_distance(known_encoding,face_encoding)
        name = 'Unknown'
        face_id = 'Unknown'
        distance = 100000
        for i in range(len(matches)):
            if matches[i].all():
                match_index = i
                temp_dist = round(np.sum(distances[match_index]),2)
                if temp_dist < distance:
                    distance = temp_dist
                    name = database[match_index]['name']
                    face_id = database[match_index]['face_id'].split("_")[1]
                    # cv2.putText(image,str(distance),(left,top-30),cv2.FONT_HERSHEY_SIMPLEX,0.75,(0,255,0),2)
        cv2.rectangle(image,(left,top),(right,bottom),(0,255,0),2)
        cv2.putText(image,name,(left,top-10),cv2.FONT_HERSHEY_SIMPLEX,0.75,(0,255,0),2)
    return image, name, face_id