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

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

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

        # Perform inference
        results = model(image_array)

        # Debug: Log the results
        print(f"Detection results: {results}")

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

        # Render results on the image
        rendered_image = results[0].plot()  # Render bounding boxes
        if rendered_image is None:
            print("Rendered image is None. Something went wrong in the plot() method.")
        
        # Debug: Log image shape after rendering
        print(f"Rendered image shape: {rendered_image.shape}")

        # Convert the rendered image to a PIL image for output
        output_image = Image.fromarray(rendered_image)

        return output_image, {cls: 1.0 for cls in detected_classes}  # Dummy scores for simplicity
    except Exception as e:
        print(f"Error during processing: {e}")
        return None, {"Error": str(e)}

# Gradio app configuration
gradio_app = gr.Interface(
    fn=predict,
    inputs=gr.Image(label="Upload an Image", type="pil"),
    outputs=[
        gr.Image(label="Predicted Image with Bounding Boxes"),  # Rendered image with bounding boxes
        gr.Label(label="Detected Classes"),  # Detected class names
    ],
    title="YOLO Object Detection App",
    description="Upload an image, and the YOLO model will detect objects in it.",
)

if __name__ == "__main__":
    gradio_app.launch()