import gradio as gr
import PIL.Image as Image
from ultralytics import YOLO
import spaces
import os
from huggingface_hub import hf_hub_download

# Helper function to download models from Hugging Face
def get_model_path(model_name):
    model_cache_path = hf_hub_download(
        repo_id="atalaydenknalbant/budgerigar_yolo_models", 
        filename=model_name
    )
    return model_cache_path

@spaces.GPU
def yolo_inference(images, model_id, conf_threshold, iou_threshold, max_detection):
    model_path = get_model_path(model_id)  # Download model
    model = YOLO(model_path)
    results = model.predict(
        source=images,
        conf=conf_threshold,
        iou=iou_threshold,
        imgsz=640,
        max_det=max_detection,
        show_labels=True,
        show_conf=True,
    )

    # Process results and convert to PIL Image
    for r in results:
        image_array = r.plot()
        image = Image.fromarray(image_array[..., ::-1])
    return image

# Define Gradio interface
interface = gr.Interface(
    fn=yolo_inference,
    inputs=[
        gr.Image(type="pil", label="Upload Image"),
        gr.Dropdown(
            choices=['budgerigar_yolo11x.pt', 'budgerigar_yolov9e.pt'],
            label="Model Name",
            value="budgerigar_yolo11x.pt",
        ),
        gr.Slider(minimum=0, maximum=1, value=0.25, label="Confidence Threshold"),
        gr.Slider(minimum=0, maximum=1, value=0.45, label="IoU Threshold"),
        gr.Slider(minimum=1, maximum=300, step=1, value=300, label="Max Detection"),
    ],
    outputs=gr.Image(type="pil", label="Annotated Image"),
    cache_examples=True,
    title="Budgerigar Gender Determination",
    description="Pretrained YOLO models for determining budgerigar gender based on cere color variations. Upload image(s) for inference",
    examples=[
        ["Male.png", "budgerigar_yolov9e.pt", 0.25, 0.45, 300],
        ["Female.png", "budgerigar_yolo11x.pt", 0.25, 0.45, 300],
    ],
)
interface.launch()