Pre-load RF-DETR models at startup and add progress bar

.gitignore

@@ -1,2 +0,0 @@
-.idea/
-.gradio/

app.py

@@ -1,19 +1,13 @@
+"""Gradio app for the trackers library — run object tracking on uploaded videos."""
+
 from __future__ import annotations
 
-import os
-import sys
+import subprocess
 import tempfile
 from pathlib import Path
 
 import cv2
 import gradio as gr
-import numpy as np
-import supervision as sv
-import torch
-from tqdm import tqdm
-from inference_models import AutoModel
-
-from trackers import ByteTrackTracker, OCSORTTracker, SORTTracker, frames_from_source
 
 MAX_DURATION_SECONDS = 30
 
@@ -28,7 +22,7 @@ MODELS = [
     "rfdetr-seg-large",
 ]
 
-TRACKERS = ["bytetrack", "sort", "ocsort"]
+TRACKERS = ["bytetrack", "sort"]
 
 COCO_CLASSES = [
     "person",
@@ -43,142 +37,18 @@ COCO_CLASSES = [
     "sports ball",
 ]
 
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-
-print(f"Loading {len(MODELS)} models on {DEVICE}...")
-LOADED_MODELS = {}
-for model_id in MODELS:
-    print(f"  Loading {model_id}...")
-    LOADED_MODELS[model_id] = AutoModel.from_pretrained(model_id, device=DEVICE)
-print("All models loaded.")
-
-COLOR_PALETTE = sv.ColorPalette.from_hex(
-    [
-        "#ffff00",
-        "#ff9b00",
-        "#ff8080",
-        "#ff66b2",
-        "#ff66ff",
-        "#b266ff",
-        "#9999ff",
-        "#3399ff",
-        "#66ffff",
-        "#33ff99",
-        "#66ff66",
-        "#99ff00",
-    ]
-)
-
-RESULTS_DIR = "results"
-os.makedirs(RESULTS_DIR, exist_ok=True)
-
-
-def _init_annotators(
-    show_boxes: bool = False,
-    show_masks: bool = False,
-    show_labels: bool = False,
-    show_ids: bool = False,
-    show_confidence: bool = False,
-) -> tuple[list, sv.LabelAnnotator | None]:
-    """Initialize supervision annotators based on display options."""
-    annotators: list = []
-    label_annotator: sv.LabelAnnotator | None = None
-
-    if show_masks:
-        annotators.append(
-            sv.MaskAnnotator(
-                color=COLOR_PALETTE,
-                color_lookup=sv.ColorLookup.TRACK,
-            )
-        )
-
-    if show_boxes:
-        annotators.append(
-            sv.BoxAnnotator(
-                color=COLOR_PALETTE,
-                color_lookup=sv.ColorLookup.TRACK,
-            )
-        )
-
-    if show_labels or show_ids or show_confidence:
-        label_annotator = sv.LabelAnnotator(
-            color=COLOR_PALETTE,
-            text_color=sv.Color.BLACK,
-            text_position=sv.Position.TOP_LEFT,
-            color_lookup=sv.ColorLookup.TRACK,
-        )
-
-    return annotators, label_annotator
-
-
-def _format_labels(
-    detections: sv.Detections,
-    class_names: list[str],
-    *,
-    show_ids: bool = False,
-    show_labels: bool = False,
-    show_confidence: bool = False,
-) -> list[str]:
-    """Generate label strings for each detection."""
-    labels = []
-
-    for i in range(len(detections)):
-        parts = []
-
-        if show_ids and detections.tracker_id is not None:
-            parts.append(f"#{int(detections.tracker_id[i])}")
-
-        if show_labels and detections.class_id is not None:
-            class_id = int(detections.class_id[i])
-            if class_names and 0 <= class_id < len(class_names):
-                parts.append(class_names[class_id])
-            else:
-                parts.append(str(class_id))
-
-        if show_confidence and detections.confidence is not None:
-            parts.append(f"{detections.confidence[i]:.2f}")
-
-        labels.append(" ".join(parts))
-
-    return labels
-
-
 VIDEO_EXAMPLES = [
     [
         "https://storage.googleapis.com/com-roboflow-marketing/supervision/video-examples/bikes-1280x720-1.mp4",
         "rfdetr-small",
-        "ocsort",
+        "bytetrack",
         0.2,
         30,
         0.3,
         3,
         0.1,
         0.6,
-        0.2,
-        3,
         [],
-        "",
-        True,
-        True,
-        False,
-        False,
-        True,
-        False,
-    ],
-    [
-        "https://storage.googleapis.com/com-roboflow-marketing/supervision/video-examples/bikes-1280x720-1.mp4",
-        "rfdetr-small",
-        "ocsort",
-        0.2,
-        30,
-        0.3,
-        3,
-        0.1,
-        0.6,
-        0.2,
-        3,
-        ["person"],
-        "",
         True,
         True,
         False,
@@ -196,10 +66,7 @@ VIDEO_EXAMPLES = [
         3,
         0.3,
         0.6,
-        0.2,
-        3,
         [],
-        "",
         True,
         True,
         False,
@@ -208,25 +75,22 @@ VIDEO_EXAMPLES = [
         True,
     ],
     [
-        "https://storage.googleapis.com/com-roboflow-marketing/supervision/video-examples/apples-1280x720-2.mp4",
-        "rfdetr-nano",
-        "sort",
+        "https://storage.googleapis.com/com-roboflow-marketing/supervision/video-examples/cars-1280x720-1.mp4",
+        "rfdetr-small",
+        "bytetrack",
         0.2,
         30,
         0.3,
         3,
         0.1,
         0.6,
-        0.2,
-        3,
-        [],
-        "",
-        True,
+        ["car"],
         True,
         True,
         False,
         True,
         False,
+        False,
     ],
     [
         "https://storage.googleapis.com/com-roboflow-marketing/supervision/video-examples/jets-1280x720-1.mp4",
@@ -238,10 +102,7 @@ VIDEO_EXAMPLES = [
         3,
         0.1,
         0.6,
-        0.2,
-        3,
         [],
-        "",
         True,
         True,
         False,
@@ -259,73 +120,25 @@ VIDEO_EXAMPLES = [
         3,
         0.1,
         0.6,
-        0.2,
-        3,
         [],
-        "",
         True,
         True,
         False,
         False,
         True,
-        True,
-    ],
-    [
-        "https://storage.googleapis.com/com-roboflow-marketing/supervision/video-examples/jets-1280x720-2.mp4",
-        "rfdetr-seg-small",
-        "bytetrack",
-        0.2,
-        30,
-        0.3,
-        3,
-        0.1,
-        0.6,
-        0.2,
-        3,
-        [],
-        "1",
-        True,
-        True,
-        False,
-        False,
-        True,
-        True,
-    ],
-    [
-        "https://storage.googleapis.com/com-roboflow-marketing/supervision/video-examples/suitcases-1280x720-4.mp4",
-        "rfdetr-small",
-        "sort",
-        0.2,
-        30,
-        0.3,
-        3,
-        0.1,
-        0.6,
-        0.2,
-        3,
-        [],
-        "",
-        True,
-        True,
-        True,
-        False,
-        True,
         False,
     ],
     [
         "https://storage.googleapis.com/com-roboflow-marketing/supervision/video-examples/vehicles-1280x720.mp4",
-        "rfdetr-medium",
-        "ocsort",
+        "rfdetr-small",
+        "bytetrack",
         0.2,
         30,
         0.3,
         3,
         0.1,
         0.6,
-        0.2,
-        3,
         [],
-        "",
         True,
         True,
         True,
@@ -336,190 +149,102 @@ VIDEO_EXAMPLES = [
 ]
 
 
-def _get_video_info(path: str) -> tuple[float, int]:
-    """Return video duration in seconds and frame count using OpenCV."""
-    video_capture = cv2.VideoCapture(path)
-    if not video_capture.isOpened():
+def _get_video_duration(path: str) -> float:
+    """Return video duration in seconds using OpenCV."""
+    cap = cv2.VideoCapture(path)
+    if not cap.isOpened():
         raise gr.Error("Could not open the uploaded video.")
-    frames_per_second = video_capture.get(cv2.CAP_PROP_FPS)
-    frame_count = int(video_capture.get(cv2.CAP_PROP_FRAME_COUNT))
-    video_capture.release()
-    if frames_per_second <= 0:
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    frame_count = cap.get(cv2.CAP_PROP_FRAME_COUNT)
+    cap.release()
+    if fps <= 0:
         raise gr.Error("Could not determine video frame rate.")
-    return frame_count / frames_per_second, frame_count
-
-
-def _resolve_class_filter(
-    classes: list[str] | None,
-    class_names: list[str],
-) -> list[int] | None:
-    """Resolve class names to integer IDs."""
-    if not classes:
-        return None
-
-    name_to_id = {name: i for i, name in enumerate(class_names)}
-    class_filter: list[int] = []
-    for name in classes:
-        if name in name_to_id:
-            class_filter.append(name_to_id[name])
-    return class_filter if class_filter else None
-
-
-def _resolve_track_id_filter(track_ids_arg: str | None) -> list[int] | None:
-    """Resolve a comma-separated string of track IDs to a list of integers.
-
-    Args:
-        track_ids_arg: Comma-separated string (e.g. `"1,3,5"`). `None` or
-            empty string means no filter.
-
-    Returns:
-        List of integer track IDs, or `None` when no valid filter remains.
-    """
-    if not track_ids_arg:
-        return None
-
-    track_ids: list[int] = []
-    for token in track_ids_arg.split(","):
-        token = token.strip()
-        try:
-            track_ids.append(int(token))
-        except ValueError:
-            print(
-                f"Warning: '{token}' is not a valid track ID, skipping.",
-                file=sys.stderr,
-            )
-    return track_ids if track_ids else None
+    return frame_count / fps
 
 
 def track(
     video_path: str,
-    model_id: str,
-    tracker_type: str,
+    model: str,
+    tracker: str,
     confidence: float,
     lost_track_buffer: int,
     track_activation_threshold: float,
     minimum_consecutive_frames: int,
     minimum_iou_threshold: float,
     high_conf_det_threshold: float,
-    direction_consistency_weight: float,
-    delta_t: int,
     classes: list[str] | None = None,
-    track_ids: str = "",
     show_boxes: bool = True,
     show_ids: bool = True,
     show_labels: bool = False,
     show_confidence: bool = False,
     show_trajectories: bool = False,
     show_masks: bool = False,
-    progress=gr.Progress(track_tqdm=True),
 ) -> str:
     """Run tracking on the uploaded video and return the output path."""
     if video_path is None:
         raise gr.Error("Please upload a video.")
 
-    duration, total_frames = _get_video_info(video_path)
+    duration = _get_video_duration(video_path)
     if duration > MAX_DURATION_SECONDS:
         raise gr.Error(
             f"Video is {duration:.1f}s long. "
-            f"Maximum allowed duration is {MAX_DURATION_SECONDS}s. "
-            f"Please use the trim tool in the Input Video player to shorten it."
+            f"Maximum allowed duration is {MAX_DURATION_SECONDS}s."
         )
 
-    detection_model = LOADED_MODELS[model_id]
-    class_names = getattr(detection_model, "class_names", [])
+    tmp_dir = tempfile.mkdtemp()
+    output_path = str(Path(tmp_dir) / "output.mp4")
+
+    cmd = [
+        "trackers",
+        "track",
+        "--source",
+        video_path,
+        "--output",
+        output_path,
+        "--overwrite",
+        "--model",
+        model,
+        "--model.device",
+        "cuda",
+        "--tracker",
+        tracker,
+        "--model.confidence",
+        str(confidence),
+        "--tracker.lost_track_buffer",
+        str(lost_track_buffer),
+        "--tracker.track_activation_threshold",
+        str(track_activation_threshold),
+        "--tracker.minimum_consecutive_frames",
+        str(minimum_consecutive_frames),
+        "--tracker.minimum_iou_threshold",
+        str(minimum_iou_threshold),
+    ]
 
-    selected_class_ids = _resolve_class_filter(classes, class_names)
-    selected_track_ids = _resolve_track_id_filter(track_ids)
+    # ByteTrack extra param
+    if tracker == "bytetrack":
+        cmd += ["--tracker.high_conf_det_threshold", str(high_conf_det_threshold)]
 
-    if tracker_type == "bytetrack":
-        tracker = ByteTrackTracker(
-            lost_track_buffer=lost_track_buffer,
-            track_activation_threshold=track_activation_threshold,
-            minimum_consecutive_frames=minimum_consecutive_frames,
-            minimum_iou_threshold=minimum_iou_threshold,
-            high_conf_det_threshold=high_conf_det_threshold,
-        )
-    elif tracker_type == "ocsort":
-        tracker = OCSORTTracker(
-            lost_track_buffer=lost_track_buffer,
-            minimum_consecutive_frames=minimum_consecutive_frames,
-            minimum_iou_threshold=minimum_iou_threshold,
-            high_conf_det_threshold=high_conf_det_threshold,
-            direction_consistency_weight=direction_consistency_weight,
-            delta_t=delta_t,
-        )
-    else:
-        tracker = SORTTracker(
-            lost_track_buffer=lost_track_buffer,
-            track_activation_threshold=track_activation_threshold,
-            minimum_consecutive_frames=minimum_consecutive_frames,
-            minimum_iou_threshold=minimum_iou_threshold,
-        )
-    tracker.reset()
+    if classes:
+        cmd += ["--classes", ",".join(classes)]
 
-    annotators, label_annotator = _init_annotators(
-        show_boxes=show_boxes,
-        show_masks=show_masks,
-        show_labels=show_labels,
-        show_ids=show_ids,
-        show_confidence=show_confidence,
-    )
-    trace_annotator = None
+    if show_boxes:
+        cmd += ["--show-boxes"]
+    else:
+        cmd += ["--no-boxes"]
+    if show_ids:
+        cmd += ["--show-ids"]
+    if show_labels:
+        cmd += ["--show-labels"]
+    if show_confidence:
+        cmd += ["--show-confidence"]
     if show_trajectories:
-        trace_annotator = sv.TraceAnnotator(
-            color=COLOR_PALETTE,
-            color_lookup=sv.ColorLookup.TRACK,
-        )
-
-    temporary_directory = tempfile.mkdtemp()
-    output_path = str(Path(temporary_directory) / "output.mp4")
-
-    video_info = sv.VideoInfo.from_video_path(video_path)
-
-    frame_generator = frames_from_source(video_path)
-
-    with sv.VideoSink(output_path, video_info=video_info) as sink:
-        for frame_idx, frame in tqdm(
-            frame_generator, total=total_frames, desc="Processing video..."
-        ):
-            predictions = detection_model(frame)
-            if predictions:
-                detections = predictions[0].to_supervision()
-
-                if len(detections) > 0 and detections.confidence is not None:
-                    confidence_mask = detections.confidence >= confidence
-                    detections = detections[confidence_mask]
-
-                if selected_class_ids is not None and len(detections) > 0:
-                    class_mask = np.isin(detections.class_id, selected_class_ids)
-                    detections = detections[class_mask]
-            else:
-                detections = sv.Detections.empty()
-
-            tracked = tracker.update(detections)
-
-            if selected_track_ids is not None and len(tracked) > 0:
-                if tracked.tracker_id is not None:
-                    track_id_mask = np.isin(tracked.tracker_id, selected_track_ids)
-                    tracked = tracked[track_id_mask]
-
-            annotated = frame.copy()
-            if trace_annotator is not None:
-                annotated = trace_annotator.annotate(annotated, tracked)
-            for annotator in annotators:
-                annotated = annotator.annotate(annotated, tracked)
-            if label_annotator is not None:
-                labeled = tracked[tracked.tracker_id != -1]
-                labels = _format_labels(
-                    labeled,
-                    class_names,
-                    show_ids=show_ids,
-                    show_labels=show_labels,
-                    show_confidence=show_confidence,
-                )
-                annotated = label_annotator.annotate(annotated, labeled, labels=labels)
+        cmd += ["--show-trajectories"]
+    if show_masks:
+        cmd += ["--show-masks"]
 
-            sink.write_frame(annotated)
+    result = subprocess.run(cmd, capture_output=True, text=True)  # noqa: S603
+    if result.returncode != 0:
+        raise gr.Error(f"Tracking failed:\n{result.stderr[-500:]}")
 
     return output_path
 
@@ -536,7 +261,7 @@ with gr.Blocks(title="Trackers Playground 🔥") as demo:
         input_video = gr.Video(label="Input Video")
         output_video = gr.Video(label="Tracked Video")
 
-    track_button = gr.Button(value="Track", variant="primary")
+    track_btn = gr.Button(value="Track", variant="primary")
 
     with gr.Row():
         model_dropdown = gr.Dropdown(
@@ -568,16 +293,6 @@ with gr.Blocks(title="Trackers Playground 🔥") as demo:
                     label="Filter Classes",
                     info="Only track selected classes. None selected means all.",
                 )
-                track_id_filter = gr.Textbox(
-                    value="",
-                    label="Filter IDs",
-                    info=(
-                        "Only display tracks with specific track IDs "
-                        "(comma-separated, e.g. 1,3,5). "
-                        "Leave empty for all."
-                    ),
-                    placeholder="e.g. 1,3,5",
-                )
 
             with gr.Column():
                 gr.Markdown("### Tracker")
@@ -587,7 +302,7 @@ with gr.Blocks(title="Trackers Playground 🔥") as demo:
                     value=30,
                     step=1,
                     label="Lost Track Buffer",
-                    info="Frames to keep a lost track before removing it (ByteTrack, SORT, OC-SORT).",
+                    info="Frames to keep a lost track before removing it.",
                 )
                 track_activation_slider = gr.Slider(
                     minimum=0.0,
@@ -595,47 +310,31 @@ with gr.Blocks(title="Trackers Playground 🔥") as demo:
                     value=0.3,
                     step=0.05,
                     label="Track Activation Threshold",
-                    info="Minimum score for a track to be activated (ByteTrack, SORT).",
+                    info="Minimum score for a track to be activated.",
                 )
-                minimum_consecutive_slider = gr.Slider(
+                min_consecutive_slider = gr.Slider(
                     minimum=1,
                     maximum=10,
                     value=2,
                     step=1,
                     label="Minimum Consecutive Frames",
-                    info="Detections needed before a track is confirmed (ByteTrack, SORT, OC-SORT).",
+                    info="Detections needed before a track is confirmed.",
                 )
-                minimum_iou_slider = gr.Slider(
+                min_iou_slider = gr.Slider(
                     minimum=0.0,
                     maximum=1.0,
                     value=0.1,
                     step=0.05,
                     label="Minimum IoU Threshold",
-                    info="Overlap required to match a detection to a track (ByteTrack, SORT, OC-SORT).",
+                    info="Overlap required to match a detection to a track.",
                 )
-                high_confidence_slider = gr.Slider(
+                high_conf_slider = gr.Slider(
                     minimum=0.0,
                     maximum=1.0,
                     value=0.6,
                     step=0.05,
                     label="High Confidence Detection Threshold",
-                    info="Detections above this are matched first (ByteTrack / OC-SORT).",
-                )
-                direction_consistency_slider = gr.Slider(
-                    minimum=0.0,
-                    maximum=1.0,
-                    value=0.2,
-                    step=0.05,
-                    label="Direction Consistency Weight",
-                    info="Weight for direction consistency in association cost (OC-SORT only).",
-                )
-                delta_t_slider = gr.Slider(
-                    minimum=1,
-                    maximum=10,
-                    value=3,
-                    step=1,
-                    label="Delta T",
-                    info="Past frames for velocity estimation during occlusion (OC-SORT only).",
+                    info="Detections above this are matched first (ByteTrack only).",
                 )
 
             with gr.Column():
@@ -682,13 +381,10 @@ with gr.Blocks(title="Trackers Playground 🔥") as demo:
             confidence_slider,
             lost_track_buffer_slider,
             track_activation_slider,
-            minimum_consecutive_slider,
-            minimum_iou_slider,
-            high_confidence_slider,
-            direction_consistency_slider,
-            delta_t_slider,
+            min_consecutive_slider,
+            min_iou_slider,
+            high_conf_slider,
             class_filter,
-            track_id_filter,
             show_boxes_checkbox,
             show_ids_checkbox,
             show_labels_checkbox,
@@ -699,7 +395,7 @@ with gr.Blocks(title="Trackers Playground 🔥") as demo:
         outputs=output_video,
     )
 
-    track_button.click(
+    track_btn.click(
         fn=track,
         inputs=[
             input_video,
@@ -708,13 +404,10 @@ with gr.Blocks(title="Trackers Playground 🔥") as demo:
             confidence_slider,
             lost_track_buffer_slider,
             track_activation_slider,
-            minimum_consecutive_slider,
-            minimum_iou_slider,
-            high_confidence_slider,
-            direction_consistency_slider,
-            delta_t_slider,
+            min_consecutive_slider,
+            min_iou_slider,
+            high_conf_slider,
             class_filter,
-            track_id_filter,
             show_boxes_checkbox,
             show_ids_checkbox,
             show_labels_checkbox,

requirements.txt

@@ -1,3 +1,3 @@
 gradio>=6.3.0,<6.4.0
-inference-models>=0.19.0
-trackers==2.3.0
+inference-models[onnx-cpu]==0.18.6rc14
+trackers==2.2.0rc1