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') # Load your trained model
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(theme=gr.themes.Soft()) as demo: # Added a theme here
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], # Add your example images
["./images.jpeg", 640, 0.4, 0.5], # Add your example images
],
fn=yolov9_inference,
inputs=[img_path, image_size, conf_threshold, iou_threshold],
outputs=[output_numpy],
cache_examples=True,
)
return demo
# Removed the separate CSS variable and added theme to gr.Blocks
demo = gr.Blocks() # Moved the theme application here
with demo:
app()
if __name__ == "__main__":
demo.launch(debug=True, share=True)