import gradio as gr
import spaces
from huggingface_hub import hf_hub_download

import yolov9
                
model = yolov9.load('best.pt')

# classes = ('ball', 'goalkeeper', 'player', 'referee')

def inference(input_img, conf_threshold, iou_threshold):
    # Set model parameters
    model.conf = conf_threshold
    model.iou = iou_threshold
    # Perform inference
    image_size = input_img.shape
    results = model(input_img, size=image_size)

    # Optionally, show detection bounding boxes on image
    output = results.render()
    
    return output[0]
    

def app():
    with gr.Blocks():
        with gr.Row():
            with gr.Column():
                input_img = gr.Image(width= 256, height=256,label="Image")
                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,
                )
                yolo_inf = gr.Button(value="Inference")

            with gr.Column():
                output_val = gr.Image(width= 256, height=256,label="Output Image")
        yolo_inf.click(
            fn= inference,
            inputs = [
                input_img,
                conf_threshold,
                iou_threshold
            ],
            outputs = [output_val],
        )
        gr.Examples([["img1.jpg",0.4, 0.6, 0.4],
                     ["img2.jpg",0.1, 0.2, 1.0]],
                   fn= inference,
                   inputs = [
                       input_img,
                       conf_threshold,
                       iou_threshold
            ],
            outputs = [output_val],
            cache_examples=True,
                   )

demo = gr.Blocks()
with demo:
    gr.HTML(
        """
    <h1 style='text-align: center'>
    YOLOv9
    </h1>
    """)
    gr.HTML(
        """
        <h3 style='text-align: center'>
        Inferencing yolov9 with custom dataset - football players dataset 
        </h3>
        """)
    with gr.Row():
        with gr.Column():
            app()

demo.launch(debug=True)