File size: 3,507 Bytes
7bed26b
36e4f2e
dc51cfe
36e4f2e
acb05a8
 
36e4f2e
 
acb05a8
e1d655d
acb05a8
7bed26b
36e4f2e
acb05a8
 
7065ee0
 
 
 
acb05a8
 
 
 
 
 
 
 
 
 
 
 
7bed26b
acb05a8
 
 
36e4f2e
acb05a8
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
e7b97f5
acb05a8
 
 
 
 
 
 
36e4f2e
acb05a8
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
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
import gradio as gr
import spaces
import supervision as sv
import PIL.Image as Image
from ultralytics import YOLO, YOLOv10
from huggingface_hub import hf_hub_download


def download_models(model_id):
    hf_hub_download("atalaydenknalbant/asl-models", filename=f"{model_id}", local_dir=f"./")
    return f"./{model_id}"



box_annotator = sv.BoxAnnotator()
category_dict = {0: 'A', 1: 'B', 2: 'C', 3: 'D', 4: 'E', 5: 'F', 6: 'G', 7: 'H', 8: 'I',
                 9: 'J', 10: 'K', 11: 'L', 12: 'M', 13: 'N', 14: 'O', 15: 'P', 16: 'Q',
                 17: 'R', 18: 'S', 19: 'T', 20: 'U', 21: 'V', 22: 'W', 23: 'X', 24: 'Y', 25: 'Z'}




@spaces.GPU(duration=200)
def yolo_inference(image, model_id, conf_threshold, iou_threshold):
    model_path = download_models(model_id)
    if model_id[:7] == 'yolov10':
        model = YOLOv10(model_path)
    else:
        model = YOLO(model_path)
    results = model(source=image, imgsz=416, iou=iou_threshold, conf=conf_threshold, verbose=False, max_det=1)[0]
    detections = sv.Detections.from_ultralytics(results)
    
    labels = [
        f"{category_dict[class_id]} {confidence:.2f}"
        for class_id, confidence in zip(detections.class_id, detections.confidence)
    ]
    annotated_image = box_annotator.annotate(image, detections=detections, labels=labels)

    return annotated_image

def app():
    with gr.Blocks():
        with gr.Row():
            with gr.Column():
                image = gr.Image(type="pil", label="Image")
                
                model_id = gr.Dropdown(
                    label="Model",
                    choices=[
                        "yolov10x.pt",
                        "yolov10s.pt",
                        "yolov9e.pt",
                        "yolov8x.pt",
                    ],
                    value="yolov10s.pt",
                )
                conf_threshold = gr.Slider(
                    label="Confidence Threshold",
                    minimum=0.1,
                    maximum=1.0,
                    step=0.1,
                    value=0.25,
                )
                iou_threshold = gr.Slider(
                    label="IoU Threshold",
                    minimum=0.1,
                    maximum=1.0,
                    step=0.1,
                    value=0.45,
                )
                yolov10_infer = gr.Button(value="Detect Objects")

            with gr.Column():
                output_image = gr.Image(type="pil", label="Annotated Image")

        yolov10_infer.click(
            fn=yolo_inference,
            inputs=[
                image,
                model_id,
                conf_threshold,
                iou_threshold,
            ],
            outputs=[output_image],
        )

        gr.Examples(
            examples=[
                [
                    "b.jpg",
                    "yolov10x.pt",
                    0.25,
                    0.45,
                ],
                [
                    "y.jpg",
                    "yolov10x.pt",
                    0.25,
                    0.45,
                ],
            ],
            fn=yolo_inference,
            inputs=[
                image,
                model_id,
                conf_threshold,
                iou_threshold,
            ],
            outputs=[output_image],
            cache_examples="lazy",
        )

gradio_app = gr.Blocks()
with gradio_app:
    with gr.Row():
        with gr.Column():
            app()

gradio_app.launch(debug=True)