import numpy as np
import gradio as gr
from PIL import Image
from sahi import AutoDetectionModel
from sahi.predict import get_sliced_prediction
from sahi.utils.cv import visualize_object_predictions
from huggingface_hub import hf_hub_download
from datasets import load_dataset

model_path = hf_hub_download("lukas-aebi/yolov8x-condenser-detection", "best.pt")

dataset = load_dataset("lukas-aebi/examples-condenser-detection")

def get_predictions(image: Image.Image, threshold: float):

    model = AutoDetectionModel.from_pretrained(
    model_type="yolov8",
    model_path=model_path,
    confidence_threshold=threshold,
    device="cpu"
    )
    
    result = get_sliced_prediction(
            image=image,
            detection_model=model,
            slice_height=358,
            slice_width=358,
            overlap_height_ratio=0.1,
            overlap_width_ratio=0.1
    )
    
    return Image.fromarray(visualize_object_predictions(
        image=np.array(result.image),
        object_prediction_list=result.object_prediction_list,
        hide_labels=True,
    )["image"])


with gr.Blocks() as demo:
    gr.Markdown(
        """
        # Condenser Detection
        * Demo Application for condenser detection on aerial images.
        * Works best with images of size 1000x1000.
        """
    )
    with gr.Row():
        with gr.Column():
            slider = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, label="Confidence Threshold")
        # with gr.Column():
        #     model = gr.Dropdown(choices=["lukas-aebi/yolov8x-condenser-detection"],
        #                         value="lukas-aebi/yolov8x-condenser-detection",
        #                         type="index",
        #                         label="Model")

    with gr.Row():
        with gr.Column():
            input_img = gr.Image(type="pil", label="Input", height=600, width=600)
        with gr.Column():
            output_img = gr.Image(type="pil", label="Output", height=600, width=600)
        
    with gr.Row():
        image_button = gr.Button("Detect Condensers")
    
    image_button.click(
        fn=get_predictions, 
        inputs=[input_img, slider], 
        outputs=output_img,
        api_name="condenser-detection"
            )

    # gr.Examples(
    #     dataset["train"]["image"],
    #     input_img,
    #     output_img,
    #     get_predictions,
    #     cache_examples=True,
    # )

demo.launch()