import gradio as gr
import torch
from PIL import Image

# Load the model using torch.hub (you must have the model locally)
model_name = "skin_burn_2022_8_21 (1).pt"
model =torch.hub.load("WongKinYiu/yolov7", 'custom',model_name)

def predict(img):
    # Run inference with YOLO model
    results = model(img)
    # Convert the result to an image
    result_img = Image.fromarray(results.render())
    return result_img

# Define a Gradio interface
iface = gr.Interface(
    fn=predict,  # the function to wrap
    inputs=gr.inputs.Image(),  # input type
    outputs=gr.outputs.Image()  # output type
)

iface.launch()