GuchaIll's picture
Update app.py
22cc8da verified
raw
history blame contribute delete
695 Bytes
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()