|
import gradio as gr |
|
import torch |
|
from yolov6 import YOLOV6 |
|
|
|
|
|
torch.hub.download_url_to_file('https://raw.githubusercontent.com/kadirnar/dethub/main/data/images/highway.jpg', 'highway.jpg') |
|
torch.hub.download_url_to_file('https://user-images.githubusercontent.com/34196005/142742872-1fefcc4d-d7e6-4c43-bbb7-6b5982f7e4ba.jpg', 'highway1.jpg') |
|
|
|
|
|
def yolov6_inference( |
|
image: gr.inputs.Image = None, |
|
model_path: gr.inputs.Dropdown = None, |
|
image_size: gr.inputs.Slider = 640, |
|
conf_threshold: gr.inputs.Slider = 0.25, |
|
iou_threshold: gr.inputs.Slider = 0.45, |
|
): |
|
""" |
|
YOLOv6 inference function |
|
Args: |
|
image: Input image |
|
model_path: Path to the model |
|
image_size: Image size |
|
conf_threshold: Confidence threshold |
|
iou_threshold: IOU threshold |
|
Returns: |
|
Rendered image |
|
""" |
|
|
|
model = YOLOV6(model_path, device="cpu", hf_model=True) |
|
model.conf_thres = conf_threshold |
|
model.iou_thresh = iou_threshold |
|
model.save_img = True |
|
model.font_path = "Arial.ttf" |
|
pred = model.predict(source=image, img_size=image_size, yaml="coco.yaml") |
|
return pred |
|
|
|
|
|
inputs = [ |
|
gr.inputs.Image(type="filepath", label="Input Image"), |
|
gr.inputs.Dropdown( |
|
label="Model", |
|
choices=[ |
|
"kadirnar/yolov6n-v3.0", |
|
"kadirnar/yolov6s-v3.0", |
|
"kadirnar/yolov6m-v3.0", |
|
"kadirnar/yolov6l-v3.0", |
|
"kadirnar/yolov6s6-v3.0", |
|
"kadirnar/yolov6m6-v3.0", |
|
"kadirnar/yolov6l6-v3.0", |
|
], |
|
default="kadirnar/yolov6s-v3.0", |
|
), |
|
gr.inputs.Slider(minimum=320, maximum=1280, default=1280, step=32, label="Image Size"), |
|
gr.inputs.Slider(minimum=0.0, maximum=1.0, default=0.25, step=0.05, label="Confidence Threshold"), |
|
gr.inputs.Slider(minimum=0.0, maximum=1.0, default=0.45, step=0.05, label="IOU Threshold"), |
|
] |
|
|
|
outputs = gr.outputs.Image(type="filepath", label="Output Image") |
|
title = "YOLOv6: a single-stage object detection framework dedicated to industrial applications." |
|
|
|
examples = [['highway1.jpg', 'kadirnar/yolov6m6-v3.0', 1280, 0.25, 0.45],['highway.jpg', 'kadirnar/yolov6s6-v3.0', 1280, 0.25, 0.45]] |
|
|
|
demo_app = gr.Interface( |
|
fn=yolov6_inference, |
|
inputs=inputs, |
|
outputs=outputs, |
|
title=title, |
|
examples=examples, |
|
cache_examples=True, |
|
theme='huggingface', |
|
) |
|
demo_app.launch(debug=True, enable_queue=True) |
|
|
