"""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 = """
🚗 Geo-trax: Aerial Vehicle Detector
Detect vehicles (car · bus · truck · motorcycle) in high-altitude bird's-eye-view drone imagery, powered by the geo-trax YOLOv8s model.
> **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.
⭐ Star on GitHub · 📺 Watch the 4-min demo · 📦 pip install geo-trax · 📄 Read the paper
"""
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)
