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)}"

@spaces.GPU
def image_to_base64(image):
    buffered = BytesIO()
    image.save(buffered, format="JPEG")
    img_str = base64.b64encode(buffered.getvalue()).decode()
    return img_str

# Define the Gradio interfaces
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 to convert it to a base64 encoded string."
)

face_crop_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__":
    gr.TabbedInterface([base64_converter_interface, face_crop_interface], ["Convert to Base64", "Crop Face"]).launch(share=True)