app.py

"""Geo-trax vehicle detector for Hugging Face Spaces.

A Gradio demo for the geo-trax YOLOv8s detector (rfonod/geo-trax). It detects vehicles in
high-altitude, top-down (bird's-eye view) aerial/drone imagery. Two tabs:

  • Image: run detection on a single uploaded image.
  • Short video: run detection frame-by-frame on a short clip (capped for the free CPU tier).

Primary classes (0–3): Car, Bus, Truck, Motorcycle — evaluated, reliable.
Experimental classes (4–5): Pedestrian, Bicycle — trained but poor performance, not evaluated;
  available as opt-in but off by default.

The full video → track → stabilize → georeference pipeline lives in the `geo-trax` package
(https://github.com/rfonod/geo-trax); this Space is a detection-only showcase of the model.
"""

import tempfile
from collections import Counter

import cv2
import gradio as gr
from huggingface_hub import hf_hub_download
from ultralytics import YOLO

# --- Model -----------------------------------------------------------------------------------

# Download the weights once into the HF hub cache (subsequent calls/imports reuse the cached
# file instead of re-downloading), then keep a single shared model instance for all requests.
MODEL_REPO, MODEL_FILE = "rfonod/geo-trax", "geotrax_hbb_yolov8s_1920_v1.pt"
model = YOLO(hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE))

# Four evaluated classes (on by default) + two experimental classes (off by default).
# Pedestrian (4) and bicycle (5) were trained but have poor performance and have not been
# formally evaluated — available as opt-in only.
SUPPORTED = {0: "Car", 1: "Bus", 2: "Truck", 3: "Motorcycle"}
EXPERIMENTAL = {4: "Pedestrian", 5: "Bicycle"}
ALL_SUPPORTED = {**SUPPORTED, **EXPERIMENTAL}
CLASS_CHOICES = list(SUPPORTED.values())            # checked by default
EXPERIMENTAL_CHOICES = list(EXPERIMENTAL.values())  # unchecked by default
ALL_CLASS_CHOICES = list(ALL_SUPPORTED.values())

# Defaults mirror geo-trax's bundled config (geotrax/cfg/default.yaml → ultralytics:).
DEFAULT_CONF, DEFAULT_IOU, MAX_DET = 0.25, 0.7, 1000
IMGSZ_CHOICES = ["640", "960", "1280", "1600", "1920"]  # 1920 = the model's native resolution.

# Slimmer annotations than Ultralytics' size-scaled defaults (which are thick on large frames).
LINE_WIDTH, FONT_SIZE = 2, 16

# Video tab is capped so a run finishes in reasonable time on the free CPU tier.
MAX_VIDEO_FRAMES = 90

# Example assets live in this Space's repo under examples/, but they are stored with Git LFS/Xet,
# so the repo tree (and the Space container) holds a small text *pointer* rather than the image —
# which left gr.Examples thumbnails broken. Reference them via the HF "resolve" URL instead: that
# endpoint always returns the real bytes (LFS resolved server-side). Same approach the Ultralytics
# demo uses with remote example URLs.
EXAMPLES_BASE = "https://huggingface.co/spaces/rfonod/geo-trax/resolve/main/examples"


# --- Helpers ---------------------------------------------------------------------------------


def _class_ids(selected_labels):
    """Map the checkbox labels back to class ids; fall back to the four primary classes if none picked."""
    ids = [cid for cid, name in ALL_SUPPORTED.items() if name in (selected_labels or [])]
    return ids or list(SUPPORTED)


def _count_rows(detected_ids, active_ids):
    """Build a [class, count] table for the active classes only (with a total row)."""
    counts = Counter(int(c) for c in detected_ids)
    rows = [[ALL_SUPPORTED.get(cid, str(cid)), counts.get(cid, 0)] for cid in active_ids]
    rows.append(["Total", sum(counts.values())])
    return rows


# --- Inference -------------------------------------------------------------------------------


def detect_image(image, conf, iou, imgsz, selected_labels, show_labels, show_conf):
    """Detect vehicles in a single image. Returns (annotated RGB image, count table)."""
    if image is None:
        return None, [["Total", 0]]

    active = _class_ids(selected_labels)
    result = model.predict(
        source=image,  # PIL image (RGB); Ultralytics handles channel order correctly.
        imgsz=int(imgsz),
        conf=float(conf),
        iou=float(iou),
        classes=active,
        max_det=MAX_DET,
        verbose=False,
    )[0]

    annotated = result.plot(
        line_width=LINE_WIDTH, font_size=FONT_SIZE, labels=show_labels, conf=show_conf
    )[:, :, ::-1]  # BGR → RGB for display.
    return annotated, _count_rows(result.boxes.cls.tolist(), active)


def detect_video(
    video_path, conf, iou, imgsz, selected_labels, show_labels, show_conf, progress=gr.Progress()
):
    """Detect vehicles frame-by-frame on a short clip (first MAX_VIDEO_FRAMES frames).

    Returns (annotated mp4 path, count table of total detections across processed frames).
    """
    if not video_path:
        return None, [["Total", 0]]

    class_ids = _class_ids(selected_labels)
    cap = cv2.VideoCapture(video_path)
    fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
    total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0)
    n_frames = min(total, MAX_VIDEO_FRAMES) if total > 0 else MAX_VIDEO_FRAMES

    out_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
    writer = None
    agg = []

    for _ in progress.tqdm(range(n_frames), desc="Processing frames"):
        ok, frame = cap.read()
        if not ok:
            break
        result = model.predict(
            source=frame,  # BGR numpy from OpenCV, Ultralytics' expected order.
            imgsz=int(imgsz),
            conf=float(conf),
            iou=float(iou),
            classes=class_ids,
            max_det=MAX_DET,
            verbose=False,
        )[0]
        annotated = result.plot(
            line_width=LINE_WIDTH, font_size=FONT_SIZE, labels=show_labels, conf=show_conf
        )  # BGR
        if writer is None:
            h, w = annotated.shape[:2]
            writer = cv2.VideoWriter(out_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h))
        writer.write(annotated)
        agg.extend(result.boxes.cls.tolist())

    cap.release()
    if writer is not None:
        writer.release()

    out_path = _to_browser_mp4(out_path)
    return out_path, _count_rows(agg, class_ids)


def _to_browser_mp4(path):
    """Best-effort re-encode to H.264/yuv420p so the clip plays inline; fall back to the input."""
    import shutil
    import subprocess

    if shutil.which("ffmpeg") is None:
        return path
    encoded = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
    try:
        subprocess.run(
            ["ffmpeg", "-y", "-i", path, "-vcodec", "libx264", "-pix_fmt", "yuv420p", encoded],
            check=True,
            capture_output=True,
        )
        return encoded
    except Exception:  # noqa: BLE001 (keep the original mp4v file if ffmpeg is unavailable)
        return path


# --- UI --------------------------------------------------------------------------------------

HEADER = """
<h2 align="center">🚗 Geo-trax: Aerial Vehicle Detector</h2>

<div style="display: flex; flex-wrap: wrap; justify-content: center; align-items: center; gap: 6px;">
<a href="https://huggingface.co/rfonod/geo-trax"><img src="https://img.shields.io/badge/%F0%9F%A4%97%20Model-rfonod%2Fgeo--trax-yellow" alt="Model"></a>
<a href="https://github.com/rfonod/geo-trax"><img src="https://img.shields.io/badge/GitHub-geo--trax-blue?logo=github" alt="GitHub"></a>
<a href="https://pypi.org/project/geo-trax/"><img src="https://img.shields.io/pypi/v/geo-trax?label=PyPI&color=blue" alt="PyPI"></a>
<a href="https://doi.org/10.1016/j.trc.2025.105205"><img src="https://img.shields.io/badge/Journal-10.1016%2Fj.trc.2025.105205-blue" alt="Paper"></a>
<a href="https://arxiv.org/abs/2411.02136"><img src="https://img.shields.io/badge/arXiv-2411.02136-b31b1b" alt="arXiv"></a>
<a href="https://youtu.be/gOGivL9FFLk"><img src="https://img.shields.io/badge/YouTube-Demo-red?logo=youtube&logoColor=white" alt="YouTube demo"></a>
</div>

<p align="center"><b>Detect vehicles (car · bus · truck · motorcycle) in high-altitude bird's-eye-view drone imagery</b>, powered by the geo-trax YOLOv8s model. </p>

> **Optimized for high-altitude, top-down (bird's-eye-view) aerial and drone footage.** This is a
> detection-only demo; the full track → stabilize → georeference pipeline lives in the
> [`geo-trax`](https://github.com/rfonod/geo-trax) package.
"""

FOOTER = """
---

### 🎬 Beyond detection: the full Geo-trax pipeline

This Space runs only the **detector**. From raw drone video (plus orthophotos), the full
**[Geo-trax](https://github.com/rfonod/geo-trax)** pipeline extracts **georeferenced vehicle
trajectories**: real-world coordinates, lane and road-section assignment, speeds and
accelerations, and estimated vehicle dimensions.

<p align="center">
<img src="https://raw.githubusercontent.com/rfonod/geo-trax/main/assets/geo-trax_visualization.webp" width="88%" alt="Geo-trax pipeline visualization">
</p>

<p align="center">⭐ <a href="https://github.com/rfonod/geo-trax">Star on GitHub</a> &nbsp;·&nbsp; 📺 <a href="https://youtu.be/gOGivL9FFLk">Watch the 4-min demo</a> &nbsp;·&nbsp; 📦 <a href="https://pypi.org/project/geo-trax/"><code>pip install geo-trax</code></a> &nbsp;·&nbsp; 📄 <a href="https://doi.org/10.1016/j.trc.2025.105205">Read the paper</a></p>
"""

COUNT_HEADERS = ["Class", "Count"]


_EXPERIMENTAL_WARNING = (
    "> ⚠️ **Experimental classes selected (Pedestrian / Bicycle):** These classes were trained "
    "but performance is poor and results have not been formally evaluated. Expect significant "
    "false positives and missed detections. See the "
    "[model card](https://huggingface.co/rfonod/geo-trax#classes-and-detection-performance) "
    "for the full class table and metrics."
)


def _controls(default_imgsz):
    """Shared conf / iou / imgsz / classes / display controls for a tab."""
    conf = gr.Slider(0.0, 1.0, value=DEFAULT_CONF, step=0.01, label="Confidence threshold")
    iou = gr.Slider(0.0, 1.0, value=DEFAULT_IOU, step=0.01, label="IoU threshold (NMS)")
    imgsz = gr.Radio(
        IMGSZ_CHOICES, value=default_imgsz, label="Inference size (px)",
        info="Higher = more accurate but slower; the model is native at 1920.",
    )
    classes = gr.CheckboxGroup(
        ALL_CLASS_CHOICES, value=CLASS_CHOICES, label="Classes",
        info="Pedestrian and Bicycle are experimental — off by default.",
    )
    exp_warning = gr.Markdown(_EXPERIMENTAL_WARNING, visible=False)
    classes.change(
        fn=lambda sel: gr.update(visible=any(c in (sel or []) for c in EXPERIMENTAL_CHOICES)),
        inputs=[classes],
        outputs=[exp_warning],
    )
    with gr.Row():
        show_labels = gr.Checkbox(value=True, label="Show labels")
        show_conf = gr.Checkbox(value=True, label="Show confidence")
    return conf, iou, imgsz, classes, show_labels, show_conf


with gr.Blocks(title="Geo-trax: Aerial Vehicle Detector") as demo:
    # sanitize_html=False keeps the inline `style="display:flex"` on the badge row (Gradio 5.x
    # strips it by default, which makes the block-displayed badge images stack vertically).
    gr.Markdown(HEADER, sanitize_html=False)

    with gr.Tab("Image"):
        with gr.Row():
            with gr.Column():
                # sources excludes "webcam": no live camera input on this tab.
                img_in = gr.Image(type="pil", sources=["upload", "clipboard"], label="Aerial BEV image")
                i_conf, i_iou, i_imgsz, i_classes, i_labels, i_show_conf = _controls("1920")
                img_btn = gr.Button("Detect", variant="primary")
            with gr.Column():
                img_out = gr.Image(label="Detections")
                img_counts = gr.Dataframe(headers=COUNT_HEADERS, label="Counts", interactive=False)
        # Multiple inputs → Gradio renders the examples as a table (thumbnail · conf · IoU · size),
        # like the Ultralytics demo. URLs (not local paths) so the real image bytes are served.
        gr.Examples(
            examples=[
                [f"{EXAMPLES_BASE}/intersection_1.jpg", 0.3, DEFAULT_IOU, "1920"],
                [f"{EXAMPLES_BASE}/intersection_2.jpg", DEFAULT_CONF, DEFAULT_IOU, "1920"],
                [f"{EXAMPLES_BASE}/intersection_3.jpg", 0.3, DEFAULT_IOU, "1920"],
            ],
            inputs=[img_in, i_conf, i_iou, i_imgsz],
            label="Example aerial images — click a row to load it",
            cache_examples=False,
        )
        img_btn.click(
            detect_image,
            inputs=[img_in, i_conf, i_iou, i_imgsz, i_classes, i_labels, i_show_conf],
            outputs=[img_out, img_counts],
        )

    with gr.Tab("Short video"):
        gr.Markdown(
            f"⏱️ Processes up to the **first {MAX_VIDEO_FRAMES} frames**. Frame-by-frame inference "
            "runs on the free **CPU** tier, so it takes a little while."
        )
        with gr.Row():
            with gr.Column():
                vid_in = gr.Video(label="Short aerial BEV clip")
                v_conf, v_iou, v_imgsz, v_classes, v_labels, v_show_conf = _controls("1280")
                vid_btn = gr.Button("Detect", variant="primary")
            with gr.Column():
                vid_out = gr.Video(label="Detections")
                vid_counts = gr.Dataframe(
                    headers=COUNT_HEADERS, label="Total detections (across processed frames)",
                    interactive=False,
                )
        gr.Examples(
            examples=[[f"{EXAMPLES_BASE}/traffic_clip.mp4", DEFAULT_CONF, DEFAULT_IOU, "1280"]],
            inputs=[vid_in, v_conf, v_iou, v_imgsz],
            label="Example clip — click to load",
            cache_examples=False,
        )
        vid_btn.click(
            detect_video,
            inputs=[vid_in, v_conf, v_iou, v_imgsz, v_classes, v_labels, v_show_conf],
            outputs=[vid_out, vid_counts],
        )

    gr.Markdown(FOOTER, sanitize_html=False)


# default_concurrency_limit=1 → only one inference runs at a time, so the single shared model is
# never called by two requests in parallel (important on the free, single-CPU tier). Set at module
# level so it applies whether HF imports `demo` or runs this script.
demo.queue(default_concurrency_limit=1)

if __name__ == "__main__":
    # ssr_mode=False avoids re-importing the app (and reloading the model) in a separate SSR worker.
    demo.launch(ssr_mode=False)