import gradio as gr
from ultralytics import YOLO
import cv2
import numpy as np
import os
import tempfile
import base64
from io import BytesIO
from PIL import Image

def perform(image, task: str="detect"):
    # Predict with the model
    results = model(image, save=True)  # predict on an image
    
    # get the saved prediction
    prediction_image = Image.open(os.path.join(results[0].save_dir, os.listdir(results[0].save_dir)[0]))
    
    # Convert PIL image to base64
    buffered = BytesIO()
    prediction_image.save(buffered, format="JPEG")
    img_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8')
    
    # initialize the output list to save the bbox
    outputs = []
    
    # Access the results
    for result in results[0]:
        xywh = result.boxes.xyxy[0].numpy().tolist()  # top-left-x, top-left-y, bottom-right-x, bottom-right-y
        name = result.names[result.boxes.cls.item()]  # classes
        confs = float(result.boxes.conf)  # confidence score of each box
        outputs.append({"xywh": xywh, "conf": confs, "name": name})
    
    
    # return the PIL image (for display) and the outputs with base64
    return prediction_image, outputs, {"image_base64": f"data:image/jpeg;base64,{img_base64}"}

def load_model(model_name: str = "yolo11n.pt"):
    """Model loading"""
    try:
        print(f"Loading the model {model_name}")
        model = YOLO(model_name)
    except Exception as e:
        print(f"Error loading the model {e}")
    return model

# loading the model
model = load_model()

# Create the Gradio interface
demo = gr.Interface(
    fn=perform,
    inputs=[gr.Image(type="pil")],  # PIL Image format
    outputs=[gr.Image(type="pil"), gr.JSON(), gr.JSON()],
    title="Object Detection using YOLO",
    description="Upload an image and get image with bbox in it",
    examples=None  # You can add example images here if needed
)

# Launch the app
if __name__ == "__main__":
    demo.launch()