import spaces
import base64
import requests
import numpy as np
import face_recognition
import gradio as gr
from io import BytesIO

@spaces.GPU
def get_face_embedding(image_input):
    # Check if the input is a URL
    if isinstance(image_input, str) and (image_input.startswith("http://") or image_input.startswith("https://")):
        response = requests.get(image_input)
        image = face_recognition.load_image_file(BytesIO(response.content))
    else:
        # Assume input is a base64 encoded string
        if ',' in image_input:
            image_input = image_input.split(',')[1]  # Remove the prefix
        img_data = base64.b64decode(image_input)
        image = face_recognition.load_image_file(BytesIO(img_data))

    # Get the face encodings for all faces in the image
    face_encodings = face_recognition.face_encodings(image)

    # If no faces are detected, return an empty list
    if not face_encodings:
        return []

    # Return the first face encoding as a list
    return face_encodings[0].tolist()


# Define the Gradio interface
interface = gr.Interface(
    fn=get_face_embedding,
    inputs="text",
    outputs="json",
    title="Face Embedding Extractor",
    description="Input a base64 encoded image or an image link to get a 128-dimensional face embedding vector. If no face is detected, an empty list is returned."
)

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