import torch
import gradio as gr
from gradio import components
from PIL import Image

model = None

def object_detection(im):
    global model
    if model is None:
        model = torch.hub.load('ultralytics/yolov5', 'custom', path='best.pt', force_reload=True)
    results = model(im)
    results.render()
    return Image.fromarray(results.ims[0])

image = components.Image(shape=(640, 640), image_mode="RGB", source="upload", label="Image")
outputs = components.Image(type="pil", label="Output Image")

iface = gr.Interface(
    fn=object_detection,
    inputs=image,
    outputs=outputs,
    title='Garbage Detection',
    description='A simple demo app for an object detection model to detect garbage in natural and urban environments.'
)
iface.launch(debug=True)