import gradio as gr
import torch
from PIL import Image
import os
import yolov9
HTML_TEMPLATE = """
"""
# The rest of the Python code remains the same
def yolov9_inference(img_path, image_size, conf_threshold, iou_threshold):
model = yolov9.load('./best.pt')
model.conf = conf_threshold
model.iou = iou_threshold
results = model(img_path, size=image_size)
output = results.render()
return output[0]
def app():
with gr.Blocks() as demo:
gr.HTML(HTML_TEMPLATE)
with gr.Row():
with gr.Column(scale=1, min_width=300):
img_path = gr.Image(type="filepath", label="Upload Image")
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)
detect_button = gr.Button("Detect Manholes", variant="primary")
with gr.Column(scale=1, min_width=300):
output_numpy = gr.Image(type="numpy", label="Detection Result")
detect_button.click(
fn=yolov9_inference,
inputs=[img_path, image_size, conf_threshold, iou_threshold],
outputs=[output_numpy]
)
gr.Examples(
examples=[
["./openmanhole.jpg", 640, 0.4, 0.5],
["./images.jpeg", 640, 0.4, 0.5],
],
fn=yolov9_inference,
inputs=[img_path, image_size, conf_threshold, iou_threshold],
outputs=[output_numpy],
cache_examples=True,
)
return demo
css = """
/* You can add any additional CSS here to style Gradio components */
.gradio-container {
max-width: 900px !important;
margin-left: auto !important;
margin-right: auto !important;
}
"""
demo = gr.Blocks(css=css)
with demo:
app()
if __name__ == "__main__":
demo.launch(debug=True, share=True)