Spaces:
Runtime error
Runtime error
import gradio as gr | |
import torch | |
from PIL import Image, ImageDraw | |
import json | |
# Model | |
model = torch.hub.load('ultralytics/yolov5', 'custom', path='best.pt', source="local") | |
def yolo(im, size=1024): | |
g = (size / max(im.size)) # gain | |
im = im.resize((int(x * g) for x in im.size), Image.BICUBIC) # resize | |
results = model(im) # inference | |
df = results.pandas().xyxy[0].to_json(orient="records") | |
res = json.loads(df) | |
detected_images = [] | |
im_draw = im.copy() | |
draw = ImageDraw.Draw(im_draw) | |
# JSONデータ内の座標に基づいて矩形を描画 | |
for item in res: | |
xmin = item['xmin']# * w | |
ymin = item['ymin']# * h | |
xmax = item['xmax']# * w | |
ymax = item['ymax']# * h | |
draw.rectangle([(xmin, ymin), (xmax, ymax)], outline="red", width=2) | |
# Extract each detected object into a separate image | |
detected_object = im.crop((xmin, ymin, xmax, ymax)) | |
detected_images.append(detected_object) | |
return [ | |
res, | |
im_draw, | |
detected_images | |
] | |
inputs = [ | |
gr.Image(type='pil', label="Original Image"), | |
] | |
outputs = [ | |
gr.JSON(label="JSON Output"), | |
gr.Image(type="pil", label="Output Image"), | |
gr.Gallery(label="Detected Objects", object_fit="contain") | |
] | |
title = "YOLOv5 Kunshujo" | |
description = "YOLOv5 Kunshujo Gradio demo for object detection. Upload an image or click an example image to use." | |
article = "<p style='text-align: center'>YOLOv5 Kunshujo is an object detection model trained on the <a href=\"https://github.com/utda/kunshujo-layout-dataset\">Kunshujo layout dataset</a>.<br/>Examples: 張交帖.[5] (国立国会図書館 National Diet Library, JAPAN)</p>" | |
examples = [ | |
['2586696_R0000008.jpg'], | |
['2586696_R0000009.jpg'] | |
] | |
demo = gr.Interface(yolo, inputs, outputs, title=title, description=description, article=article, examples=examples, allow_flagging="never") | |
demo.launch() |