import gradio as gr
import torch


title = "Object Detection with YOLOv5"
description = "Detect objects in images with YOLOv5 (You Only Look Once v5)"
article = "<p style='text-align: center'><a href='https://news.machinelearning.sg/posts/object_detection_with_yolov5/' target='blank' class='footer'>Blog</a> | <a href='https://github.com/eugenesiow/practical-ml' target='_blank' class='footer'>Github Repo</a></p>"


def inference(img):
    results = model([img], size=640)
    results.render()
    return results.imgs[0]


torch.hub.download_url_to_file('https://images.unsplash.com/photo-1526497127495-3b388dc87620?auto=format&fit=crop&w=640&q=80',
                               'baseball.jpg')
torch.hub.download_url_to_file('https://images.unsplash.com/photo-1608278047522-58806a6ac85b?auto=format&fit=crop&w=640&q=80', 'hulk.jpg')
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')


gr.Interface(
    inference,
    gr.inputs.Image(type="pil", label="Input"),
    gr.outputs.Image(type="pil", label="Output"),
    title=title,
    description=description,
    article=article,
    examples=[['baseball.jpg'], ['hulk.jpg']],
    enable_queue=True,
    allow_flagging=False,
).launch(debug=False)