import cv2
import numpy as np
import torch
from PIL import Image
from facenet_pytorch import MTCNN, InceptionResnetV1
import gradio as gr
import os

# Load pre-trained model for face detection and face recognition
mtcnn = MTCNN(image_size=160, margin=0, min_face_size=20)
resnet = InceptionResnetV1(pretrained='vggface2').eval()

reference_image_embedding = None

# Function to preprocess image and extract face embeddings
def preprocess_image(image):
    image = Image.fromarray(image)
    image_cropped = mtcnn(image)
    if image_cropped is not None:
        image_embedding = resnet(image_cropped.unsqueeze(0))
        return image_embedding
    return None

def set_reference_image(image):
    global reference_image_embedding
    if image is not None:
        reference_image_embedding = preprocess_image(image)
        if reference_image_embedding is not None:
            return "Reference image set successfully."
        else:
            return "No face detected in the reference image."
    else:
        return "Failed to set reference image."

def capture_and_compare(captured_image):
    if reference_image_embedding is None:
        return "Please set the reference image first."

    captured_image_embedding = preprocess_image(captured_image)
    if captured_image_embedding is None:
        return "No face detected in the captured image."

    similarity = torch.nn.functional.cosine_similarity(captured_image_embedding, reference_image_embedding)
    if similarity.item() > 0.6:  # Adjust the threshold as necessary
        return "Faces match! Access granted."
    else:
        return "Faces do not match. Access denied."

# Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Face Recognition Lock")

    with gr.Row():
        with gr.Column():
            reference_image_input = gr.Image(label="Upload Reference Image")
            reference_image_button = gr.Button("Set Reference Image")
            reference_image_status = gr.Textbox(label="Status")

        with gr.Column():
            captured_image_input = gr.Image(label="Capture or Upload Image to Compare")
            compare_faces_button = gr.Button("Compare Faces")
            compare_faces_status = gr.Textbox(label="Result")

    reference_image_button.click(set_reference_image, inputs=[reference_image_input], outputs=[reference_image_status])
    compare_faces_button.click(capture_and_compare, inputs=[captured_image_input], outputs=[compare_faces_status])

# Launch the Gradio interface
demo.launch()