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

@spaces.GPU
def crop_face(base64_image):
    try:
        # 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:
            return "Could not decode the image or no data in buffer"
    
        # 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 message
        if len(faces) == 0:
            return "No faces found"

        # 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

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

def image_to_base64(image):
    # Convert PIL Image to bytes
    buffered = io.BytesIO()
    image.save(buffered, format="JPEG")
    # Encode bytes to Base64 string
    img_str = base64.b64encode(buffered.getvalue()).decode()
    return img_str


# Define the Gradio interface using the updated syntax
base64_converter_interface = gr.Interface(
    fn=image_to_base64, 
    inputs=gr.Image(type="pil"), 
    outputs=gr.Textbox(),
    title="Image to Base64 Encoder",
    description="Upload an image and convert it to a Base64 encoded string."
)

face_crop_interface = gr.Interface(
    fn=crop_face,
    inputs=gr.Textbox(),
    outputs=gr.Textbox(),
    title="Face Cropper",
    description="Input a base64 encoded image to get a base64 encoded cropped face."
)

if __name__ == "__main__":
    gr.TabbedInterface([base64_converter_interface, face_crop_interface], ["Convert to Base64","Crop Face"]).launch()