import gradio as gr
import numpy as np
import tensorflow as tf
from PIL import Image
import cv2

# Load TensorFlow Lite model
interpreter = tf.lite.Interpreter(model_path="facenet.tflite")
interpreter.allocate_tensors()

# Get input and output details
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

def preprocess_image(image):
    """
    Preprocess the input image for the FaceNet model.
    """
    image = Image.fromarray(image)
    image = image.resize((160, 160))  # Resize to the model's input size
    image_array = np.asarray(image).astype(np.float32)
    image_array = (image_array - 127.5) / 127.5  # Normalize to [-1, 1]
    image_array = np.expand_dims(image_array, axis=0)  # Add batch dimension
    return image_array

def create_face_embedding(image):
    """
    Generate a face embedding for the given image.
    """
    processed_image = preprocess_image(image)

    # Run the model
    interpreter.set_tensor(input_details[0]['index'], processed_image)
    interpreter.invoke()

    # Extract the embedding
    embedding = interpreter.get_tensor(output_details[0]['index'])
    return embedding.flatten()

# Gradio interface
def generate_embedding(image):
    """
    Gradio function to process the image and return full embeddings.
    """
    try:
        embedding = create_face_embedding(image)
        return embedding.tolist()  # Convert numpy array to list
    except Exception as e:
        return f"Error: {e}"

# Gradio interface setup
iface = gr.Interface(
    fn=generate_embedding,
    inputs=gr.Image(type="numpy", label="Upload Face Image"),
    outputs=gr.JSON(label="Face Embedding"),
    title="Face Embedding Generator",
    description="Upload a face image to generate a 512-dimensional embedding using the FaceNet model."
)

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