Spaces:
Sleeping
Sleeping
import gradio as gr | |
import json | |
from ultralyticsplus import YOLO, render_result | |
# Model Heading and Description | |
model_heading = "YOLOv8 NDL-DocL Datasets" | |
description = """YOLOv8 NDL-DocL Datasets Gradio demo for object detection. Upload an image or click an example image to use.""" | |
article = "<p style='text-align: center'>YOLOv5 NDL-DocL Datasets is an object detection model trained on the <a href=\"https://github.com/ndl-lab/layout-dataset\">NDL-DocL Datasets</a>.</p>" | |
image_path= [ | |
['『源氏物語』(東京大学総合図書館所蔵).jpg', 0.25, 0.45], | |
['『源氏物語』(京都大学所蔵).jpg', 0.25, 0.45], | |
['『平家物語』(国文学研究資料館提供).jpg', 0.25, 0.45] | |
] | |
# Load YOLO model | |
model = YOLO('nakamura196/yolov8-ndl-layout') | |
def yolov8_img_inference( | |
image: gr.Image = None, | |
conf_threshold: gr.Slider = 0.25, | |
iou_threshold: gr.Slider = 0.45, | |
): | |
""" | |
YOLOv8 inference function | |
Args: | |
image: Input image | |
conf_threshold: Confidence threshold | |
iou_threshold: IOU threshold | |
Returns: | |
Rendered image | |
""" | |
results = model.predict(image, conf=conf_threshold, iou=iou_threshold, device="cpu") | |
render = render_result(model=model, image=image, result=results[0]) | |
json_data = json.loads(results[0].tojson()) | |
return render, json_data | |
inputs_image = [ | |
gr.Image(type="filepath", label="Input Image"), | |
gr.Slider(minimum=0.0, maximum=1.0, value=0.25, step=0.05, label="Confidence Threshold"), | |
gr.Slider(minimum=0.0, maximum=1.0, value=0.45, step=0.05, label="IOU Threshold"), | |
] | |
outputs_image =[ | |
gr.Image(type="filepath", label="Output Image"), | |
gr.JSON(label="Output JSON") | |
] | |
demo = gr.Interface( | |
fn=yolov8_img_inference, | |
inputs=inputs_image, | |
outputs=outputs_image, | |
title=model_heading, | |
description=description, | |
examples=image_path, | |
article=article, | |
cache_examples=False | |
) | |
demo.launch(share=False) |