from PIL import Image
import numpy as np
from ultralytics import YOLO
import gradio as gr

# Load the YOLO model
#MODEL_URL = 'https://huggingface.co/ayoubsa/yolo_model/resolve/main/best.pt'
MODEL_URL = 'best.pt'
model = YOLO(MODEL_URL, cache_dir="/root/.cache")

# Define the prediction function for Gradio
def predict(image):
    try:
        # Convert PIL image to NumPy array
        image_array = np.array(image)

        # Perform prediction
        results = model(image_array)

        # Access the first result
        result = results[0]

        # Extract detected classes
        detected_classes = [model.names[int(cls)] for cls in result.boxes.cls]
        print(f"Detected classes: {detected_classes}")

        # Render bounding boxes on the image
        annotated_image = result.plot()

        # Convert the annotated image to PIL format
        output_image = Image.fromarray(annotated_image)

        # Return the annotated image and detected classes as output
        return output_image, detected_classes

    except Exception as e:
        print("Error during prediction:", str(e))
        return None, ["Error during processing"]

# Create the Gradio interface
iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Upload an Image"),  # Input image as PIL
    outputs=[
        gr.Image(type="pil", label="Predicted Image with Bounding Boxes"),  # Annotated image
        gr.Label(label="Detected Classes")  # Detected classes
    ],
    title="YOLO Object Detection App",
    description="Upload an image, and the YOLO model will detect objects and annotate them with bounding boxes and class labels."
)

# Launch the Gradio app
iface.launch()