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

@spaces.GPU
def crop_face(base64_image):
    try:
        # Decode the base64 image to an OpenCV format
        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:
            return "Image decoding failed. Check the input format."

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

        # Convert the image to grayscale for face detection
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        faces = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5, minSize=(30, 30))

        if len(faces) == 0:
            return "No faces detected in the image."

        # Crop the first detected face
        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

    except Exception as e:
        return f"An error occurred: {str(e)}"

interface = gr.Interface(
    fn=crop_face,
    inputs=gr.Textbox(),
    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)