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

@spaces.GPU
def get_face_embedding(base64_image):
    # Decode the base64 image
    img_data = base64.b64decode(base64_image)
    np_arr = np.frombuffer(img_data, np.uint8)
    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 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)