import gradio as gr
import torch

###############

def yolov7_inference(
    image: gr.inputs.Image = None,

):

    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    path = 'y7-prdef.pt'
    model = torch.hub.load("WongKinYiu/yolov7","custom",f"{path}")

    results = model([image], size=640)
    return results.render()[0]
        

inputs = [
    gr.inputs.Image(type="pil", label="Input Image"),
]


demo_app = gr.Interface(
    fn=yolov7_inference,
    inputs=inputs,
    outputs=gr.outputs.Image(type="filepath", label="Output Image"),
    title="Yolov7 | Jar lid product defects",
    examples=['t1.JPG'],
    cache_examples=True,
)
demo_app.launch(debug=True, enable_queue=True)