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

# Load YOLOv8 model
model = YOLO("yolov8n.pt")

# Detection function
def detect_objects(image):
    if image is None:
        return None
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)
    results = model.predict(image)
    result_image = results[0].plot()
    return Image.fromarray(result_image)

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## 🧠 YOLOv8 Object Detection (Image Upload or Camera)")

    image_input = gr.Image(label="📷 Upload or Capture Image", type="numpy", height=300)
    result_output = gr.Image(label="🔍 Detection Result", type="pil", height=300)

    detect_button = gr.Button("🚀 Detect Objects")

    detect_button.click(fn=detect_objects, inputs=image_input, outputs=result_output)

demo.launch()