import spaces  # MUST come before torch / any CUDA-touching import
import os
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import tempfile
import gradio as gr
import numpy as np
import PIL.Image as Image
from huggingface_hub import hf_hub_download
from ultralytics import YOLO

# VisDrone class names (10 classes detected by this model)
VISDRONE_CLASSES = [
    "pedestrian", "people", "bicycle", "car", "van",
    "truck", "tricycle", "awning-tricycle", "bus", "motor",
]

# Load model weights at module scope — ZeroGPU intercepts .to("cuda")
weights_path = hf_hub_download(repo_id="dronefreak/visdrone-yolov26n", filename="best.pt")
model = YOLO(weights_path)


@spaces.GPU(duration=30)
def detect(
    image,
    conf_threshold: float = 0.25,
    iou_threshold: float = 0.7,
    show_labels: bool = True,
    show_conf: bool = True,
    imgsz: int = 640,
):
    """Detect objects in an aerial/drone image using YOLOv26n fine-tuned on VisDrone.

    Args:
        image: Input image (PIL or numpy array).
        conf_threshold: Minimum confidence score for detections (0–1).
        iou_threshold: IoU NMS threshold (0–1).
        show_labels: Whether to draw class labels on the output.
        show_conf: Whether to draw confidence scores on the output.
        imgsz: Inference image size in pixels.
    """
    results = model.predict(
        source=image,
        conf=conf_threshold,
        iou=iou_threshold,
        imgsz=imgsz,
        verbose=False,
    )

    # Build detection summary
    detection_counts = {}
    if results[0].boxes is not None and len(results[0].boxes) > 0:
        cls_ids = results[0].boxes.cls.cpu().numpy().astype(int)
        for cid in cls_ids:
            name = VISDRONE_CLASSES[cid] if cid < len(VISDRONE_CLASSES) else str(cid)
            detection_counts[name] = detection_counts.get(name, 0) + 1

    summary_lines = [f"**Total detections:** {sum(detection_counts.values())}"]
    for name, count in sorted(detection_counts.items(), key=lambda x: -x[1]):
        summary_lines.append(f"- {name}: {count}")
    summary = "\n".join(summary_lines)

    # Annotated image
    annotated = results[0].plot(labels=show_labels, conf=show_conf)
    annotated_pil = Image.fromarray(annotated[..., ::-1])

    return annotated_pil, summary


CSS = """
#col-container { max-width: 1200px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks(title="YOLOv26n VisDrone Detection") as demo:
    gr.Markdown(
        """
        # YOLOv26n VisDrone Object Detection

        A lightweight 2.6M-parameter YOLOv26n model fine-tuned on the [VisDrone](https://github.com/VisDrone/VisDrone-Dataset) benchmark for aerial/drone imagery.
        It detects 10 classes: *pedestrian, people, bicycle, car, van, truck, tricycle, awning-tricycle, bus, motor*.

        Upload a drone/aerial image and run detection, or try one of the examples below.
        """
    )

    with gr.Row():
        with gr.Column():
            input_image = gr.Image(type="pil", label="Input Image")
            with gr.Accordion("Detection Settings", open=False):
                conf_slider = gr.Slider(minimum=0.0, maximum=1.0, value=0.25, step=0.01, label="Confidence Threshold")
                iou_slider = gr.Slider(minimum=0.0, maximum=1.0, value=0.7, step=0.01, label="IoU (NMS) Threshold")
                imgsz_radio = gr.Radio(choices=[320, 640, 1024], value=640, label="Inference Image Size")
                labels_checkbox = gr.Checkbox(value=True, label="Show Labels")
                conf_show_checkbox = gr.Checkbox(value=True, label="Show Confidence Scores")
            detect_btn = gr.Button("Detect Objects", variant="primary")

        with gr.Column():
            output_image = gr.Image(type="pil", label="Detection Result")
            detection_summary = gr.Markdown(label="Detection Summary")

    detect_btn.click(
        fn=detect,
        inputs=[input_image, conf_slider, iou_slider, labels_checkbox, conf_show_checkbox, imgsz_radio],
        outputs=[output_image, detection_summary],
    )

    gr.Examples(
        examples=[
            ["examples/0000001_02999_d_0000005.jpg"],
            ["examples/0000002_00005_d_0000014.jpg"],
            ["examples/0000006_00159_d_0000001.jpg"],
        ],
        inputs=[input_image],
        outputs=[output_image, detection_summary],
        fn=detect,
        cache_examples=True,
        cache_mode="lazy",
    )

demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)