Spaces:
Sleeping
Sleeping
import gradio as gr | |
from ultralytics import YOLO | |
# Import YOLOv9 | |
import yolov9 | |
# Define function to perform prediction with YOLOv9 model | |
def predict_image(image, model_path, image_size, conf_threshold, iou_threshold): | |
# Load YOLO model | |
model = YOLO(model_path) | |
# Perform inference with YOLO model | |
results = model(image, size=image_size, conf=conf_threshold, iou=iou_threshold) | |
# Render the output | |
output_image = results.render() | |
return output_image[0] | |
# Define Gradio interface | |
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=[ | |
"yolov9c-seg.pt", | |
"yolov5m.pt", | |
"yolov5l.pt", | |
"yolov5x.pt", | |
], | |
value="yolov5s.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="Submit") | |
with gr.Column(): | |
output_numpy = gr.Image(type="numpy", label="Output") | |
yolov9_infer.click( | |
fn=predict_image, | |
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) | |