File size: 4,004 Bytes
97afc34
 
 
e6c16b3
97afc34
e6c16b3
97afc34
e6c16b3
 
 
 
 
 
 
97afc34
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
e7f051a
97afc34
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
dd79218
97afc34
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c3f9ae7
 
 
 
 
 
 
 
 
 
 
97afc34
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c3f9ae7
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
120
121
122
import gradio as gr
import spaces
from huggingface_hub import hf_hub_download
import os

# Function to download models from Hugging Face
def download_models(model_id):
    # Check if the model file exists locally
    local_model_path = "./gelan-c-seg.pt"
    if not os.path.exists(local_model_path):
        # Download the model from Hugging Face if it doesn't exist locally
        hf_hub_download("merve/yolov9", filename="gelan-c-seg.pt", local_dir="./")
    return local_model_path


@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-seg.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: Learning What You Want to Learn Using Programmable Gradient Information
    </h1>
    """)
    gr.HTML(
        """
        <h3 style='text-align: center'>
        Follow me for more!
        <a href='https://twitter.com/kadirnar_ai' target='_blank'>Twitter</a> | <a href='https://github.com/kadirnar' target='_blank'>Github</a> | <a href='https://www.linkedin.com/in/kadir-nar/' target='_blank'>Linkedin</a>  | <a href='https://www.huggingface.co/kadirnar/' target='_blank'>HuggingFace</a>
        </h3>
        """)
    with gr.Row():
        with gr.Column():
            app()

gradio_app.launch(debug=True)