import gradio as gr
import torch
from PIL import Image

# Load the trained YOLO model (or any other model you're using)
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')  # Replace with your model

# Define a prediction function
def predict(image):
    results = model(image)
    return results.render()[0]  # Returns the annotated image

# Create a Gradio interface
iface = gr.Interface(fn=predict, inputs=gr.Image(), outputs=gr.Image())

# Launch the interface
iface.launch()