import glob
import gradio as gr
from ultralytics import YOLO
from PIL import Image

# Load the ONNX model with task explicitly set to "detect"
model = YOLO("best.onnx", task="detect")  # Specify 'task=detect' for detection models

# Prediction function
def predict(image, confidence_threshold):
    # Run inference on the image using the ONNX model
    results = model(image, conf=confidence_threshold)
    
    # Get and return the result image with bounding boxes
    result_image = results[0].plot()[:, :, ::-1]  # Convert BGR to RGB for Gradio
    return result_image

# Function to clear the image input and result
def clear():
    return None, None  # Clears both the input and output image components

# Gradio app configuration
app_title = "🐟Fish Detector (Grayscale) ONNX Gradio Demo"
app_description = """
Upload an image to detect fish using an ONNX-optimized YOLO model. Adjust the confidence threshold to refine detection sensitivity.

**Links:**
- [Model on Hugging Face](https://huggingface.co/akridge/yolo11-fish-detector-grayscale)
- [Dataset on Hugging Face](https://huggingface.co/datasets/akridge/MOUSS_fish_segment_dataset_grayscale)
"""

# Load example images if they are available in an images folder
examples = glob.glob("images/*.[jp][pn]g")

# Custom CSS for larger text
custom_css = """
    .gradio-container h1 {
        font-size: 2.5em;  /* Increase title font size */
    }
    .gradio-container p {
        font-size: 1.25em;  /* Increase paragraph font size */
    }
    .gradio-container .gr-button {
        font-size: 1.25em;  /* Increase button text size */
    }
    .gradio-container .gr-input, .gradio-container .gr-slider {
        font-size: 1.25em;  /* Increase input and slider text size */
    }
"""

# Set up the Gradio interface with Ocean theme and custom title and favicon
with gr.Blocks(theme=gr.themes.Ocean(), css=custom_css, title="Fish Detector (Grayscale) ONNX Gradio Demo") as interface:
    gr.Markdown(f"<h1 style='text-align: center;'>{app_title}</h1>")
    gr.Markdown(app_description)

    # Layout for inputs and outputs side by side
    with gr.Row():
        image_input = gr.Image(type="pil", label="Upload an Image")
        result_output = gr.Image(type="numpy", label="Detection Results")

    # Confidence slider and buttons below the images
    with gr.Column():
        confidence_slider = gr.Slider(
            minimum=0.0, maximum=1.0, value=0.35,  # Default confidence threshold
            label="Detection Confidence Threshold"
        )

        # Add "Run" and "Clear" buttons side by side
        with gr.Row():
            clear_button = gr.Button("Clear", variant="secondary")
            run_button = gr.Button("Run Detection", variant="primary")  # Make 'Run' prominent

        # Add example images for testing
        gr.Examples(
            examples=examples,
            inputs=image_input,
            outputs=result_output,
            examples_per_page=10,
            label="Sample Images"
        )

    # Link the buttons to their respective functions
    run_button.click(
        fn=predict,
        inputs=[image_input, confidence_slider],
        outputs=result_output
    )

    clear_button.click(
        fn=clear,
        inputs=None,
        outputs=[image_input, result_output]
    )

# Launch the app with custom browser tab title and favicon
interface.launch()