burn-detection / app.py
michaelwja's picture
Update app.py
fbd683f
raw
history blame
618 Bytes
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()