import gradio as gr
import tensorflow as tf
import numpy as np
from scipy.spatial.distance import cosine
import cv2
import os

# Load the embedding model
embedding_model = tf.keras.models.load_model('embedding_model.h5')

# Database to store embeddings and user IDs
user_embeddings = {}

# Preprocess the image
def preprocess_image(image):
    image = cv2.resize(image, (200, 200))  # Assuming your model expects 200x200 input
    image = tf.keras.applications.resnet50.preprocess_input(image)
    return np.expand_dims(image, axis=0)

# Generate embedding
def generate_embedding(image):
    preprocessed_image = preprocess_image(image)
    return embedding_model.predict(preprocessed_image)[0]

# Register new user
def register_user(image, user_id):
    embedding = generate_embedding(image)
    user_embeddings[user_id] = embedding
    return f"User {user_id} registered successfully."

# Recognize user
def recognize_user(image):
    new_embedding = generate_embedding(image)
    min_distance = float('inf')
    recognized_user_id = "Unknown"
    
    for user_id, embedding in user_embeddings.items():
        distance = cosine(new_embedding, embedding)
        if distance < min_distance:
            min_distance = distance
            recognized_user_id = user_id

    return f"Recognized User: {recognized_user_id}"

# Gradio interface for registering users
register_interface = gr.Interface(
    fn=register_user,
    inputs=[gr.inputs.Image(shape=(200, 200)), gr.inputs.Textbox(label="User ID")],
    outputs="text",
    live=True
)

# Gradio interface for recognizing users
recognize_interface = gr.Interface(
    fn=recognize_user,
    inputs=gr.inputs.Image(shape=(200, 200)),
    outputs="text",
    live=True
)

if __name__ == "__main__":
    register_interface.launch(share=True)
    recognize_interface.launch(share=True)