app.py

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)