chanelcolgate's picture
Update app.py
6f21361 verified
import json
import glob
from collections import Counter
import requests
import gradio as gr
from ultralyticsplus import YOLO, download_from_hub, render_result
hf_model_ids = [
"chanelcolgate/chamdiemgianhang-vsk",
"chanelcolgate/chamdiemgianhang-vsk-v2",
"chanelcolgate/chamdiemgianhang-vsk-v4",
"chanelcolgate/chamdiemgianhang-vsk-v5",
"chanelcolgate/chamdiemgianhang-vsk-v6",
]
image_paths = [
[image_path, "chanelcolgate/chamdiemgianhang-vsk-v2", 640, 0.25, 0.45]
for image_path in glob.glob("./tmp/*.jpg")
]
def detection_image(
image=None,
hf_model_id="chanelcolgate/chamdiemgianhang-vsk-v2",
image_size=640,
conf_threshold=0.25,
iou_threshold=0.45,
):
model_path = download_from_hub(hf_model_id)
model = YOLO(model_path)
results = model(image, imgsz=image_size, conf=conf_threshold, iou=iou_threshold)
json_result = json.loads(results[0].tojson())
class_counts = Counter(detection["name"] for detection in json_result)
render = render_result(model=model, image=image, result=results[0])
return render, class_counts
def detection_image_link(
image=None,
hf_model_id="chanelcolgate/chamdiemgianhang-vsk-v2",
image_size=640,
conf_threshold=0.25,
iou_threshold=0.45,
):
model_path = download_from_hub(hf_model_id)
model = YOLO(model_path)
results = model(image, imgsz=image_size, conf=conf_threshold, iou=iou_threshold)
json_result = json.loads(results[0].tojson())
class_counts = Counter(detection["name"] for detection in json_result)
render = render_result(model=model, image=image, result=results[0])
return render, class_counts
title = "Cham Diem Gian Hang VSK"
interface = gr.Interface(
fn=detection_image,
inputs=[
gr.Image(type="pil"),
gr.Dropdown(hf_model_ids),
gr.Slider(minimum=320, maximum=1280, value=640, step=32, label="Image Size"),
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=[gr.Image(type="pil"), gr.Textbox(show_label=False)],
title=title,
examples=image_paths,
cache_examples=True if image_paths else False,
)
interface_link = gr.Interface(
fn=detection_image,
inputs=[
gr.Textbox(label="Image Link"),
gr.Dropdown(hf_model_ids),
gr.Slider(minimum=320, maximum=1280, value=640, step=32, label="Image Size"),
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=[gr.Image(type="pil"), gr.Textbox(show_label=False)],
title=title,
)
gr.TabbedInterface(
[interface, interface_link], tab_names=["Image inference", "Image link inference"]
).queue().launch()