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):
    try:
        buffered = BytesIO()
        image.save(buffered, format="JPEG")
        img_str = base64.b64encode(buffered.getvalue()).decode('utf-8')
        print("Encoded String:", img_str[:100])  # Print the first 100 characters
        return img_str
    except Exception as e:
        print("Error:", str(e))
        return f"An error occurred: {str(e)}"



# Define the Gradio interfaces


base64_converter_interface = gr.Interface(
    fn=image_to_base64,
    inputs=gr.Image(),
    outputs="text",
    title="Image to Base64 Converter",
    description="Upload an image to convert it to a Base64 encoded string."
)

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