import gradio as gr
import torch
import numpy as np
from PIL import Image

# Load the custom-trained YOLOv5 model
model = torch.hub.load('ultralytics/yolov5', 'custom', path='best.pt', _verbose=False)

def gradio_wrapper(img):
    global model
    # Run the model on the input image and get the results
    results = model(img)
    # Render the results and return the annotated image
    return results.render()[0]

# Set up the Gradio image input component with the option to upload an image or use the webcam
s = gr.Radio(["upload", "webcam"], type="index")

if s == "upload":
    image = gr.inputs.Image(source="upload")
else:
    image = gr.inputs.Image(source="webcam")


# Create the Gradio interface with the gradio_wrapper function, the image input component, and an image output component
demo = gr.Interface(
    gradio_wrapper,
    image,
    'image',
    live=True,
    title="CiclopeIA",
    description="App based on the CiclopeIA project from Saturdays AI. Identifies the value of Euro banknotes."
)

# Launch the Gradio interface
demo.launch()