import gradio as gr
from huggingface_hub import hf_hub_download
import yolov9

class Inference_Nascent_Spawning_Deriving_From_YOLOv9:
    def __init__(self):
        self.model = None
        self.model_path = None
        self.image_size = None 
        self.conf_threshold = None 
        self.iou_threshold = None 
        
    # Object behavior / Method -> 1
    def download_models(self, model_id):
        hf_hub_download("merve/yolov9", filename=f"{model_id}", local_dir=f"./")
        return f"./{model_id}"
    
    # Object behavior / Method -> 2
    def load_model(self, model_id):
        self.model_path = self.download_models(model_id)
        self.model = yolov9.load(self.model_path, device="cpu") # Inference generated from CPU 
        #self.model = yolov9.load(self.model_path, device="cuda:0")
        
    # Object behavior / Method -> 3
    def configure_model(self, conf_threshold, iou_threshold):
        self.conf_threshold = conf_threshold 
        self.iou_threshold = iou_threshold 
        self.model.conf = self.conf_threshold 
        self.model.iou = self.iou_threshold 
        
    # Object behavior / Method -> 4
    def perform_inference(self, img_path,model_id,image_size, conf_threshold, iou_threshold):
        self.image_size = image_size
        self.load_model(model_id) # Load the model before performing inference
        self.configure_model(conf_threshold, iou_threshold)
        results = self.model(img_path, size=self.image_size)
        output = results.render()
        return output[0]
    
    # Object behavior / Method -> 5 
    # Note: 5 is a method deriving from within the class with the name 
    # Inference_Nascent_Spawning_Deriving_From_YOLOv9
    # One can also declare outside of the OOP as a function, which in turn,
    # calls the methods inside of the OOP leveraging the functionality
    # fostering from each unique Object behavior / Method
    # Personal preference -> This instantiation from within OOP
    def launch_gradio_app(self):
        with gr.Blocks() as gradio_app: 
            with gr.Row():
                with gr.Column():
                    img_path = gr.Image(type='filepath', label='Image')
                    model_id = gr.Dropdown(
                        label="Model",
                        choices=[
                            "gelan-c.pt",
                            "gelan-e.pt",
                            "yolov9-c.pt",
                            "yolov9-e.pt",
                        ],
                        value="gelan-e.pt",
                    )
                    image_size = gr.Slider(
                        label="Image Size",
                        minimum=320,
                        maximum=1280,
                        step=32,
                        value=640, # Default value of 640 foments higher percentage obverse the image detection
                    )
                    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,
                    )
                    yolov9_infer = gr.Button(value="Inference")
                    
                with gr.Column():
                    output_numpy = gr.Image(type="numpy", label="Output")
                    
            # yolov9_infer leveraging click functionality
            # Resembles iface = gr.Interface( 
            #fn=... 
            #inputs=[], 
            #outputs=[],
            yolov9_infer.click(
                fn=self.perform_inference,
                inputs=[
                    img_path,
                    model_id,
                    image_size,
                    conf_threshold,
                    iou_threshold,
                ],
                outputs=[output_numpy],
             
            )
            
            gr.Examples(
                examples=[
                    ["cow.jpeg", "gelan-e.pt", 640, 0.4, 0.5], 
                    ["techengue_GTA.png", "yolov9-c.pt", 640, 0.4, 0.5],
                ],
                fn=self.perform_inference,
                inputs=[
                    img_path,
                    model_id,
                    image_size,
                    conf_threshold,
                    iou_threshold,
                ],
                outputs=[output_numpy],
                cache_examples=True,
            )
            
            gr.HTML(
                """
                <h1 style='text-align: center'>
                YOLOv9: Learning What You Want to Learn Using Programmable Gradient Information
                </h1>
                """
            )
            
            gr.HTML(
            """
            <h3 style='text-align: center'>
            Expound further notions regarding this topic at: 
            https://doi.org/10.48550/arXiv.2402.13616
            """
            )
            
        gradio_app.launch(debug=True)
        
# Instantiate the class and launch the Gradio app
yolo_inference = Inference_Nascent_Spawning_Deriving_From_YOLOv9()
yolo_inference.launch_gradio_app()