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

# Load YOLO model
model = YOLO("best.pt")  # keep best.pt in the same folder

def predict(image):
    # Run YOLO inference
    results = model.predict(image, save=False)
    
    # Plot results (draw bounding boxes)
    result_image = Image.fromarray(results[0].plot()[:, :, ::-1])  # BGR → RGB
    
    # Collect labels and confidence scores
    labels = []
    for box in results[0].boxes:
        cls = results[0].names[int(box.cls)]
        conf = float(box.conf)
        labels.append(f"{cls}: {conf:.2f}")
    
    return result_image, "\n".join(labels) if labels else "No objects detected"

# Gradio UI
demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=[gr.Image(type="pil"), gr.Textbox(label="Detections")],
    title="YOLOv8 Object Detection",
    description="Upload an image to detect objects with bounding boxes and confidence scores."
)

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