File size: 3,343 Bytes
97afc34
2ee64bf
97afc34
 
e6c16b3
2ee64bf
 
 
97afc34
2ee64bf
 
97afc34
2ee64bf
 
 
 
 
 
 
 
 
97afc34
2ee64bf
 
 
97afc34
2ee64bf
0496ae0
2ee64bf
 
 
 
 
97afc34
2ee64bf
5d7a51d
2ee64bf
 
 
 
b00892e
7ddfbe4
97afc34
2ee64bf
 
 
 
 
 
 
 
 
 
fd732f2
2ee64bf
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
cf34ca3
f4f13b0
 
 
 
 
 
d98a392
f4f13b0
 
 
 
 
 
 
 
 
 
cf34ca3
2ee64bf
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
import gradio as gr
import spaces
from huggingface_hub import hf_hub_download


def download_models(model_id):
    hf_hub_download("merve/yolov9", filename=f"{model_id}", local_dir=f"./")
    return f"./{model_id}"

@spaces.GPU
def yolov9_inference(img_path, model_id, image_size, conf_threshold, iou_threshold):
    """
    Load a YOLOv9 model, configure it, perform inference on an image, and optionally adjust 
    the input size and apply test time augmentation.
    
    :param model_path: Path to the YOLOv9 model file.
    :param conf_threshold: Confidence threshold for NMS.
    :param iou_threshold: IoU threshold for NMS.
    :param img_path: Path to the image file.
    :param size: Optional, input size for inference.
    :return: A tuple containing the detections (boxes, scores, categories) and the results object for further actions like displaying.
    """
    # Import YOLOv9
    import yolov9
    
    # Load the model
    model_path = download_models(model_id)
    model = yolov9.load(model_path, device="cpu")
    
    # Set model parameters
    model.conf = conf_threshold
    model.iou = iou_threshold
    
    # Perform inference
    results = model(img_path, size=image_size)

    # Optionally, show detection bounding boxes on image
    output = results.render()
    
    return output[0]


def app():
    with gr.Blocks():
        with gr.Row():
            with gr.Column():
                img_path = gr.Image(type="filepath", label="Image")
                model_path = gr.Dropdown(
                    label="Model",
                    choices=[
                        "gelan-c.pt",
                        "gelan-e.pt",
                        "yolov9-c.pt",
                        "yolov9-e.pt",                                              
                    ],
                    value="gelan-e.pt",
                )
                image_size = gr.Slider(
                    label="Image Size",
                    minimum=320,
                    maximum=1280,
                    step=32,
                    value=640,
                )
                conf_threshold = gr.Slider(
                    label="Confidence Threshold",
                    minimum=0.1,
                    maximum=1.0,
                    step=0.1,
                    value=0.4,
                )
                iou_threshold = gr.Slider(
                    label="IoU Threshold",
                    minimum=0.1,
                    maximum=1.0,
                    step=0.1,
                    value=0.5,
                )
                yolov9_infer = gr.Button(value="Inference")

            with gr.Column():
                output_numpy = gr.Image(type="numpy",label="Output")

        yolov9_infer.click(
            fn=yolov9_inference,
            inputs=[
                img_path,
                model_path,
                image_size,
                conf_threshold,
                iou_threshold,
            ],
            outputs=[output_numpy],
        )   

gradio_app = gr.Blocks()
with gradio_app:
    gr.HTML(
        """
    <h1 style='text-align: center'>
    YOLOv9 Base Model
    </h1>
    """)
    gr.HTML(
        """
        <h3 style='text-align: center'>
        </h3>
        """)
    with gr.Row():
        with gr.Column():
            app()    
    
gradio_app.launch(debug=True)