app.py

"""Gradio app for the trackers library — run object tracking on uploaded videos."""

from __future__ import annotations

import subprocess
import tempfile
from pathlib import Path

import cv2
import gradio as gr

MAX_DURATION_SECONDS = 30

MODELS = [
    "rfdetr-nano",
    "rfdetr-small",
    "rfdetr-medium",
    "rfdetr-large",
    "rfdetr-seg-nano",
    "rfdetr-seg-small",
    "rfdetr-seg-medium",
    "rfdetr-seg-large",
]

TRACKERS = ["bytetrack", "sort"]

COCO_CLASSES = [
    "person",
    "bicycle",
    "car",
    "motorcycle",
    "airplane",
    "bus",
    "truck",
    "cat",
    "dog",
    "sports ball",
]

VIDEO_EXAMPLES = [
    [
        "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,
        [],
        True,
        True,
        False,
        False,
        True,
        False,
    ],
    [
        "https://storage.googleapis.com/com-roboflow-marketing/supervision/video-examples/bikes-1280x720-2.mp4",
        "rfdetr-seg-small",
        "sort",
        0.2,
        30,
        0.3,
        3,
        0.3,
        0.6,
        [],
        True,
        True,
        False,
        False,
        True,
        True,
    ],
    [
        "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,
        ["car"],
        True,
        True,
        False,
        True,
        False,
        False,
    ],
    [
        "https://storage.googleapis.com/com-roboflow-marketing/supervision/video-examples/jets-1280x720-1.mp4",
        "rfdetr-small",
        "bytetrack",
        0.2,
        30,
        0.3,
        3,
        0.1,
        0.6,
        [],
        True,
        True,
        False,
        False,
        False,
        False,
    ],
    [
        "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,
        [],
        True,
        True,
        False,
        False,
        True,
        False,
    ],
    [
        "https://storage.googleapis.com/com-roboflow-marketing/supervision/video-examples/vehicles-1280x720.mp4",
        "rfdetr-small",
        "bytetrack",
        0.2,
        30,
        0.3,
        3,
        0.1,
        0.6,
        [],
        True,
        True,
        True,
        False,
        True,
        False,
    ],
]


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.")
    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 / fps


def track(
    video_path: 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,
    classes: list[str] | None = None,
    show_boxes: bool = True,
    show_ids: bool = True,
    show_labels: bool = False,
    show_confidence: bool = False,
    show_trajectories: bool = False,
    show_masks: bool = False,
) -> 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 = _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."
        )

    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),
    ]

    # ByteTrack extra param
    if tracker == "bytetrack":
        cmd += ["--tracker.high_conf_det_threshold", str(high_conf_det_threshold)]

    if classes:
        cmd += ["--classes", ",".join(classes)]

    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:
        cmd += ["--show-trajectories"]
    if show_masks:
        cmd += ["--show-masks"]

    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


with gr.Blocks(title="Trackers Playground 🔥") as demo:
    gr.Markdown(
        "# Trackers Playground 🔥\n\n"
        "Upload a video, detect COCO objects with "
        "[RF-DETR](https://github.com/roboflow-ai/rf-detr) and track them with "
        "[Trackers](https://github.com/roboflow/trackers)."
    )

    with gr.Row():
        input_video = gr.Video(label="Input Video")
        output_video = gr.Video(label="Tracked Video")

    track_btn = gr.Button(value="Track", variant="primary")

    with gr.Row():
        model_dropdown = gr.Dropdown(
            choices=MODELS,
            value="rfdetr-small",
            label="Detection Model",
        )
        tracker_dropdown = gr.Dropdown(
            choices=TRACKERS,
            value="bytetrack",
            label="Tracker",
        )

    with gr.Accordion("Configuration", open=False):
        with gr.Row():
            with gr.Column():
                gr.Markdown("### Model")
                confidence_slider = gr.Slider(
                    minimum=0.0,
                    maximum=1.0,
                    value=0.2,
                    step=0.05,
                    label="Detection Confidence",
                    info="Minimum score for a detection to be kept.",
                )
                class_filter = gr.CheckboxGroup(
                    choices=COCO_CLASSES,
                    value=[],
                    label="Filter Classes",
                    info="Only track selected classes. None selected means all.",
                )

            with gr.Column():
                gr.Markdown("### Tracker")
                lost_track_buffer_slider = gr.Slider(
                    minimum=1,
                    maximum=120,
                    value=30,
                    step=1,
                    label="Lost Track Buffer",
                    info="Frames to keep a lost track before removing it.",
                )
                track_activation_slider = gr.Slider(
                    minimum=0.0,
                    maximum=1.0,
                    value=0.3,
                    step=0.05,
                    label="Track Activation Threshold",
                    info="Minimum score for a track to be activated.",
                )
                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.",
                )
                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.",
                )
                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 only).",
                )

            with gr.Column():
                gr.Markdown("### Visualization")
                show_boxes_checkbox = gr.Checkbox(
                    value=True,
                    label="Show Boxes",
                    info="Draw bounding boxes around detections.",
                )
                show_ids_checkbox = gr.Checkbox(
                    value=True,
                    label="Show IDs",
                    info="Display track ID for each object.",
                )
                show_labels_checkbox = gr.Checkbox(
                    value=False,
                    label="Show Labels",
                    info="Display class name for each detection.",
                )
                show_confidence_checkbox = gr.Checkbox(
                    value=False,
                    label="Show Confidence",
                    info="Display detection confidence score.",
                )
                show_trajectories_checkbox = gr.Checkbox(
                    value=False,
                    label="Show Trajectories",
                    info="Draw motion path for each tracked object.",
                )
                show_masks_checkbox = gr.Checkbox(
                    value=False,
                    label="Show Masks",
                    info="Draw segmentation masks (seg models only).",
                )

    gr.Examples(
        examples=VIDEO_EXAMPLES,
        fn=track,
        cache_examples=True,
        inputs=[
            input_video,
            model_dropdown,
            tracker_dropdown,
            confidence_slider,
            lost_track_buffer_slider,
            track_activation_slider,
            min_consecutive_slider,
            min_iou_slider,
            high_conf_slider,
            class_filter,
            show_boxes_checkbox,
            show_ids_checkbox,
            show_labels_checkbox,
            show_confidence_checkbox,
            show_trajectories_checkbox,
            show_masks_checkbox,
        ],
        outputs=output_video,
    )

    track_btn.click(
        fn=track,
        inputs=[
            input_video,
            model_dropdown,
            tracker_dropdown,
            confidence_slider,
            lost_track_buffer_slider,
            track_activation_slider,
            min_consecutive_slider,
            min_iou_slider,
            high_conf_slider,
            class_filter,
            show_boxes_checkbox,
            show_ids_checkbox,
            show_labels_checkbox,
            show_confidence_checkbox,
            show_trajectories_checkbox,
            show_masks_checkbox,
        ],
        outputs=output_video,
    )

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