import gradio as gr
import yolov5
from PIL import Image

# Load YOLOv5 model
model = yolov5.load("keremberke/yolov5n-garbage")

# Set model parameters
model.conf = 0.25  # Confidence threshold
model.iou = 0.45  # IoU threshold

def predict(img):
    # Convert image to PIL format
    img = Image.fromarray(img)
    # Perform inference
    results = model(img, size=640)
    # Show results
    results.save(save_dir="results/")
    return "results/image0.jpg"

# Gradio UI
iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(),
    outputs=gr.Image(),
    title="Garbage Object Detection",
    description="Upload an image and the model will detect garbage objects in it."
)

iface.launch()