|
import gradio as gr |
|
import torch |
|
|
|
from PIL import Image |
|
|
|
model = torch.hub.load('ultralytics/yolov5', 'custom', path='best.pt', force_reload=True) |
|
|
|
def detect(image): |
|
results = model(image) |
|
results.render() |
|
return Image.fromarray(results.imgs[0]) |
|
|
|
|
|
inputs = gr.inputs.Image(type='pil', label="Original Image") |
|
outputs = gr.outputs.Image(type="pil", label="Output Image") |
|
|
|
title = "Object detection from Infrared image using YOLOv5n" |
|
|
|
|
|
gr.Interface(detect, inputs, outputs, title=title, theme="huggingface").launch(debug=True) |
|
|
|
|