added track-id filter to demo (#293)

* added track-id filter to demo

* Fix demo track-id filter bug, add example variants, and apply code review fixes

---------

Co-authored-by: SkalskiP <piotr.skalski92@gmail.com>

Author: AlexBodner

demo/app.py

@@ -4,11 +4,10 @@
 # Licensed under the Apache License, Version 2.0 [see LICENSE for details]
 # ------------------------------------------------------------------------
 
-"""Gradio app for the trackers library — run object tracking on uploaded videos."""
-
 from __future__ import annotations
 
 import os
+import sys
 import tempfile
 from pathlib import Path
 
@@ -162,6 +161,26 @@ def _format_labels(
         0.1,
         0.6,
         [],
+        "",
+        True,
+        True,
+        False,
+        False,
+        True,
+        False,
+    ],
+    [
+        "https://storage.googleapis.com/com-roboflow-marketing/supervision/video-examples/bikes-1280x720-1.mp4",
+        "rfdetr-small",
+        "bytetrack",
+        0.2,
+        30,
+        0.3,
+        3,
+        0.1,
+        0.6,
+        ["person"],
+        "",
         True,
         True,
         False,
@@ -180,6 +199,7 @@ def _format_labels(
         0.3,
         0.6,
         [],
+        "",
         True,
         True,
         False,
@@ -188,21 +208,22 @@ def _format_labels(
         True,
     ],
     [
-        "https://storage.googleapis.com/com-roboflow-marketing/supervision/video-examples/cars-1280x720-1.mp4",
-        "rfdetr-small",
-        "bytetrack",
+        "https://storage.googleapis.com/com-roboflow-marketing/supervision/video-examples/apples-1280x720-2.mp4",
+        "rfdetr-nano",
+        "sort",
         0.2,
         30,
         0.3,
         3,
         0.1,
         0.6,
-        ["car"],
+        [],
+        "",
         True,
         True,
-        False,
         True,
         False,
+        True,
         False,
     ],
     [
@@ -216,6 +237,7 @@ def _format_labels(
         0.1,
         0.6,
         [],
+        "",
         True,
         True,
         False,
@@ -234,11 +256,50 @@ def _format_labels(
         0.1,
         0.6,
         [],
+        "",
         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,
+        [],
+        "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,
+        [],
+        "",
+        True,
+        True,
+        True,
+        False,
+        True,
         False,
     ],
     [
@@ -252,6 +313,7 @@ def _format_labels(
         0.1,
         0.6,
         [],
+        "",
         True,
         True,
         True,
@@ -291,6 +353,32 @@ def _resolve_class_filter(
     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
+
+
 def track(
     video_path: str,
     model_id: str,
@@ -302,6 +390,7 @@ def track(
     minimum_iou_threshold: float,
     high_conf_det_threshold: float,
     classes: list[str] | None = None,
+    track_ids: str = "",
     show_boxes: bool = True,
     show_ids: bool = True,
     show_labels: bool = False,
@@ -318,14 +407,17 @@ def track(
     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"Maximum allowed duration is {MAX_DURATION_SECONDS}s. "
+            f"Please use the trim tool in the Input Video player to shorten it."
         )
 
     detection_model = LOADED_MODELS[model_id]
     class_names = getattr(detection_model, "class_names", [])
 
     class_filter = _resolve_class_filter(classes, class_names)
 
+    track_id_filter = _resolve_track_id_filter(track_ids)
+
     tracker: ByteTrackTracker | SORTTracker
     if tracker_type == "bytetrack":
         tracker = ByteTrackTracker(
@@ -385,6 +477,11 @@ def track(
 
             tracked = tracker.update(detections)
 
+            if track_id_filter is not None and len(tracked) > 0:
+                if tracked.tracker_id is not None:
+                    mask = np.isin(tracked.tracker_id, track_id_filter)
+                    tracked = tracked[mask]
+
             annotated = frame.copy()
             if trace_annotator is not None:
                 annotated = trace_annotator.annotate(annotated, tracked)
@@ -409,9 +506,10 @@ def track(
 with gr.Blocks(title="Trackers Playground 🔥") as demo:
     gr.Markdown(
         "# Trackers Playground 🔥\n\n"
-        "Upload a video, detect COCO objects with "
+        "Upload a video, detect objects with "
         "[RF-DETR](https://github.com/roboflow-ai/rf-detr) and track them with "
-        "[Trackers](https://github.com/roboflow/trackers)."
+        "[Trackers](https://github.com/roboflow/trackers). This demo uses models "
+        "pretrained on 80 COCO classes, but Trackers works with any detection model."
     )
 
     with gr.Row():
@@ -450,6 +548,16 @@ def track(
                     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")
@@ -542,6 +650,7 @@ def track(
             min_iou_slider,
             high_conf_slider,
             class_filter,
+            track_id_filter,
             show_boxes_checkbox,
             show_ids_checkbox,
             show_labels_checkbox,
@@ -565,6 +674,7 @@ def track(
             min_iou_slider,
             high_conf_slider,
             class_filter,
+            track_id_filter,
             show_boxes_checkbox,
             show_ids_checkbox,
             show_labels_checkbox,