import gradio as gr
import onnxruntime as ort
from PIL import Image
import numpy as np

# Load the YOLOv8 ONNX model
session = ort.InferenceSession("yolov8n.pt")

def predict(image):
    # Preprocess the image
    img = np.array(image).astype(np.float32)
    img = np.expand_dims(img, axis=0)

    # Run inference
    outputs = session.run(None, {"images": img})[0]

    return outputs  # Modify to return detections

interface = gr.Interface(
    fn=predict, 
    inputs=gr.Image(type="pil"), 
    outputs="label"
)

interface.launch()