burn-detection / app.py
michaelwja's picture
Create app.py
3cbf1d1
raw
history blame
604 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.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()