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

def crop_face(base64_image):
    # Decode the base64 image
    img_data = base64.b64decode(base64_image)
    np_arr = np.frombuffer(img_data, np.uint8)
    image = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
    
    if image is None:
        print("Could not decode the image")
        return None

    # Load the pre-trained face detector
    face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')

    # Convert the image to grayscale
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    # Detect faces in the image
    faces = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5, minSize=(30, 30))

    # If no faces are detected, return None
    if len(faces) == 0:
        print("No faces found")
        return None

    # Crop the first face found
    x, y, w, h = faces[0]
    face_crop = image[y:y+h, x:x+w]

    # Encode the cropped face to base64
    _, buffer = cv2.imencode('.jpg', face_crop)
    face_base64 = base64.b64encode(buffer).decode('utf-8')

    return face_base64

# Define the Gradio interface
interface = gr.Interface(
    fn=crop_face,
    inputs="text",
    outputs="text",
    title="Face Cropper",
    description="Input a base64 encoded image to get a base64 encoded cropped face."
)

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