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

# Load YOLOv11 Model
model_path = "best.pt" 
model = YOLO(model_path)

def predict(image):
    image = np.array(image)
    results = model(image)

    # Draw bounding boxes
    for result in results:
        for box in result.boxes:
            x1, y1, x2, y2 = map(int, box.xyxy[0])
            conf = box.conf[0]
            cls = int(box.cls[0])
            label = f"{model.names[cls]} {conf:.2f}"
            cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
            cv2.putText(image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

    return Image.fromarray(image)

# Gradio Interface
iface = gr.Interface(fn=predict, inputs="image", outputs="image", title="YOLOv11 Object Detection")
iface.launch()