File size: 2,966 Bytes
0be45a9
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
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",
]

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()