yolo9 / app.py
Parthebhan's picture
Update app.py
cc749dd verified
raw
history blame
1.78 kB
import gradio as gr
from huggingface_hub import hf_hub_download
import os
import cv2 # Import OpenCV
# 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
def yolov9_inference(img_path):
"""
Perform inference on an image using the YOLOv9 model.
:param img_path: Path to the image file.
:return: Output image with detections.
"""
# Load the image from the file path
image = cv2.imread(img_path) # Use OpenCV to read the image
# Load the model
model_path = download_models("gelan-c-seg.pt")
model = cv2.dnn.readNetFromDarknet(model_path)
# Perform inference
# (Add your inference code here)
# Optionally, show detection bounding boxes on image
output_image = image # Placeholder for the output image
return output_image
def app():
with gr.Blocks():
with gr.Row():
with gr.Column():
img_path = gr.Image(type="file", label="Upload Image")
inference_button = gr.Button(label="Run Inference")
with gr.Column():
output_image = gr.Image(label="Output Image")
inference_button.click(
fn=yolov9_inference,
inputs=[img_path],
outputs=[output_image],
)
gradio_app = gr.Interface(
app,
title="YOLOv9 Inference",
description="Perform object detection using the YOLOv9 model.",
)
gradio_app.launch(debug=True)