"""
SAR Dark Ship Detection — Gradio Space (Week 1 demo).

Loads YOLOv11n-OBB weights from HF Hub at startup, runs inference on uploaded SAR images.
"""
import gradio as gr
from huggingface_hub import hf_hub_download
from ultralytics import YOLO
from PIL import Image
import numpy as np

GITHUB_URL    = "https://github.com/tejassnaikk/sar-dark-ship-detection"
MODEL_CARD_URL = "https://huggingface.co/tejassnaikk/rsdd-yolo11n-obb-v1"

# Download weights once at startup (~5.5 MB, cached on Space instance)
_weights = hf_hub_download(repo_id="tejassnaikk/rsdd-yolo11n-obb-v1", filename="best.pt")
model = YOLO(_weights)


def detect_ships(image: Image.Image, conf_threshold: float = 0.25):
    if image is None:
        return None, "Upload a SAR image to begin."

    results = model.predict(image, conf=conf_threshold, imgsz=512, verbose=False)
    annotated_bgr = results[0].plot()   # numpy array, BGR
    annotated_rgb = annotated_bgr[..., ::-1]

    obb = results[0].obb
    n = len(obb) if obb is not None else 0

    if n == 0:
        msg = f"No ships detected above confidence {conf_threshold:.0%}. Try lowering the threshold."
    else:
        msg = (
            f"**{n} ship candidate{'s' if n != 1 else ''} detected** "
            f"(confidence ≥ {conf_threshold:.0%}).\n\n"
            "⚠️ These are model predictions on a research demo — not verified vessel identifications."
        )

    return Image.fromarray(annotated_rgb), msg


with gr.Blocks(title="SAR Dark Ship Detection") as demo:
    gr.Markdown(
        "# 🛰️ SAR Dark Ship Detection — Week 1 Demo\n\n"
        "YOLOv11n-OBB trained on RSDD-SAR (7,000 high-resolution SAR ship images, ~3 m resolution). "
        "Upload a SAR chip to see rotated bounding box predictions.\n\n"
        "**Trained on RSDD-SAR; not yet adapted to Sentinel-1 GRD (Week 3 of the project).**"
    )

    with gr.Row():
        with gr.Column():
            inp  = gr.Image(type="pil", label="Input SAR image")
            conf = gr.Slider(0.05, 0.95, value=0.25, step=0.05, label="Confidence threshold")
            btn  = gr.Button("Detect ships", variant="primary")
        with gr.Column():
            out_img = gr.Image(type="pil", label="Detections")
            out_msg = gr.Markdown(label="Status")

    btn.click(detect_ships, inputs=[inp, conf], outputs=[out_img, out_msg])

    gr.Markdown(
        "---\n"
        f"[Project repo]({GITHUB_URL}) · "
        f"[Model card]({MODEL_CARD_URL}) · "
        "Outputs are *AIS-dark vessel candidates*, not confirmed detections — "
        "research only, not for operational use."
    )

if __name__ == "__main__":
    demo.launch()