diff --git a/.github/workflows/ci-integrations.yml b/.github/workflows/ci-integrations.yml index 1def972fa..a1c060160 100644 --- a/.github/workflows/ci-integrations.yml +++ b/.github/workflows/ci-integrations.yml @@ -27,7 +27,7 @@ jobs: activate-environment: true - name: 🚀 Install Packages - run: uv sync --frozen --group dev + run: uv sync --frozen --group dev --extra reid - name: 🧪 Run Integration Tests diff --git a/README.md b/README.md index 31b984165..544f3f5d3 100644 --- a/README.md +++ b/README.md @@ -22,6 +22,7 @@ Keeping track of objects across video frames is one of those problems that sound - **Benchmarked across four datasets.** MOT17, SportsMOT, SoccerNet, and DanceTrack — at default parameters and after hyperparameter tuning, so you know what to expect before you deploy. - **Tunable out of the box.** Built-in Optuna-based hyperparameter search via `trackers tune` so you can optimize for your specific scene and detector. - **Camera motion compensation.** BoT-SORT handles moving cameras natively, keeping track IDs stable even when the whole frame shifts. +- **Optional appearance ReID.** BoT-SORT can fuse visual embeddings with motion for harder association scenes: install `trackers[reid]` (pulls in the standalone [`reid`](https://github.com/roboflow/re-ID) package) and pass a `reid.ReIDModel` as `reid_model`. ## Install diff --git a/docs/api/reid.md b/docs/api/reid.md new file mode 100644 index 000000000..57ab06507 --- /dev/null +++ b/docs/api/reid.md @@ -0,0 +1,33 @@ +--- +description: Python API reference for the ReID encoder protocol, feature bank, and appearance association utilities in Roboflow Trackers. +--- + +# ReID API + +Requires the optional extra: + +```bash +pip install 'trackers[reid]' +``` + +This page covers the `ReIDEncoder` protocol, `FeatureBank`, and appearance +association helpers in `trackers.core.reid`. For enabling appearance on +BoT-SORT, threshold selection, and benchmark results, see the +[ReID appearance guide](../learn/reid.md). Model loading and gallery evaluation +are in the standalone [`reid`](https://github.com/roboflow/re-ID) package. + +## ReIDEncoder + +::: trackers.core.reid.encoder.ReIDEncoder + +## FeatureBank + +::: trackers.core.reid.feature_bank.FeatureBank + +## appearance_similarity + +::: trackers.core.reid.appearance.appearance_similarity + +## extract_detection_embeddings + +::: trackers.core.reid.appearance.extract_detection_embeddings diff --git a/docs/assets/reid/mot17-fastreid-appearance-distances.png b/docs/assets/reid/mot17-fastreid-appearance-distances.png new file mode 100644 index 000000000..f92a71a76 Binary files /dev/null and b/docs/assets/reid/mot17-fastreid-appearance-distances.png differ diff --git a/docs/assets/reid/soccernet-osnet-appearance-distances.png b/docs/assets/reid/soccernet-osnet-appearance-distances.png new file mode 100644 index 000000000..880be1427 Binary files /dev/null and b/docs/assets/reid/soccernet-osnet-appearance-distances.png differ diff --git a/docs/learn/install.md b/docs/learn/install.md index 4bbff27d0..c5ded5d0e 100644 --- a/docs/learn/install.md +++ b/docs/learn/install.md @@ -70,6 +70,29 @@ The `detection` extra installs `inference-models`, enabling the CLI to run detec uv pip install "trackers[detection]" ``` +### ReID (BoT-SORT appearance) + +The `reid` extra installs PyTorch, timm, Hugging Face Hub, safetensors, Pillow, +and gdown for ReID model loading (OSNet, FastReID SBS, and `timm:` backbones) +and BoT-SORT appearance association. + +=== "pip" + + ```bash + pip install "trackers[reid]" + ``` + +=== "uv" + + ```bash + uv pip install "trackers[reid]" + ``` + +Use via `from reid import ReIDModel` and +`BoTSORTTracker(reid_model=...)`, or via CLI flags such as +`--tracker.reid.enable` and `--tracker.reid.architecture` on `trackers track` command +(BoT-SORT only). + !!! tip "GPU Acceleration" For GPU support, ensure PyTorch is installed with CUDA or MPS. diff --git a/docs/learn/reid.md b/docs/learn/reid.md new file mode 100644 index 000000000..ed2aa9012 --- /dev/null +++ b/docs/learn/reid.md @@ -0,0 +1,125 @@ +--- +description: Use ReID appearance association with BoT-SORT in Roboflow Trackers, from model loading to appearance threshold selection, with MOT17 and SoccerNet results. +--- + +# ReID Appearance + +BoT-SORT can fuse appearance embeddings with IoU during association. Embeddings +come from a model in the standalone [`reid`](https://github.com/roboflow/re-ID) +package; the association helpers are in `trackers.core.reid` (see the +[ReID API](../api/reid.md)). + +**What you'll learn:** + +- How to enable appearance association on BoT-SORT +- Which parameters control the appearance gate +- How to pick `appearance_threshold` for your encoder and domain +- What ReID changes on MOT17 and SoccerNet + +--- + +## Install + +The `reid` extra installs PyTorch, timm, Hugging Face Hub, safetensors, Pillow, +and gdown for ReID model loading. + +=== "pip" + ```bash + pip install "trackers[reid]" + ``` + +=== "uv" + ```bash + uv pip install "trackers[reid]" + ``` + +--- + +## Quickstart + +```python +from reid import ReIDModel + +from trackers import BoTSORTTracker + +reid_model = ReIDModel.from_pretrained("fastreid_mot17_sbs50") +tracker = BoTSORTTracker(reid_model=reid_model, appearance_threshold=0.2) +``` + +For the model catalog and fine-tuning, see the +[`reid` training guide](https://github.com/roboflow/re-ID/blob/main/docs/learn/train.md). + +--- + +## Parameters + +| Parameter | Default | Purpose | +| ---------------------- | ------- | ---------------------------------------------------------------------------------------------------- | +| `reid_ema_alpha` | 0.9 | EMA momentum for a track's appearance feature. | +| `appearance_threshold` | 0.25 | Max `d_app` for an appearance match (BoT-SORT paper default; the MOT17 setup below uses 0.2). | +| `proximity_threshold` | 0.5 | IoU gate before appearance (`IoU ≥ 1 - proximity_threshold`), from true IoU even with GIoU/DIoU/CIoU. | + +--- + +## Choosing an appearance threshold + +BoT-SORT rejects an appearance match when `d_app = 0.5 * (1 - cos_sim)` exceeds +`appearance_threshold` (paper default 0.25). Pick θ on a labeled split with the +encoder you will track with: + +1. Embed GT crops. +2. Histogram `d_app` for same-ID vs different-ID pairs. +3. Choose θ so most same-ID pairs fall below it and most different-ID pairs fall + above it. + +**MOT17 val, `fastreid_mot17_sbs50`.** Same-ID peaks near 0, different-ID near +0.4. θ=0.2 keeps ~88% of same-ID pairs and ~1% of different-ID pairs +([MOT17 re-ID study](https://www-sop.inria.fr/members/Francois.Bremond/Postscript/Tomasz__SCCAI_2025.pdf) +Table 8; stricter than the paper default 0.25). + +![FastReID MOT17 SBS on MOT17 val GT](../assets/reid/mot17-fastreid-appearance-distances.png) + +**SoccerNet test, `osnet_x1_0_msmt17_combineall`.** Same-ID and different-ID +overlap heavily (similar kits). θ=0.2 passes ~50% of different-ID pairs and +stays flat vs CMC-only; θ=0.1 rejects too many same-ID pairs and costs HOTA and +IDF1 (see the SoccerNet table below). + +![OSNet MSMT17 on SoccerNet test GT](../assets/reid/soccernet-osnet-appearance-distances.png) + +--- + +## BoT-SORT with and without ReID + +### MOT17 test + +YOLOX detections, CMC on, Codabench MOT17 test (same protocol as the +[tracker comparison](../trackers/comparison.md) default table). ReID: +`fastreid_mot17_sbs50`, `appearance_threshold=0.2` +([MOT17 re-ID study](https://www-sop.inria.fr/members/Francois.Bremond/Postscript/Tomasz__SCCAI_2025.pdf) +Table 8). + +| Config | HOTA | IDF1 | MOTA | +| :-------------- | :----: | :----: | :----: | +| BoT-SORT | 63.7 | 78.7 | **79.2** | +| BoT-SORT + ReID | **63.9** | **79.2** | **79.2** | + +### MOT17 val-half + +YOLOX detections, CMC on, MOT17 val-half split, same encoder and threshold. + +| Config | HOTA | MOTA | IDF1 | +| :-------------- | :----: | :----: | :----: | +| BoT-SORT | 68.9 | 78.3 | 81.2 | +| BoT-SORT + ReID | **69.1** | **78.4** | **81.9** | + +### SoccerNet test (OSNet MSMT17) + +Oracle detections, CMC on, SoccerNet-tracking test (same protocol as the +[tracker comparison](../trackers/comparison.md) default table). ReID: +`osnet_x1_0_msmt17_combineall` (MSMT17 pretrained). + +| Config | HOTA | IDF1 | MOTA | +| :---------------------------------- | :----: | :----: | :----: | +| BoT-SORT | 84.5 | 79.3 | **96.6** | +| BoT-SORT + OSNet MSMT17 (θ=0.2) | **84.6** | **79.4** | **96.6** | +| BoT-SORT + OSNet MSMT17 (θ=0.1) | 82.9 | 77.7 | 96.5 | diff --git a/docs/trackers/botsort.md b/docs/trackers/botsort.md index 6638114bd..f2e61e046 100644 --- a/docs/trackers/botsort.md +++ b/docs/trackers/botsort.md @@ -50,6 +50,12 @@ BoT-SORT keeps the same tracking-by-detection backbone as [ByteTrack](bytetrack. | `high_conf_det_threshold` | Confidence split between stage-1 and stage-2 detections. | 0.5-0.7 common. Higher shifts more detections to recovery stage; lower gives stage-1 broader coverage. | | `enable_cmc` | Enables camera motion compensation before association. | Keep enabled for moving-camera footage (sports, drone, handheld). Disable mainly for static cameras if you need maximal speed. | +## ReID appearance (optional) + +BoT-SORT can fuse appearance embeddings with IoU during association via an +optional `reid_model`. Install, usage, parameters, and MOT17 with/without ReID +scores are on the [ReID appearance](../learn/reid.md) page. + ## Run on video, webcam, or RTSP stream These examples use `opencv-python` for decoding and display. Replace ``, ``, and `` with your inputs. `` is usually 0 for the default camera. diff --git a/docs/trackers/comparison.md b/docs/trackers/comparison.md index 729480d5a..78405a119 100644 --- a/docs/trackers/comparison.md +++ b/docs/trackers/comparison.md @@ -28,6 +28,9 @@ Pedestrian tracking with crowded scenes and frequent occlusions. Strongly tests Parameters were tuned on the validation set. Results are reported on the test set via Codabench submission. Detections come from a YOLOX model. + BoT-SORT rows are CMC without appearance; for CMC + FastReID on MOT17 and + OSNet MSMT17 on SoccerNet see + [BoT-SORT with and without ReID](../learn/reid.md#bot-sort-with-and-without-reid). === "Default" diff --git a/mkdocs.yml b/mkdocs.yml index f1ebd59f1..cbad042ae 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -104,6 +104,7 @@ plugins: - mkdocstrings: handlers: python: + paths: [src] options: # Controls whether to show symbol types in the table of contents show_symbol_type_toc: true @@ -135,6 +136,7 @@ nav: - Detection Quality Matters: learn/detection-quality.md - State Estimators: learn/state-estimators.md - IoU Variants: learn/iou.md + - ReID Appearance: learn/reid.md - About: learn/about.md - Trackers: - Comparison: trackers/comparison.md @@ -147,6 +149,7 @@ nav: - API Reference: - Trackers: api/trackers.md - Motion: api/motion.md + - ReID: api/reid.md - Datasets: api/datasets.md - Evals: api/evals.md - I/O: api/io.md diff --git a/notebooks/eval_trackers_reid.ipynb b/notebooks/eval_trackers_reid.ipynb new file mode 100644 index 000000000..fb829f9cc --- /dev/null +++ b/notebooks/eval_trackers_reid.ipynb @@ -0,0 +1,952 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "2d522414", + "metadata": {}, + "source": [ + "# Tracker ReID evaluation on MOT17 val\n", + "\n", + "Compare BoT-SORT with and without a **reid** encoder on the MOT17 val-half split (YOLOX detections, TrackEval metrics).\n", + "\n", + "| Config | Tracker | CMC | ReID | Fusion |\n", + "|---|---|---|---|---|\n", + "| Baseline | BoT-SORT | yes (sparseOptFlow) | no | geometry + CMC |\n", + "| + ReID | BoT-SORT | yes | yes | min-cost (x5) |\n", + "\n", + "**Defaults:** `fastreid_mot17_sbs50`, `REID_APPEARANCE_THRESHOLD=0.2` (MOT17 re-ID study Table 8).\n", + "\n", + "**Data:** `trackers download mot17 --split val --asset annotations,frames` + [YOLOX val detections](https://drive.google.com/file/d/1BuXtPWf8QbPU_y1i2xY2IbTE-rj3l6qT).\n", + "\n", + "**Outputs:** `trackers_reid_outputs/` - set `RERUN[name]=False` to reuse cached preds.\n", + "\n", + "Section 8 builds a side-by-side baseline vs +ReID video for the sequence with the largest \u0394HOTA (Colab auto-download).\n", + "\n", + "> **Runtime:** T4 GPU (`Runtime -> Change runtime type`).\n" + ] + }, + { + "cell_type": "markdown", + "id": "7bec6c65", + "metadata": {}, + "source": [ + "## 1. Setup\n", + "\n", + "Install **trackers** with the ReID extra (pulls in the standalone `reid` package),\n", + "plus notebook-only deps (`matplotlib`, `scikit-learn`, `gdown`).\n", + "\n", + "**Local:** skip the install cell if your venv already has `trackers[reid]`\n", + "(`pip install 'trackers[reid]'` or `uv sync --extra reid`).\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6bc2b8d8", + "metadata": {}, + "outputs": [], + "source": [ + "import subprocess\n", + "import sys\n", + "\n", + "try:\n", + " import google.colab # noqa: F401\n", + "\n", + " IN_COLAB_INSTALL = True\n", + "except ImportError:\n", + " IN_COLAB_INSTALL = False\n", + "\n", + "if IN_COLAB_INSTALL:\n", + " cmds = [\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"-q\", \"--upgrade\", \"pip\"],\n", + " [\n", + " sys.executable,\n", + " \"-m\",\n", + " \"pip\",\n", + " \"install\",\n", + " \"-q\",\n", + " \"trackers[reid]\",\n", + " \"matplotlib\",\n", + " \"scikit-learn\",\n", + " \"gdown\",\n", + " ],\n", + " ]\n", + " for cmd in cmds:\n", + " subprocess.run(cmd, check=True) # noqa: S603\n", + " print(\"Installed trackers[reid] and notebook deps.\")\n", + "else:\n", + " print(\"Local kernel: skipping install.\")\n", + " print(\"Ensure trackers[reid] is importable (pip install 'trackers[reid]' or uv sync --extra reid).\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6c2e60ad", + "metadata": {}, + "outputs": [], + "source": [ + "import shutil\n", + "import subprocess\n", + "import sys\n", + "import warnings\n", + "import zipfile\n", + "from pathlib import Path\n", + "\n", + "import cv2\n", + "import gdown\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "import supervision as sv\n", + "import torch\n", + "from IPython.display import Video\n", + "from IPython.display import display as ipy_display\n", + "from reid import DEFAULT_MODEL, FASTREID_MOT17_SBS50, ReIDModel\n", + "from sklearn.decomposition import PCA\n", + "\n", + "from trackers import BoTSORTTracker\n", + "from trackers.eval import evaluate_mot_sequences\n", + "from trackers.eval.box import box_iou\n", + "from trackers.eval.results import BenchmarkResult\n", + "from trackers.io.frames import load_mot_frame_image\n", + "from trackers.io.mot import _MOTOutput, load_mot_file\n", + "\n", + "warnings.filterwarnings(\"ignore\")\n", + "\n", + "try:\n", + " from google.colab import files\n", + "\n", + " IN_COLAB = True\n", + " REPO_ROOT = Path(\"/content\")\n", + "except ImportError:\n", + " files = None\n", + " IN_COLAB = False\n", + " REPO_ROOT = Path(\"..\").resolve()\n", + "\n", + "VAL_SEQUENCES = [\n", + " \"MOT17-02-FRCNN\",\n", + " \"MOT17-04-FRCNN\",\n", + " \"MOT17-05-FRCNN\",\n", + " \"MOT17-09-FRCNN\",\n", + " \"MOT17-10-FRCNN\",\n", + " \"MOT17-11-FRCNN\",\n", + " \"MOT17-13-FRCNN\",\n", + "]\n", + "\n", + "device = torch.cuda.get_device_name(0) if torch.cuda.is_available() else \"cpu\"\n", + "print(f\"PyTorch {torch.__version__} | CUDA {torch.cuda.is_available()} | {device}\")" + ] + }, + { + "cell_type": "markdown", + "id": "a54bb5ed", + "metadata": {}, + "source": [ + "## 2. ReID model\n", + "\n", + "| `REID_ENCODER` | Training | Input |\n", + "|---|---|---|\n", + "| `fastreid_mot17_sbs50` (default) | MOT17 train-half | 384\u00d7128 |\n", + "| `osnet_msmt17` | MSMT17 combineall | 256\u00d7128 |\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bbc892d6", + "metadata": {}, + "outputs": [], + "source": [ + "REID_ENCODER = \"fastreid_mot17_sbs50\"\n", + "REID_APPEARANCE_THRESHOLD = 0.2 # MOT17 re-ID study Table 8; BoT-SORT paper default 0.25\n", + "\n", + "if REID_ENCODER == FASTREID_MOT17_SBS50:\n", + " reid_model = ReIDModel.from_pretrained(FASTREID_MOT17_SBS50)\n", + "elif REID_ENCODER in (DEFAULT_MODEL, \"osnet_msmt17\"):\n", + " reid_model = ReIDModel.from_pretrained()\n", + "else:\n", + " raise ValueError(f\"Unknown REID_ENCODER: {REID_ENCODER!r}\")\n", + "\n", + "print(f\"Encoder: {REID_ENCODER} | \u03b8_emb: {REID_APPEARANCE_THRESHOLD}\")\n", + "print(reid_model.preprocessing.describe())" + ] + }, + { + "cell_type": "markdown", + "id": "29afb2e0", + "metadata": {}, + "source": [ + "## 3. Download data\n", + "\n", + "MOT17 val GT + frames via `trackers download`. YOLOX val detections via gdown\n", + "(BoT-SORT / ByteTrack eval protocol). YOLOX frame IDs are remapped to 1...N.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5ea423a9", + "metadata": {}, + "outputs": [], + "source": [ + "FORCE_DOWNLOAD = False\n", + "\n", + "MOT17_VAL = REPO_ROOT / \"mot17\" / \"val\"\n", + "YOLOX_DIR = REPO_ROOT / \"MOT17_yolox_dets\"\n", + "YOLOX_VAL_DIR = YOLOX_DIR / \"val\"\n", + "YOLOX_ZIP = YOLOX_DIR / \"yolox_detections_MOT17.zip\"\n", + "YOLOX_GDRIVE_ID = \"1BuXtPWf8QbPU_y1i2xY2IbTE-rj3l6qT\"\n", + "OUTPUT_ROOT = REPO_ROOT / \"trackers_reid_outputs\"\n", + "OUTPUT_ROOT.mkdir(parents=True, exist_ok=True)\n", + "\n", + "\n", + "def yolox_det_path(seq: str) -> Path:\n", + " return YOLOX_VAL_DIR / f\"{seq.replace('-FRCNN', '')}_val.txt\"\n", + "\n", + "\n", + "def mot17_val_ready() -> bool:\n", + " return all(\n", + " (MOT17_VAL / seq / \"gt\" / \"gt.txt\").is_file() and (MOT17_VAL / seq / \"img1\").is_dir() for seq in VAL_SEQUENCES\n", + " )\n", + "\n", + "\n", + "def yolox_ready() -> bool:\n", + " return YOLOX_VAL_DIR.is_dir() and len(list(YOLOX_VAL_DIR.glob(\"MOT17-*_val.txt\"))) >= len(VAL_SEQUENCES)\n", + "\n", + "\n", + "if FORCE_DOWNLOAD or not mot17_val_ready():\n", + " subprocess.run( # noqa: S603\n", + " [\n", + " sys.executable,\n", + " \"-m\",\n", + " \"trackers.scripts\",\n", + " \"download\",\n", + " \"mot17\",\n", + " \"--split\",\n", + " \"val\",\n", + " \"--asset\",\n", + " \"annotations,frames\",\n", + " \"-o\",\n", + " str(REPO_ROOT),\n", + " ],\n", + " check=True,\n", + " )\n", + "else:\n", + " print(\"MOT17 val already present.\")\n", + "\n", + "if FORCE_DOWNLOAD or not yolox_ready():\n", + " YOLOX_DIR.mkdir(parents=True, exist_ok=True)\n", + " print(\"Downloading YOLOX val detections...\")\n", + " gdown.download(id=YOLOX_GDRIVE_ID, output=str(YOLOX_ZIP), quiet=False)\n", + " with zipfile.ZipFile(YOLOX_ZIP) as zf:\n", + " zf.extractall(YOLOX_DIR)\n", + "else:\n", + " print(\"YOLOX detections already present.\")\n", + "\n", + "SEQUENCE_PATHS: dict[str, dict] = {}\n", + "for seq in VAL_SEQUENCES:\n", + " gt = MOT17_VAL / seq / \"gt\" / \"gt.txt\"\n", + " img = MOT17_VAL / seq / \"img1\"\n", + " det = yolox_det_path(seq)\n", + " if not (gt.is_file() and img.is_dir() and det.is_file()):\n", + " print(f\" skip {seq}: missing gt, img1, or YOLOX det\")\n", + " continue\n", + " n_frames = len(list(img.glob(\"*.jpg\")))\n", + " SEQUENCE_PATHS[seq] = {\"gt\": gt, \"img\": img, \"det\": det, \"n_frames\": n_frames}\n", + " print(f\" {seq}: {n_frames} frames\")\n", + "\n", + "ACTIVE_SEQUENCES = list(SEQUENCE_PATHS)\n", + "if not ACTIVE_SEQUENCES:\n", + " raise RuntimeError(\"No sequences ready - re-run downloads above.\")\n", + "\n", + "SEQMAP_PATH = OUTPUT_ROOT / \"MOT17-val.txt\"\n", + "SEQMAP_PATH.write_text(\"name\\n\" + \"\\n\".join(ACTIVE_SEQUENCES) + \"\\n\")\n", + "print(f\"\\n{len(ACTIVE_SEQUENCES)} sequences -> outputs in {OUTPUT_ROOT}\")" + ] + }, + { + "cell_type": "markdown", + "id": "b57560b2", + "metadata": {}, + "source": [ + "## 4. Tracking helpers\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f38487ea", + "metadata": {}, + "outputs": [], + "source": [ + "RERUN = {\n", + " \"botsort_baseline\": True,\n", + " \"botsort_reid\": True,\n", + "}\n", + "\n", + "\n", + "def _yolox_frame_offset(det_path: Path) -> int:\n", + " min_frame = None\n", + " with det_path.open() as f:\n", + " for line in f:\n", + " parts = line.strip().split(\",\")\n", + " if len(parts) < 6:\n", + " continue\n", + " frame = int(float(parts[0]))\n", + " min_frame = frame if min_frame is None else min(min_frame, frame)\n", + " return (min_frame - 1) if min_frame and min_frame > 1 else 0\n", + "\n", + "\n", + "def load_yolox_dets(det_path: Path) -> dict[int, sv.Detections]:\n", + " offset = _yolox_frame_offset(det_path)\n", + " by_frame: dict[int, list[list[float]]] = {}\n", + " with det_path.open() as f:\n", + " for line in f:\n", + " parts = line.strip().split(\",\")\n", + " if len(parts) < 6:\n", + " continue\n", + " frame = int(float(parts[0])) - offset\n", + " if frame < 1:\n", + " continue\n", + " x1, y1, x2, y2, score = map(float, parts[1:6])\n", + " if score <= 0:\n", + " continue\n", + " by_frame.setdefault(frame, []).append([x1, y1, x2, y2, score])\n", + " return {\n", + " frame: sv.Detections(\n", + " xyxy=np.array(boxes, dtype=np.float32)[:, :4],\n", + " confidence=np.array(boxes, dtype=np.float32)[:, 4],\n", + " )\n", + " for frame, boxes in by_frame.items()\n", + " }\n", + "\n", + "\n", + "def fmt_metrics(result: BenchmarkResult) -> tuple[float, float, float, float, float, int]:\n", + " a = result.aggregate\n", + " return (\n", + " (a.HOTA.HOTA * 100 if a.HOTA else float(\"nan\")),\n", + " (a.HOTA.AssA * 100 if a.HOTA else float(\"nan\")),\n", + " (a.HOTA.DetA * 100 if a.HOTA else float(\"nan\")),\n", + " (a.CLEAR.MOTA * 100 if a.CLEAR else float(\"nan\")),\n", + " (a.Identity.IDF1 * 100 if a.Identity else float(\"nan\")),\n", + " (a.CLEAR.IDSW if a.CLEAR else 0),\n", + " )\n", + "\n", + "\n", + "def print_metrics(label: str, result: BenchmarkResult) -> None:\n", + " hota, _assa, _deta, mota, idf1, idsw = fmt_metrics(result)\n", + " print(f\"{label}: HOTA {hota:6.2f} MOTA {mota:6.2f} IDF1 {idf1:6.2f} IDSW {idsw}\")\n", + "\n", + "\n", + "def run_tracking(name: str, factory, *, use_frames: bool) -> Path:\n", + " pred_dir = OUTPUT_ROOT / name / \"preds\"\n", + " pred_dir.mkdir(parents=True, exist_ok=True)\n", + "\n", + " for seq in ACTIVE_SEQUENCES:\n", + " spec = SEQUENCE_PATHS[seq]\n", + " dets = load_yolox_dets(spec[\"det\"])\n", + " images = sorted(spec[\"img\"].glob(\"*.jpg\"))\n", + " tracker = factory()\n", + "\n", + " with _MOTOutput(pred_dir / f\"{seq}.txt\") as out:\n", + " for frame_idx in range(1, spec[\"n_frames\"] + 1):\n", + " frame = None\n", + " if use_frames and frame_idx <= len(images):\n", + " frame = cv2.imread(str(images[frame_idx - 1]))\n", + " tracked = tracker.update(dets.get(frame_idx, sv.Detections.empty()), frame)\n", + " if tracked.tracker_id is not None:\n", + " tracked = tracked[tracked.tracker_id != -1]\n", + " out.write(frame_idx, tracked)\n", + " print(f\" {seq}: {spec['n_frames']} frames\")\n", + "\n", + " return pred_dir\n", + "\n", + "\n", + "def evaluate(name: str, pred_dir: Path) -> BenchmarkResult:\n", + " result = evaluate_mot_sequences(\n", + " gt_dir=MOT17_VAL,\n", + " tracker_dir=pred_dir,\n", + " seqmap=SEQMAP_PATH,\n", + " metrics=[\"CLEAR\", \"HOTA\", \"Identity\"],\n", + " )\n", + " cache = OUTPUT_ROOT / name / \"eval_results.json\"\n", + " cache.parent.mkdir(parents=True, exist_ok=True)\n", + " result.save(cache)\n", + " return result\n", + "\n", + "\n", + "def load_or_run(name: str, factory, *, use_frames: bool) -> BenchmarkResult:\n", + " pred_dir = OUTPUT_ROOT / name / \"preds\"\n", + " cache = OUTPUT_ROOT / name / \"eval_results.json\"\n", + " preds_ok = pred_dir.exists() and all((pred_dir / f\"{s}.txt\").exists() for s in ACTIVE_SEQUENCES)\n", + "\n", + " ran = False\n", + " if RERUN.get(name, True) or not preds_ok:\n", + " print(f\"Running {name}...\")\n", + " pred_dir = run_tracking(name, factory, use_frames=use_frames)\n", + " ran = True\n", + " else:\n", + " print(f\"Using cached preds: {pred_dir}\")\n", + "\n", + " if not ran and cache.exists():\n", + " print(f\"Using cached eval: {cache}\")\n", + " return BenchmarkResult.load(cache)\n", + "\n", + " print(f\"Evaluating {name}...\")\n", + " return evaluate(name, pred_dir)\n", + "\n", + "\n", + "def match_dets_to_gt(gt_frame, det_xyxy: np.ndarray, min_iou: float = 0.5) -> np.ndarray:\n", + " if len(det_xyxy) == 0:\n", + " return np.array([], dtype=np.int64)\n", + " gt_xyxy = sv.xywh_to_xyxy(gt_frame.boxes)\n", + " keep = (gt_frame.confidences > 0) & (gt_frame.classes == 1)\n", + " gt_xyxy, gt_ids = gt_xyxy[keep], gt_frame.ids[keep]\n", + " if len(gt_xyxy) == 0:\n", + " return np.full(len(det_xyxy), -1, dtype=np.int64)\n", + " ious = box_iou(det_xyxy.astype(np.float64), gt_xyxy.astype(np.float64))\n", + " out = np.full(len(det_xyxy), -1, dtype=np.int64)\n", + " for i in range(len(det_xyxy)):\n", + " j = int(np.argmax(ious[i]))\n", + " if ious[i, j] >= min_iou:\n", + " out[i] = int(gt_ids[j])\n", + " return out" + ] + }, + { + "cell_type": "markdown", + "id": "823b2696", + "metadata": {}, + "source": [ + "## 5. Run trackers\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d09332f1", + "metadata": {}, + "outputs": [], + "source": [ + "EXPERIMENTS = [\n", + " (\n", + " \"botsort_baseline\",\n", + " \"BoT-SORT (baseline)\",\n", + " lambda: BoTSORTTracker(enable_cmc=True),\n", + " True,\n", + " ),\n", + " (\n", + " \"botsort_reid\",\n", + " \"BoT-SORT + ReID\",\n", + " lambda: BoTSORTTracker(\n", + " enable_cmc=True,\n", + " reid_model=reid_model,\n", + " reid_ema_alpha=0.9,\n", + " appearance_threshold=REID_APPEARANCE_THRESHOLD,\n", + " ),\n", + " True,\n", + " ),\n", + "]\n", + "\n", + "results: dict[str, BenchmarkResult] = {}\n", + "for name, label, factory, use_frames in EXPERIMENTS:\n", + " results[name] = load_or_run(name, factory, use_frames=use_frames)\n", + " print_metrics(label, results[name])\n", + " print()\n", + "\n", + "result_baseline = results[\"botsort_baseline\"]\n", + "result_reid = results[\"botsort_reid\"]" + ] + }, + { + "cell_type": "markdown", + "id": "ad28e88f", + "metadata": {}, + "source": [ + "## 6. ReID embedding visualization (optional)\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6612c281", + "metadata": {}, + "outputs": [], + "source": [ + "VIZ_SEQ = \"MOT17-02-FRCNN\"\n", + "VIZ_STRIDE, VIZ_MAX_FRAMES, VIZ_MAX_POINTS, VIZ_MAX_CROPS = 5, 40, 300, 24\n", + "\n", + "spec = SEQUENCE_PATHS[VIZ_SEQ]\n", + "gt_by_frame = load_mot_file(spec[\"gt\"])\n", + "dets_by_frame = load_yolox_dets(spec[\"det\"])\n", + "images = sorted(spec[\"img\"].glob(\"*.jpg\"))\n", + "\n", + "crops, embeddings, gt_ids = [], [], []\n", + "for frame_idx in list(range(1, spec[\"n_frames\"] + 1, VIZ_STRIDE))[:VIZ_MAX_FRAMES]:\n", + " dets = dets_by_frame.get(frame_idx)\n", + " gt = gt_by_frame.get(frame_idx)\n", + " if dets is None or gt is None or len(dets) == 0:\n", + " continue\n", + " dets = dets[dets.confidence >= 0.5]\n", + " if len(dets) == 0:\n", + " continue\n", + " bgr = cv2.imread(str(images[frame_idx - 1]))\n", + " if bgr is None:\n", + " continue\n", + " matched = match_dets_to_gt(gt, dets.xyxy)\n", + " feats = reid_model.extract_features(dets, bgr)\n", + " for i in range(len(dets)):\n", + " if matched[i] < 0:\n", + " continue\n", + " crop = sv.crop_image(bgr, dets.xyxy[i].astype(int))\n", + " if crop.size == 0:\n", + " continue\n", + " crops.append(crop[:, :, ::-1])\n", + " embeddings.append(feats[i])\n", + " gt_ids.append(int(matched[i]))\n", + "\n", + "if not embeddings:\n", + " raise RuntimeError(\"No matched crops - try another sequence or lower confidence threshold\")\n", + "\n", + "emb = np.stack(embeddings)\n", + "labels = np.array(gt_ids)\n", + "if len(emb) > VIZ_MAX_POINTS:\n", + " idx = np.linspace(0, len(emb) - 1, VIZ_MAX_POINTS, dtype=int)\n", + " emb, labels, crops = emb[idx], labels[idx], [crops[i] for i in idx]\n", + "\n", + "coords = PCA(n_components=2, random_state=0).fit_transform(emb)\n", + "unique = np.unique(labels)\n", + "colors = {pid: plt.colormaps[\"tab20\"](i % 20) for i, pid in enumerate(unique)}\n", + "\n", + "fig, (ax_pca, ax_crop) = plt.subplots(1, 2, figsize=(14, 6))\n", + "for pid in unique:\n", + " m = labels == pid\n", + " ax_pca.scatter(coords[m, 0], coords[m, 1], s=28, alpha=0.85, color=colors[pid], label=f\"id {pid}\")\n", + "ax_pca.set(title=f\"{VIZ_SEQ} - PCA by GT id\", xlabel=\"PC1\", ylabel=\"PC2\")\n", + "ax_pca.grid(True, alpha=0.3)\n", + "if len(unique) <= 12:\n", + " ax_pca.legend(fontsize=8)\n", + "\n", + "n_show = min(len(crops), VIZ_MAX_CROPS)\n", + "ncols, nrows = 6, int(np.ceil(n_show / 6))\n", + "mosaic = np.full((nrows * 64, ncols * 32, 3), 255, dtype=np.uint8)\n", + "for k in range(n_show):\n", + " r, c = divmod(k, ncols)\n", + " tile = cv2.resize(crops[k], (32, 64))\n", + " y, x = r * 64, c * 32\n", + " mosaic[y : y + 64, x : x + 32] = tile\n", + " rgb = (np.array(colors[labels[k]])[:3] * 255).astype(np.uint8)\n", + " mosaic[y : y + 2, x : x + 32] = rgb\n", + " mosaic[y + 62 : y + 64, x : x + 32] = rgb\n", + "\n", + "ax_crop.imshow(mosaic)\n", + "ax_crop.set(title=f\"Sample crops ({n_show})\")\n", + "ax_crop.axis(\"off\")\n", + "plt.tight_layout()\n", + "plt.show()\n", + "print(f\"{len(coords)} points, {len(unique)} GT ids\")" + ] + }, + { + "cell_type": "markdown", + "id": "b0ea623e", + "metadata": {}, + "source": [ + "## 7. Results\n", + "\n", + "**7.1-7.2** BoT-SORT vs published references.\n" + ] + }, + { + "cell_type": "markdown", + "id": "43321292", + "metadata": {}, + "source": [ + "### 7.1 BoT-SORT - reference targets\n", + "\n", + "**Primary - [*Does Re-ID Really Help in Multi-Object Tracking?*](https://www-sop.inria.fr/members/Francois.Bremond/Postscript/Tomasz__SCCAI_2025.pdf) (2025). BoT-SORT + YOLOX + MOT17 FastReID, app th=0.2. Combined val scores from **Table 8 (HOTA)** and **Table 13 (IDF1)**; MOTA is not reported for this YOLOX setup.\n", + "\n", + "| Config | HOTA | IDF1 |\n", + "|---|---:|---:|\n", + "| No re-ID | 68.43 | 80.92 |\n", + "| MOT17 FastReID, app th=0.2 | 68.95 | 81.98 |\n", + "| **ReID d (reference)** | **+0.52** | **+1.06** |\n", + "\n", + "**Secondary - [BoT-SORT paper](https://arxiv.org/abs/2206.14651)** (Table 1, MOT17 val):\n", + "\n", + "| Method | HOTA | MOTA | IDF1 |\n", + "|---|---:|---:|---:|\n", + "| BoT-SORT | 69.11 | 78.39 | 81.53 |\n", + "| BoT-SORT + ReID | 69.17 | 78.46 | 82.07 |\n", + "| **ReID d (BoT-SORT paper)** | **+0.06** | **+0.07** | **+0.54** |\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d16f6483", + "metadata": {}, + "outputs": [], + "source": [ + "# MOT17 re-ID study reference - Table 8 (HOTA) + Table 13 (IDF1), COMBINED row.\n", + "# MOTA is not reported for the YOLOX setup in that study.\n", + "REID_STUDY_NO_REID = {\"hota\": 68.428, \"mota\": None, \"idf1\": 80.92}\n", + "REID_STUDY_MOT17_TH02 = {\"hota\": 68.951, \"mota\": None, \"idf1\": 81.984}\n", + "REID_STUDY_REID_DELTA = {k: REID_STUDY_MOT17_TH02[k] - REID_STUDY_NO_REID[k] for k in (\"hota\", \"idf1\")}\n", + "\n", + "# BoT-SORT paper Table 1 (MOT17 val, YOLOX).\n", + "BOTSORT_PAPER = {\"hota\": 69.11, \"mota\": 78.39, \"idf1\": 81.53}\n", + "BOTSORT_PAPER_REID = {\"hota\": 69.17, \"mota\": 78.46, \"idf1\": 82.07}\n", + "BOTSORT_PAPER_REID_DELTA = {k: BOTSORT_PAPER_REID[k] - BOTSORT_PAPER[k] for k in BOTSORT_PAPER}\n", + "\n", + "\n", + "def fmt_ref_metric(value: float | None) -> str:\n", + " return f\"{value:6.2f}\" if value is not None else \" -\"\n", + "\n", + "\n", + "def seq_metrics(result: BenchmarkResult, seq: str) -> tuple[float, float, float, int]:\n", + " s = result.sequences.get(seq)\n", + " if s is None:\n", + " return float(\"nan\"), float(\"nan\"), float(\"nan\"), 0\n", + " return (\n", + " s.HOTA.HOTA * 100 if s.HOTA else float(\"nan\"),\n", + " s.HOTA.AssA * 100 if s.HOTA else float(\"nan\"),\n", + " s.Identity.IDF1 * 100 if s.Identity else float(\"nan\"),\n", + " s.CLEAR.IDSW if s.CLEAR else 0,\n", + " )\n", + "\n", + "\n", + "botsort_rows = [\n", + " (\"BoT-SORT (baseline)\", result_baseline),\n", + " (\"BoT-SORT + ReID\", result_reid),\n", + "]\n", + "\n", + "print(\"BoT-SORT - trackers (aggregate, all val sequences)\")\n", + "print(f\"{'Config':<28} {'HOTA':>6} {'AssA':>6} {'DetA':>6} {'MOTA':>6} {'IDF1':>6} {'IDSW':>5}\")\n", + "print(\"-\" * 72)\n", + "for label, res in botsort_rows:\n", + " hota, assa, deta, mota, idf1, idsw = fmt_metrics(res)\n", + " print(f\"{label:<28} {hota:6.2f} {assa:6.2f} {deta:6.2f} {mota:6.2f} {idf1:6.2f} {idsw:5d}\")\n", + "\n", + "b = fmt_metrics(result_baseline)\n", + "r = fmt_metrics(result_reid)\n", + "print(\n", + " f\"\\nBoT-SORT ReID uplift (trackers): \"\n", + " f\"\u0394HOTA {r[0] - b[0]:+6.2f} \u0394MOTA {r[3] - b[3]:+6.2f} \"\n", + " f\"\u0394IDF1 {r[4] - b[4]:+6.2f} \u0394IDSW {int(r[5] - b[5]):+5d}\"\n", + ")\n", + "\n", + "print(\"\\nBoT-SORT vs MOT17 re-ID study (primary - Table 8 + Table 13)\")\n", + "print(f\"{'':28} {'HOTA':>6} {'MOTA':>6} {'IDF1':>6}\")\n", + "print(\"-\" * 52)\n", + "print(\n", + " f\"{'Reference (no re-ID)':<28} \"\n", + " f\"{REID_STUDY_NO_REID['hota']:6.2f} {fmt_ref_metric(REID_STUDY_NO_REID['mota'])} \"\n", + " f\"{REID_STUDY_NO_REID['idf1']:6.2f}\"\n", + ")\n", + "print(\n", + " f\"{'trackers (baseline)':<28} {b[0]:6.2f} {b[3]:6.2f} {b[4]:6.2f} \"\n", + " f\" d {b[0] - REID_STUDY_NO_REID['hota']:+5.2f} \"\n", + " f\"{'-':>6} {b[4] - REID_STUDY_NO_REID['idf1']:+5.2f}\"\n", + ")\n", + "print(\n", + " f\"{'Reference (MOT17 th=0.2)':<28} \"\n", + " f\"{REID_STUDY_MOT17_TH02['hota']:6.2f} {fmt_ref_metric(REID_STUDY_MOT17_TH02['mota'])} \"\n", + " f\"{REID_STUDY_MOT17_TH02['idf1']:6.2f}\"\n", + ")\n", + "print(\n", + " f\"{'trackers (+ ReID)':<28} {r[0]:6.2f} {r[3]:6.2f} {r[4]:6.2f} \"\n", + " f\" d {r[0] - REID_STUDY_MOT17_TH02['hota']:+5.2f} \"\n", + " f\"{'-':>6} {r[4] - REID_STUDY_MOT17_TH02['idf1']:+5.2f}\"\n", + ")\n", + "print(\n", + " f\"\\nReID uplift vs reference study\\n\"\n", + " f\" \u0394HOTA trackers {r[0] - b[0]:+6.2f} reference {REID_STUDY_REID_DELTA['hota']:+6.2f} \"\n", + " f\"gap {(r[0] - b[0]) - REID_STUDY_REID_DELTA['hota']:+6.2f}\\n\"\n", + " f\" \u0394MOTA trackers {r[3] - b[3]:+6.2f} reference -\\n\"\n", + " f\" \u0394IDF1 trackers {r[4] - b[4]:+6.2f} reference {REID_STUDY_REID_DELTA['idf1']:+6.2f} \"\n", + " f\"gap {(r[4] - b[4]) - REID_STUDY_REID_DELTA['idf1']:+6.2f}\"\n", + ")\n", + "\n", + "print(\"\\nBoT-SORT vs BoT-SORT paper Table 1 (secondary)\")\n", + "print(f\"{'':28} {'HOTA':>6} {'MOTA':>6} {'IDF1':>6}\")\n", + "print(\"-\" * 52)\n", + "print(\n", + " f\"{'BoT-SORT paper':<28} {BOTSORT_PAPER['hota']:6.2f} {BOTSORT_PAPER['mota']:6.2f} {BOTSORT_PAPER['idf1']:6.2f}\"\n", + ")\n", + "print(\n", + " f\"{'trackers (baseline)':<28} {b[0]:6.2f} {b[3]:6.2f} {b[4]:6.2f} \"\n", + " f\" d {b[0] - BOTSORT_PAPER['hota']:+5.2f} \"\n", + " f\"{b[3] - BOTSORT_PAPER['mota']:+5.2f} {b[4] - BOTSORT_PAPER['idf1']:+5.2f}\"\n", + ")\n", + "print(\n", + " f\"{'BoT-SORT paper + ReID':<28} {BOTSORT_PAPER_REID['hota']:6.2f} \"\n", + " f\"{BOTSORT_PAPER_REID['mota']:6.2f} {BOTSORT_PAPER_REID['idf1']:6.2f}\"\n", + ")\n", + "print(\n", + " f\"{'trackers (+ ReID)':<28} {r[0]:6.2f} {r[3]:6.2f} {r[4]:6.2f} \"\n", + " f\" d {r[0] - BOTSORT_PAPER_REID['hota']:+5.2f} \"\n", + " f\"{r[3] - BOTSORT_PAPER_REID['mota']:+5.2f} {r[4] - BOTSORT_PAPER_REID['idf1']:+5.2f}\"\n", + ")\n", + "print(\n", + " f\"\\nReID uplift vs BoT-SORT paper\\n\"\n", + " f\" \u0394HOTA trackers {r[0] - b[0]:+6.2f} BoT-SORT paper {BOTSORT_PAPER_REID_DELTA['hota']:+6.2f} \"\n", + " f\"gap {(r[0] - b[0]) - BOTSORT_PAPER_REID_DELTA['hota']:+6.2f}\\n\"\n", + " f\" \u0394MOTA trackers {r[3] - b[3]:+6.2f} BoT-SORT paper {BOTSORT_PAPER_REID_DELTA['mota']:+6.2f} \"\n", + " f\"gap {(r[3] - b[3]) - BOTSORT_PAPER_REID_DELTA['mota']:+6.2f}\\n\"\n", + " f\" \u0394IDF1 trackers {r[4] - b[4]:+6.2f} BoT-SORT paper {BOTSORT_PAPER_REID_DELTA['idf1']:+6.2f} \"\n", + " f\"gap {(r[4] - b[4]) - BOTSORT_PAPER_REID_DELTA['idf1']:+6.2f}\"\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "8ee1ac84", + "metadata": {}, + "source": [ + "### 7.2 BoT-SORT - per-sequence vs reference (Table 8 HOTA + Table 13 IDF1)\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d448e555", + "metadata": {}, + "outputs": [], + "source": [ + "REID_STUDY_PER_SEQ = {\n", + " \"MOT17-02\": {\n", + " \"no_reid\": {\"hota\": 47.131, \"idf1\": 56.968},\n", + " \"mot17_th02\": {\"hota\": 49.304, \"idf1\": 60.0},\n", + " },\n", + " \"MOT17-04\": {\n", + " \"no_reid\": {\"hota\": 78.976, \"idf1\": 91.021},\n", + " \"mot17_th02\": {\"hota\": 79.046, \"idf1\": 90.864},\n", + " },\n", + " \"MOT17-05\": {\n", + " \"no_reid\": {\"hota\": 60.078, \"idf1\": 75.124},\n", + " \"mot17_th02\": {\"hota\": 61.469, \"idf1\": 77.969},\n", + " },\n", + " \"MOT17-09\": {\n", + " \"no_reid\": {\"hota\": 67.941, \"idf1\": 79.985},\n", + " \"mot17_th02\": {\"hota\": 65.878, \"idf1\": 78.832},\n", + " },\n", + " \"MOT17-10\": {\n", + " \"no_reid\": {\"hota\": 57.204, \"idf1\": 76.157},\n", + " \"mot17_th02\": {\"hota\": 59.565, \"idf1\": 81.087},\n", + " },\n", + " \"MOT17-11\": {\n", + " \"no_reid\": {\"hota\": 66.697, \"idf1\": 77.326},\n", + " \"mot17_th02\": {\"hota\": 66.699, \"idf1\": 77.326},\n", + " },\n", + " \"MOT17-13\": {\n", + " \"no_reid\": {\"hota\": 69.833, \"idf1\": 89.533},\n", + " \"mot17_th02\": {\"hota\": 69.791, \"idf1\": 89.431},\n", + " },\n", + "}\n", + "\n", + "\n", + "def ref_seq_key(seq: str) -> str:\n", + " parts = seq.split(\"-\")\n", + " return f\"{parts[0]}-{parts[1]}\"\n", + "\n", + "\n", + "for seq in ACTIVE_SEQUENCES:\n", + " key = ref_seq_key(seq)\n", + " ref = REID_STUDY_PER_SEQ.get(key, {})\n", + " print(seq)\n", + " print(f\" {'Config':<28} {'HOTA':>6} {'IDF1':>6} {'IDSW':>5} {'Ref H':>6} {'dH':>6} {'Ref I':>6} {'dI':>6}\")\n", + " for label, res in botsort_rows:\n", + " hota, assa, idf1, idsw = seq_metrics(res, seq)\n", + " ref_key = \"no_reid\" if \"baseline\" in label else \"mot17_th02\"\n", + " ref_vals = ref.get(ref_key, {})\n", + " ref_hota = ref_vals.get(\"hota\", float(\"nan\"))\n", + " ref_idf1 = ref_vals.get(\"idf1\", float(\"nan\"))\n", + " delta_h = hota - ref_hota if ref_hota == ref_hota else float(\"nan\")\n", + " delta_i = idf1 - ref_idf1 if ref_idf1 == ref_idf1 else float(\"nan\")\n", + " ref_h_s = f\"{ref_hota:6.2f}\" if ref_hota == ref_hota else \" n/a\"\n", + " ref_i_s = f\"{ref_idf1:6.2f}\" if ref_idf1 == ref_idf1 else \" n/a\"\n", + " delta_h_s = f\"{delta_h:+6.2f}\" if delta_h == delta_h else \" n/a\"\n", + " delta_i_s = f\"{delta_i:+6.2f}\" if delta_i == delta_i else \" n/a\"\n", + " print(f\" {label:<28} {hota:6.2f} {idf1:6.2f} {idsw:5d} {ref_h_s} {delta_h_s} {ref_i_s} {delta_i_s}\")\n", + " print()" + ] + }, + { + "cell_type": "markdown", + "id": "8de54c38", + "metadata": {}, + "source": [ + "### 8. Visual comparison - largest ReID gain sequence\n", + "\n", + "Side-by-side **baseline vs +ReID** video for the val sequence with the largest \u0394HOTA\n", + "(from the runs above). On Colab the mp4 is downloaded automatically.\n", + "" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9a4f4194", + "metadata": {}, + "outputs": [], + "source": [ + "# Auto-pick the sequence with the largest HOTA gain (override with COMPARE_SEQ = \"MOT17-02-FRCNN\").\n", + "COMPARE_SEQ: str | None = None\n", + "COMPARE_FPS = 30\n", + "COMPARE_MAX_FRAMES: int | None = None # None = full sequence\n", + "\n", + "\n", + "def _mot_frame_to_detections(mot: dict, frame_idx: int) -> sv.Detections:\n", + " frame = mot.get(frame_idx)\n", + " if frame is None:\n", + " return sv.Detections.empty()\n", + " active = frame.ids >= 0\n", + " if not np.any(active):\n", + " return sv.Detections.empty()\n", + " return sv.Detections(\n", + " xyxy=sv.xywh_to_xyxy(frame.boxes[active]).astype(np.float32),\n", + " tracker_id=frame.ids[active].astype(int),\n", + " confidence=frame.confidences[active].astype(np.float32),\n", + " )\n", + "\n", + "\n", + "def _annotate_tracks(frame_bgr: np.ndarray, detections: sv.Detections) -> np.ndarray:\n", + " if len(detections) == 0:\n", + " return frame_bgr\n", + " palette, lookup = sv.ColorPalette.DEFAULT, sv.ColorLookup.TRACK\n", + " scene = sv.BoxAnnotator(color=palette, color_lookup=lookup, thickness=2).annotate(frame_bgr, detections)\n", + " labels = [str(int(tid)) for tid in detections.tracker_id]\n", + " return sv.LabelAnnotator(\n", + " color=palette,\n", + " color_lookup=lookup,\n", + " text_color=sv.Color.BLACK,\n", + " text_scale=0.5,\n", + " ).annotate(scene, detections, labels=labels)\n", + "\n", + "\n", + "def _panel_badge(frame: np.ndarray, text: str, accent: tuple[int, int, int]) -> np.ndarray:\n", + " out = frame.copy()\n", + " font, scale, thick = cv2.FONT_HERSHEY_SIMPLEX, 0.7, 2\n", + " (tw, th), _ = cv2.getTextSize(text, font, scale, thick)\n", + " x, y, pad, bar = 12, 12, 10, 6\n", + " w, h = tw + 2 * pad + bar + 8, th + 2 * pad\n", + " cv2.rectangle(out, (x, y), (x + w, y + h), (24, 24, 28), -1)\n", + " cv2.rectangle(out, (x + 6, y + 6), (x + 6 + bar, y + h - 6), accent, -1)\n", + " cv2.putText(out, text, (x + bar + pad + 4, y + pad + th), font, scale, (245, 245, 245), thick, cv2.LINE_AA)\n", + " return out\n", + "\n", + "\n", + "seq_gains: list[tuple[str, float, float, float]] = []\n", + "for seq in ACTIVE_SEQUENCES:\n", + " h_b, _, i_b, _ = seq_metrics(result_baseline, seq)\n", + " h_r, _, i_r, _ = seq_metrics(result_reid, seq)\n", + " if h_b == h_b and h_r == h_r:\n", + " seq_gains.append((seq, h_r - h_b, i_r - i_b, h_r))\n", + "\n", + "if not seq_gains:\n", + " raise RuntimeError(\"No per-sequence metrics - run sections 5-7 first.\")\n", + "\n", + "seq_gains.sort(key=lambda row: row[1], reverse=True)\n", + "print(\"Per-sequence ReID \u0394HOTA (largest first):\")\n", + "for seq, dh, di, _ in seq_gains:\n", + " print(f\" {seq:<20} \u0394HOTA {dh:+6.2f} \u0394IDF1 {di:+6.2f}\")\n", + "\n", + "COMPARE_SEQ = COMPARE_SEQ or seq_gains[0][0]\n", + "delta_hota, delta_idf1 = next((dh, di) for s, dh, di, _ in seq_gains if s == COMPARE_SEQ)\n", + "print(f\"\\nRendering comparison for {COMPARE_SEQ} (\u0394HOTA {delta_hota:+.2f}, \u0394IDF1 {delta_idf1:+.2f})\")\n", + "\n", + "pred_base = OUTPUT_ROOT / \"botsort_baseline\" / \"preds\" / f\"{COMPARE_SEQ}.txt\"\n", + "pred_reid = OUTPUT_ROOT / \"botsort_reid\" / \"preds\" / f\"{COMPARE_SEQ}.txt\"\n", + "if not pred_base.is_file() or not pred_reid.is_file():\n", + " raise FileNotFoundError(f\"Missing preds for {COMPARE_SEQ}:\\n {pred_base}\\n {pred_reid}\")\n", + "\n", + "mot_base = load_mot_file(pred_base)\n", + "mot_reid = load_mot_file(pred_reid)\n", + "img_dir = SEQUENCE_PATHS[COMPARE_SEQ][\"img\"]\n", + "n_frames = SEQUENCE_PATHS[COMPARE_SEQ][\"n_frames\"]\n", + "if COMPARE_MAX_FRAMES is not None:\n", + " n_frames = min(n_frames, COMPARE_MAX_FRAMES)\n", + "\n", + "compare_fps = COMPARE_FPS\n", + "seqinfo = img_dir.parent / \"seqinfo.ini\"\n", + "if seqinfo.is_file():\n", + " for line in seqinfo.read_text().splitlines():\n", + " if line.startswith(\"frameRate=\"):\n", + " compare_fps = int(line.split(\"=\", 1)[1])\n", + " break\n", + "\n", + "sample = load_mot_frame_image(img_dir, 1)\n", + "h, w = sample.shape[:2]\n", + "\n", + "out_path = OUTPUT_ROOT / f\"compare_{COMPARE_SEQ}_baseline_vs_reid.mp4\"\n", + "video_info = sv.VideoInfo(width=w * 2, height=h, fps=compare_fps, total_frames=n_frames)\n", + "\n", + "with sv.VideoSink(target_path=str(out_path), video_info=video_info) as sink:\n", + " for frame_idx in range(1, n_frames + 1):\n", + " frame = load_mot_frame_image(img_dir, frame_idx)\n", + " left = _panel_badge(\n", + " _annotate_tracks(frame.copy(), _mot_frame_to_detections(mot_base, frame_idx)),\n", + " \"BASELINE (NO REID)\",\n", + " (0, 165, 255),\n", + " )\n", + " right = _panel_badge(\n", + " _annotate_tracks(frame.copy(), _mot_frame_to_detections(mot_reid, frame_idx)),\n", + " \"BOT-SORT + REID\",\n", + " (80, 200, 120),\n", + " )\n", + " sink.write_frame(np.hstack([left, right]))\n", + "\n", + "ffmpeg = shutil.which(\"ffmpeg\")\n", + "if ffmpeg is not None:\n", + " tmp = out_path.with_name(f\"{out_path.stem}.web.mp4\")\n", + " result = subprocess.run( # noqa: S603\n", + " [\n", + " ffmpeg,\n", + " \"-y\",\n", + " \"-i\",\n", + " str(out_path),\n", + " \"-c:v\",\n", + " \"libx264\",\n", + " \"-pix_fmt\",\n", + " \"yuv420p\",\n", + " \"-movflags\",\n", + " \"+faststart\",\n", + " \"-an\",\n", + " str(tmp),\n", + " ],\n", + " capture_output=True,\n", + " text=True,\n", + " )\n", + " if result.returncode == 0:\n", + " tmp.replace(out_path)\n", + "\n", + "print(f\"Wrote {out_path} ({n_frames} frames @ {compare_fps} fps)\")\n", + "ipy_display(Video(str(out_path), embed=True, width=min(960, w * 2)))\n", + "if IN_COLAB:\n", + " files.download(str(out_path))" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.13" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/pyproject.toml b/pyproject.toml index ed02d2a4f..f92696a56 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -47,6 +47,8 @@ dependencies = [ [project.optional-dependencies] detection = ["inference-models>=0.19.0"] tune = ["optuna>=3.0.0"] +# Temporary git pin until reid is on PyPI. Swap to reid>=0.1.0,<0.2 before merge. +reid = ["reid @ git+https://github.com/roboflow/re-ID.git@main"] [project.scripts] trackers = "trackers.scripts.__main__:main" @@ -196,7 +198,7 @@ markers = [ [tool.codespell] skip = "*.pth" -ignore-words-list = "mot" +ignore-words-list = "mot,STrack" [tool.mypy] python_version = "3.10" @@ -219,5 +221,7 @@ module = [ "rfdetr.*", "supervision", "supervision.*", + "reid", + "reid.*", ] ignore_missing_imports = true diff --git a/src/trackers/core/base.py b/src/trackers/core/base.py index c60329c21..a070fb11f 100644 --- a/src/trackers/core/base.py +++ b/src/trackers/core/base.py @@ -37,7 +37,11 @@ class ParameterInfo: class TrackerParameters(dict[str, ParameterInfo]): - """Tracker parameter mapping with CLI-only filtering for IoU metrics.""" + """Tracker parameter mapping that hides non-flaggable constructor args from the CLI. + + Omits IoU metric objects and injection-only parameters such as ``reid_model`` + from CLI flag generation while keeping them available on the tracker itself. + """ def items(self) -> Iterator[tuple[str, ParameterInfo]]: # type: ignore[override] try: @@ -47,6 +51,8 @@ def items(self) -> Iterator[tuple[str, ParameterInfo]]: # type: ignore[override return for name, param_info in super().items(): + if name == "reid_model": + continue param_type = param_info.param_type if isinstance(param_type, type) and issubclass(param_type, BaseIoU): continue diff --git a/src/trackers/core/botsort/fusion.py b/src/trackers/core/botsort/fusion.py new file mode 100644 index 000000000..80c70e288 --- /dev/null +++ b/src/trackers/core/botsort/fusion.py @@ -0,0 +1,47 @@ +# ------------------------------------------------------------------------ +# Trackers +# Copyright (c) 2026 Roboflow. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ +# +# Adapted from NirAharon/BoT-SORT (MIT) +# Copyright (c) 2022 Nir Aharon +# Source: https://github.com/NirAharon/BoT-SORT +# Reference: tracker/bot_sort.py (ReID appearance-IoU cost fusion) +# ------------------------------------------------------------------------ + +"""Appearance-IoU fusion for BoT-SORT ReID association.""" + +from __future__ import annotations + +import numpy as np + + +def fuse_botsort_reid_association( + association_similarity: np.ndarray, + appearance_similarity: np.ndarray, + *, + proximity_threshold: float, + appearance_threshold: float, + proximity_iou_similarity: np.ndarray | None = None, +) -> np.ndarray: + """Fuse IoU and appearance the way BoT-SORT ``bot_sort.py`` does. + + Computes ``min(association_cost, capped_appearance_cost)`` with proximity + and appearance gates, then returns the corresponding similarity matrix + (``1 - cost``). + + ``proximity_iou_similarity`` is the standard-IoU gate (defaults to + ``association_similarity``). Pass it separately when association uses + GIoU/DIoU/CIoU so proximity still uses plain IoU. + """ + if proximity_iou_similarity is None: + proximity_iou_similarity = association_similarity + + d_iou = 1.0 - association_similarity + d_iou_proximity = 1.0 - proximity_iou_similarity + d_app = 0.5 * (1.0 - appearance_similarity) + d_app = np.where(d_app > appearance_threshold, 1.0, d_app) + d_app = np.where(d_iou_proximity > proximity_threshold, 1.0, d_app) + fused_cost = np.minimum(d_iou, d_app) + return 1.0 - fused_cost diff --git a/src/trackers/core/botsort/tracker.py b/src/trackers/core/botsort/tracker.py index 700ffe651..296927f72 100644 --- a/src/trackers/core/botsort/tracker.py +++ b/src/trackers/core/botsort/tracker.py @@ -12,8 +12,12 @@ from scipy.optimize import linear_sum_assignment from trackers.core.base import BaseTracker +from trackers.core.botsort.fusion import fuse_botsort_reid_association from trackers.core.botsort.tracklet import BoTSORTTracklet from trackers.core.botsort.utils import _fuse_score, get_alive_tracklets +from trackers.core.reid.appearance import appearance_similarity, extract_detection_embeddings +from trackers.core.reid.encoder import ReIDEncoder +from trackers.core.reid.feature_bank import FeatureBank from trackers.utils.cmc import CMC, CMCConfig, CMCMethod from trackers.utils.detections import default_confidences from trackers.utils.iou import BaseIoU, IoU @@ -34,11 +38,11 @@ class BoTSORTTracker(BaseTracker): 3) Split tracks into confirmed, unconfirmed, and lost 4) Apply camera motion compensation to predicted tracks 5) Associate high-confidence detections to confirmed + lost tracks - (IoU fused with detection scores + assignment) + (IoU fused with detection scores, optional appearance) 6) Associate low-confidence detections to remaining tracks - (excluding lost tracks) + (excluding lost tracks; geometry only) 7) Match remaining unmatched high-confidence detections to unconfirmed tracks - and remove unmatched unconfirmed tracks + (optional appearance) and remove unmatched unconfirmed tracks 8) Spawn new tracks from still unmatched high-confidence detections (instantly activated on the very first frame) 9) Remove tracks that have been lost for too long @@ -84,13 +88,25 @@ class BoTSORTTracker(BaseTracker): Passing ``None`` (the default) is equivalent to ``IoU()`` and is provided for backward compatibility with existing code that did not supply an ``iou`` argument. + reid_model: Optional appearance encoder (``ReIDEncoder``) for appearance + association. Pass a ``reid.ReIDModel`` in normal use. Requires + ``frame`` in :meth:`update`. When ``None`` (default), behaviour + matches the geometry-only BoT-SORT baseline. + reid_ema_alpha: EMA momentum for track appearance features. Default ``0.9``. + appearance_threshold: Appearance distance gate. Rejects matches when the + halved cosine distance ``0.5 * (1 - cos_sim)`` exceeds this value. + Default ``0.25`` (BoT-SORT ``appearance_thresh``). + proximity_threshold: Standard-IoU distance gate applied before appearance + is used. Computed from true IoU even when ``iou`` is GIoU/DIoU/CIoU. + Default ``0.5`` (BoT-SORT ``proximity_thresh``; requires + ``IoU >= 1 - proximity_threshold``). Notes: - Positive `maximum_frames_without_update` values are scaled by ``frame_rate`` and rounded up to at least one missed frame. Explicit zero-buffer configurations remain zero. - - When CMC is enabled, pass the current video frame via the ``frame`` - argument of :meth:`update`. + - When CMC or ReID is enabled, pass the current video frame via the + ``frame`` argument of :meth:`update`. """ tracker_id = "botsort" @@ -124,6 +140,10 @@ def __init__( instant_first_frame_activation: bool = True, state_estimator_class: type[BaseStateEstimator] = XCYCWHStateEstimator, iou: BaseIoU | None = None, + reid_model: ReIDEncoder | None = None, + reid_ema_alpha: float = 0.9, + appearance_threshold: float = 0.25, + proximity_threshold: float = 0.5, ) -> None: self.maximum_frames_without_update = self._compute_maximum_frames_without_update( lost_track_buffer=lost_track_buffer, @@ -146,6 +166,17 @@ def __init__( self.enable_cmc = enable_cmc self.cmc = CMC(CMCConfig(method=cmc_method, downscale=cmc_downscale)) if enable_cmc else None + self.reid_model = reid_model + if not 0.0 <= reid_ema_alpha <= 1.0: + raise ValueError(f"reid_ema_alpha must be in [0, 1], got {reid_ema_alpha}") + self.reid_ema_alpha = reid_ema_alpha + if not 0.0 <= appearance_threshold <= 1.0: + raise ValueError(f"appearance_threshold must be in [0, 1], got {appearance_threshold}") + if not 0.0 <= proximity_threshold <= 1.0: + raise ValueError(f"proximity_threshold must be in [0, 1], got {proximity_threshold}") + self.appearance_threshold = appearance_threshold + self.proximity_threshold = proximity_threshold + self._init_timestamp_state(frame_rate) def update( @@ -182,10 +213,11 @@ def update( the last state. Notes: - - If CMC is enabled, pass the current video frame via ``frame`` so the - tracker can estimate a global affine transform and warp predicted - track states before association. When ``frame=None`` and - ``enable_cmc=True``, CMC is silently skipped for that step. + - If CMC or ReID is enabled, pass the current video frame via ``frame`` + so the tracker can estimate a global affine transform and/or extract + appearance embeddings before association. When ``frame=None`` and + ``enable_cmc=True``, CMC is silently skipped for that step; ReID + requires ``frame`` and raises when it is ``None``. """ timing = self._predict_timing(timestamp) if timing.skip_update: @@ -247,38 +279,49 @@ def update( mask_boxes = high_boxes if len(high_boxes) > 0 else None H = self.cmc.estimate(frame, mask_boxes) CMC.apply_batch(H, self.tracks) + + det_embeddings: np.ndarray | None = None + if self.reid_model is not None: + if frame is None: + raise ValueError(f"{type(self).__name__}.update() requires frame when reid_model is set.") + if len(high_boxes) > 0: + det_embeddings = extract_detection_embeddings(self.reid_model, frame, high_boxes) + # Step 1: associate high-confidence detections to confirmed + lost tracks. # Lost tracks are included here (following the original ByteTrack), and # IoU is fused with detection scores. strack_pool = confirmed_tracks + lost_tracks - iou_matrix = self._get_iou_matrix(strack_pool, high_boxes) - iou_matrix = _fuse_score(self.iou.normalize_for_fusion(iou_matrix), high_scores) + similarity_matrix = self._association_similarity(strack_pool, high_boxes, high_scores, det_embeddings) + matched, unmatched_pool, unmatched_high = self._get_associated_indices( - iou_matrix, self.minimum_iou_threshold_first_assoc + similarity_matrix, self.minimum_iou_threshold_first_assoc ) for row, col in matched: - track = strack_pool[row] - track.update(high_boxes[col]) - if track.number_of_successful_updates >= self.minimum_consecutive_frames and track.tracker_id == -1: - track.tracker_id = self._allocate_tracker_id() - out_det_indices.append(int(high_indices[col])) - out_tracker_ids.append(track.tracker_id) + self._assign_track_detection( + strack_pool[row], + high_boxes[col], + det_embeddings[col] if det_embeddings is not None else None, + int(high_indices[col]), + out_det_indices, + out_tracker_ids, + ) # Step 2: associate low-confidence detections to remaining *tracked* tracks # only (excluding lost tracks, following the original ByteTrack). - # No score fusing in second association. remaining_tracked = [strack_pool[i] for i in unmatched_pool if strack_pool[i].time_since_update == 1] iou_matrix = self._get_iou_matrix(remaining_tracked, low_boxes) matched, _, unmatched_low = self._get_associated_indices(iou_matrix, self.minimum_iou_threshold_second_assoc) for row, col in matched: - track = remaining_tracked[row] - track.update(low_boxes[col]) - if track.number_of_successful_updates >= self.minimum_consecutive_frames and track.tracker_id == -1: - track.tracker_id = self._allocate_tracker_id() - out_det_indices.append(int(low_indices[col])) - out_tracker_ids.append(track.tracker_id) + self._assign_track_detection( + remaining_tracked[row], + low_boxes[col], + None, + int(low_indices[col]), + out_det_indices, + out_tracker_ids, + ) # Unmatched low-confidence detections for det_local_idx in sorted(unmatched_low): @@ -294,21 +337,23 @@ def update( if len(unconfirmed_tracks) > 0 and len(unmatched_high_list) > 0: uh_boxes = high_boxes[unmatched_high_list] uh_scores = high_scores[unmatched_high_list] + uh_embeddings = det_embeddings[unmatched_high_list] if det_embeddings is not None else None + similarity_matrix = self._association_similarity(unconfirmed_tracks, uh_boxes, uh_scores, uh_embeddings) - iou_matrix = self._get_iou_matrix(unconfirmed_tracks, uh_boxes) - iou_matrix = _fuse_score(self.iou.normalize_for_fusion(iou_matrix), uh_scores) matched_uc, unmatched_uc_indices, remaining_uh = self._get_associated_indices( - iou_matrix, self.minimum_iou_threshold_unconfirmed_assoc + similarity_matrix, self.minimum_iou_threshold_unconfirmed_assoc ) for row, col in matched_uc: - track = unconfirmed_tracks[row] orig_high_idx = unmatched_high_list[col] - track.update(high_boxes[orig_high_idx]) - if track.number_of_successful_updates >= self.minimum_consecutive_frames and track.tracker_id == -1: - track.tracker_id = self._allocate_tracker_id() - out_det_indices.append(int(high_indices[orig_high_idx])) - out_tracker_ids.append(track.tracker_id) + self._assign_track_detection( + unconfirmed_tracks[row], + high_boxes[orig_high_idx], + det_embeddings[orig_high_idx] if det_embeddings is not None else None, + int(high_indices[orig_high_idx]), + out_det_indices, + out_tracker_ids, + ) # Only remaining unmatched high-conf dets proceed to spawning unmatched_high = [unmatched_high_list[i] for i in remaining_uh] @@ -328,6 +373,7 @@ def update( out_det_indices, out_tracker_ids, is_first_frame=(self.frame_id == 1), + det_embeddings=det_embeddings, ) # Full lifecycle prune: removes immature+unmatched and any remaining expired @@ -349,12 +395,61 @@ def update( result.tracker_id = np.array(out_tracker_ids, dtype=int) return result - def _get_iou_matrix(self, tracklets: list[BoTSORTTracklet], detections: np.ndarray) -> np.ndarray: + def _assign_track_detection( + self, + track: BoTSORTTracklet, + bbox: np.ndarray, + embedding: np.ndarray | None, + global_det_index: int, + out_det_indices: list[int], + out_tracker_ids: list[int], + ) -> None: + """Update a track from a matched detection and record output indices.""" + track.update(bbox) + if track.feature_bank is not None and embedding is not None: + track.feature_bank.update(embedding) + if track.number_of_successful_updates >= self.minimum_consecutive_frames and track.tracker_id == -1: + track.tracker_id = self._allocate_tracker_id() + out_det_indices.append(global_det_index) + out_tracker_ids.append(track.tracker_id) + + def _association_similarity( + self, + tracklets: list[BoTSORTTracklet], + boxes: np.ndarray, + scores: np.ndarray, + embeddings: np.ndarray | None, + ) -> np.ndarray: + """Score-fused association similarity, with optional BoT-SORT ReID fusion.""" + iou_sim_raw = self.iou.normalize_for_fusion(self._get_iou_matrix(tracklets, boxes)) + iou_sim_fused = _fuse_score(iou_sim_raw, scores) + if embeddings is None or len(tracklets) == 0: + return iou_sim_fused + + track_feats = [None if t.feature_bank is None else t.feature_bank.feature for t in tracklets] + proximity_iou = ( + iou_sim_raw if isinstance(self.iou, IoU) else self._get_iou_matrix(tracklets, boxes, metric=IoU()) + ) + return fuse_botsort_reid_association( + iou_sim_fused, + appearance_similarity(track_feats, embeddings), + proximity_iou_similarity=proximity_iou, + proximity_threshold=self.proximity_threshold, + appearance_threshold=self.appearance_threshold, + ) + + def _get_iou_matrix( + self, + tracklets: list[BoTSORTTracklet], + detections: np.ndarray, + *, + metric: BaseIoU | None = None, + ) -> np.ndarray: if len(tracklets) == 0: tracklet_boxes = np.empty((0, 4)) else: tracklet_boxes = np.array([tracklet.get_state_bbox() for tracklet in tracklets]) - return self.iou.compute(tracklet_boxes, detections) + return (metric or self.iou).compute(tracklet_boxes, detections) def _get_associated_indices( self, @@ -405,6 +500,7 @@ def _spawn_new_tracks( out_det_indices: list[int], out_tracker_ids: list[int], is_first_frame: bool = False, + det_embeddings: np.ndarray | None = None, ) -> None: """Create new tracklets from unmatched high-confidence detections. @@ -422,6 +518,10 @@ def _spawn_new_tracks( initial_bbox=detection_boxes[global_idx], state_estimator_class=self.state_estimator_class, ) + if self.reid_model is not None: + tracklet.feature_bank = FeatureBank(self.reid_ema_alpha) + if det_embeddings is not None: + tracklet.feature_bank.update(det_embeddings[det_local_idx]) if is_first_frame and self.instant_first_frame_activation: tracklet.tracker_id = self._allocate_tracker_id() self.tracks.append(tracklet) diff --git a/src/trackers/core/botsort/tracklet.py b/src/trackers/core/botsort/tracklet.py index 1f384f15d..7ebf11071 100644 --- a/src/trackers/core/botsort/tracklet.py +++ b/src/trackers/core/botsort/tracklet.py @@ -6,6 +6,8 @@ from __future__ import annotations +from typing import TYPE_CHECKING + import numpy as np from trackers.utils.base_tracklet import BaseTracklet @@ -19,6 +21,9 @@ XYXYStateEstimator, ) +if TYPE_CHECKING: + from trackers.core.reid.feature_bank import FeatureBank + class BoTSORTTracklet(BaseTracklet): """Tracklet for the BoT-SORT tracker. @@ -53,6 +58,8 @@ def __init__( # Count initial bbox as first successful update so that # number_of_successful_updates starts at 1. self.number_of_successful_updates = 1 + # Optional appearance feature bank, set by BoTSORTTracker when ReID is enabled. + self.feature_bank: FeatureBank | None = None def _configure_initial_noise(self, bbox: np.ndarray) -> None: """Set initial P, Q, R based on the first detection's size.""" diff --git a/src/trackers/core/reid/__init__.py b/src/trackers/core/reid/__init__.py new file mode 100644 index 000000000..c8fbe85fa --- /dev/null +++ b/src/trackers/core/reid/__init__.py @@ -0,0 +1,20 @@ +# ------------------------------------------------------------------------ +# Trackers +# Copyright (c) 2026 Roboflow. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ + +"""Appearance-ReID association helpers for multi-object trackers.""" + +from __future__ import annotations + +from trackers.core.reid.appearance import appearance_similarity, extract_detection_embeddings +from trackers.core.reid.encoder import ReIDEncoder +from trackers.core.reid.feature_bank import FeatureBank + +__all__ = [ + "FeatureBank", + "ReIDEncoder", + "appearance_similarity", + "extract_detection_embeddings", +] diff --git a/src/trackers/core/reid/appearance.py b/src/trackers/core/reid/appearance.py new file mode 100644 index 000000000..7341058a7 --- /dev/null +++ b/src/trackers/core/reid/appearance.py @@ -0,0 +1,100 @@ +# ------------------------------------------------------------------------ +# Trackers +# Copyright (c) 2026 Roboflow. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ + +"""Appearance embedding helpers for tracker association.""" + +from __future__ import annotations + +from collections.abc import Sequence + +import numpy as np +import supervision as sv + +from trackers.core.reid.encoder import ReIDEncoder + +_NORM_EPS = 1e-12 + + +def _require_embedding_matrix(embeddings: np.ndarray) -> np.ndarray: + """Return a finite float32 embedding matrix.""" + cleaned = np.asarray(embeddings, dtype=np.float32) + if cleaned.ndim != 2: + raise ValueError(f"embeddings must be 2-D, got shape {cleaned.shape}") + if cleaned.size > 0 and not np.all(np.isfinite(cleaned)): + raise ValueError("embeddings must contain only finite values") + return cleaned + + +def _l2_normalize(embedding: np.ndarray) -> np.ndarray: + """Return an L2-normalised 1-D vector.""" + flat = np.asarray(embedding, dtype=np.float64).reshape(-1) + if flat.size == 0: + raise ValueError("embedding must be non-empty") + if not np.all(np.isfinite(flat)): + raise ValueError("embedding must contain only finite values") + norm = float(np.linalg.norm(flat)) + return (flat / max(norm, _NORM_EPS)).astype(np.float32) + + +def _l2_normalize_rows(embeddings: np.ndarray) -> np.ndarray: + """L2-normalise each row in an embedding matrix.""" + if embeddings.size == 0: + return embeddings + mat = embeddings.astype(np.float64) + norms = np.linalg.norm(mat, axis=1, keepdims=True) + return (mat / np.maximum(norms, _NORM_EPS)).astype(np.float32) + + +def extract_detection_embeddings( + model: ReIDEncoder, + frame: np.ndarray, + boxes: np.ndarray, +) -> np.ndarray: + """Extract appearance embeddings for detection boxes.""" + if len(boxes) == 0: + return np.empty((0, 0), dtype=np.float32) + embeddings = _require_embedding_matrix(model.extract_features(sv.Detections(xyxy=boxes), frame)) + if embeddings.shape[0] != len(boxes): + raise ValueError(f"embedding rows ({embeddings.shape[0]}) must match detection boxes ({len(boxes)})") + return embeddings + + +def appearance_similarity( + track_features: Sequence[np.ndarray | None], + det_embeddings: np.ndarray, +) -> np.ndarray: + """Cosine similarity between track and detection embeddings.""" + n_tracks = len(track_features) + det_embeddings = _l2_normalize_rows(_require_embedding_matrix(det_embeddings)) + n_dets = det_embeddings.shape[0] + similarity = np.zeros((n_tracks, n_dets), dtype=np.float32) + + if n_tracks == 0 or n_dets == 0: + return similarity + + embed_dim = det_embeddings.shape[1] + track_rows: list[np.ndarray] = [] + kept_indices: list[int] = [] + for track_idx, feature in enumerate(track_features): + if feature is None: + continue + flat = np.asarray(feature, dtype=np.float32).reshape(-1) + if flat.shape[0] != embed_dim: + raise ValueError( + f"track feature dim {flat.shape[0]} does not match detection " + f"embedding dim {embed_dim} (track index {track_idx})" + ) + track_rows.append(_l2_normalize(flat)) + kept_indices.append(track_idx) + + if not track_rows: + return similarity + + cosine_similarities = (np.stack(track_rows) @ det_embeddings.T).astype(np.float32) + for local_idx, track_idx in enumerate(kept_indices): + similarity[track_idx] = cosine_similarities[local_idx] + + return similarity diff --git a/src/trackers/core/reid/encoder.py b/src/trackers/core/reid/encoder.py new file mode 100644 index 000000000..bfa2efc08 --- /dev/null +++ b/src/trackers/core/reid/encoder.py @@ -0,0 +1,30 @@ +# ------------------------------------------------------------------------ +# Trackers +# Copyright (c) 2026 Roboflow. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ + +"""Encoder protocol for ReID association.""" + +from __future__ import annotations + +from typing import Protocol + +import numpy as np +import supervision as sv + + +class ReIDEncoder(Protocol): + """Encoder with ``extract_features(detections, frame)``.""" + + def extract_features(self, detections: sv.Detections, frame: np.ndarray) -> np.ndarray: + """Return appearance embeddings for each detection box. + + Args: + detections: Boxes to embed (``xyxy``). + frame: BGR frame the detections were produced on. + + Returns: + Float32 array of shape ``(N, D)``, or ``(0, 0)`` when empty. + """ + ... diff --git a/src/trackers/core/reid/feature_bank.py b/src/trackers/core/reid/feature_bank.py new file mode 100644 index 000000000..5f5439285 --- /dev/null +++ b/src/trackers/core/reid/feature_bank.py @@ -0,0 +1,48 @@ +# ------------------------------------------------------------------------ +# Trackers +# Copyright (c) 2026 Roboflow. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ + +"""Per-track exponential moving average feature bank.""" + +from __future__ import annotations + +import numpy as np + +from trackers.core.reid.appearance import _l2_normalize + + +class FeatureBank: + """Per-track EMA unit embedding (L2 before and after blend). + + Follows BoT-SORT ``STrack.update_features`` + (https://github.com/NirAharon/BoT-SORT/blob/main/tracker/bot_sort.py). + """ + + def __init__(self, alpha: float = 0.9) -> None: + if not 0.0 <= alpha <= 1.0: + raise ValueError(f"alpha must be in [0, 1], got {alpha}") + self._alpha = alpha + self._feature: np.ndarray | None = None + + @property + def feature(self) -> np.ndarray | None: + """Current stored unit embedding, or ``None`` if never updated.""" + return None if self._feature is None else self._feature.copy() + + def update(self, embedding: np.ndarray) -> None: + """Blend an L2-normalized embedding into the stored unit feature.""" + cleaned = _l2_normalize(embedding) + + if self._feature is None: + self._feature = cleaned + return + + if self._feature.shape != cleaned.shape: + raise ValueError( + f"embedding shape {cleaned.shape} does not match stored feature shape {self._feature.shape}" + ) + + blended = self._alpha * self._feature + (1.0 - self._alpha) * cleaned + self._feature = _l2_normalize(blended) diff --git a/src/trackers/scripts/track.py b/src/trackers/scripts/track.py index 539a3a237..5568bcd2a 100644 --- a/src/trackers/scripts/track.py +++ b/src/trackers/scripts/track.py @@ -150,6 +150,46 @@ def add_track_subparser(subparsers: argparse._SubParsersAction) -> None: # Add dynamic tracker parameters _add_tracker_params(tracker_group) + reid_group = parser.add_argument_group("reid (BoT-SORT only; requires trackers[reid])") + reid_group.add_argument( + "--tracker.reid.enable", + action="store_true", + dest="tracker_reid_enable", + help="Enable appearance-based ReID association for BoT-SORT.", + ) + reid_group.add_argument( + "--tracker.reid.model", + type=str, + default=None, + dest="tracker_reid_model", + metavar="SOURCE", + help=( + "ReID checkpoint source (curated alias, hf:// URL, local path, or " + "save_pretrained directory). Implies --tracker.reid.enable. " + "Default alias when omitted: osnet_x1_0_msmt17_combineall." + ), + ) + reid_group.add_argument( + "--tracker.reid.device", + type=str, + default=DEFAULT_DEVICE, + dest="tracker_reid_device", + metavar="DEVICE", + help=f"ReID compute device: auto, cpu, cuda, mps. Default: {DEFAULT_DEVICE}", + ) + reid_group.add_argument( + "--tracker.reid.architecture", + type=str, + default=None, + dest="tracker_reid_architecture", + metavar="NAME", + help=( + "Backbone architecture for bare local .pth/.safetensors weights " + "(e.g. osnet_x1_0, fastreid_sbs_resnest50). Required when " + "--tracker.reid.model points to a bare weights file." + ), + ) + # Output options output_group = parser.add_argument_group("output") output_group.add_argument( @@ -310,6 +350,10 @@ def run_track(args: argparse.Namespace) -> int: # Create tracker tracker_params = _extract_tracker_params(args.tracker, args) + tracker_params, reid_error = _apply_reid_tracker_params(args.tracker, args, tracker_params) + if reid_error is not None: + print(reid_error, file=sys.stderr) + return 1 tracker = _init_tracker(args.tracker, **tracker_params) if args.source is not None: @@ -363,7 +407,7 @@ def _run_frameless( mask = np.isin(detections.class_id, class_filter) detections = detections[mask] # type: ignore[assignment] - tracked = tracker.update(detections) + tracked = tracker.update(detections, frame=None) if track_id_filter is not None and len(tracked) > 0: if tracked.tracker_id is not None: @@ -600,6 +644,61 @@ def _run_model(model: AnyModel, frame: np.ndarray, confidence: float) -> sv.Dete return detections +def _reid_requested(args: argparse.Namespace) -> bool: + return bool(args.tracker_reid_enable) or args.tracker_reid_model is not None + + +def _apply_reid_tracker_params( + tracker_id: str, + args: argparse.Namespace, + params: dict[str, object], +) -> tuple[dict[str, object], str | None]: + """Attach a CLI-instantiated ReID model to BoT-SORT tracker params.""" + if not _reid_requested(args): + return params, None + + if tracker_id != "botsort": + return params, (f"Error: --tracker.reid.* options apply only to --tracker botsort, got {tracker_id!r}.") + + if args.source is None: + return params, ( + "Error: ReID-enabled BoT-SORT requires --source (video/webcam/images) " + "so appearance embeddings can be extracted from frames." + ) + + try: + from reid import ReIDModel + except ImportError: + return params, ( + "Error: ReID tracking requires the optional `trackers[reid]` extra.\n" + "Install with: pip install 'trackers[reid]'" + ) + + model_source = args.tracker_reid_model + architecture = args.tracker_reid_architecture + if architecture is not None and model_source is None: + return params, ( + "Error: --tracker.reid.architecture requires --tracker.reid.model (bare weights need a checkpoint path)." + ) + + load_kwargs: dict[str, object] = {"device": args.tracker_reid_device} + if model_source is not None: + load_kwargs["source"] = model_source + if architecture is not None: + load_kwargs["architecture"] = architecture + + try: + reid_model = ReIDModel.from_pretrained(**load_kwargs) + except KeyboardInterrupt: + raise + except (OSError, ValueError, RuntimeError) as exc: + return params, f"Error: Failed to load ReID model: {exc}" + + params = dict(params) + params["reid_model"] = reid_model + return params, None + + def _extract_tracker_params(tracker_id: str, args: argparse.Namespace) -> dict[str, object]: """Extract tracker parameters from CLI args. diff --git a/tests/core/test_botsort_reid.py b/tests/core/test_botsort_reid.py new file mode 100644 index 000000000..4e899e1bb --- /dev/null +++ b/tests/core/test_botsort_reid.py @@ -0,0 +1,352 @@ +# ------------------------------------------------------------------------ +# Trackers +# Copyright (c) 2026 Roboflow. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ + +"""BoT-SORT ReID association and fusion tests.""" + +from __future__ import annotations + +import subprocess +import sys + +import numpy as np +import pytest +import supervision as sv + +from trackers.core.botsort.fusion import fuse_botsort_reid_association +from trackers.core.botsort.tracker import BoTSORTTracker +from trackers.core.reid.appearance import ( + appearance_similarity, + extract_detection_embeddings, +) +from trackers.core.reid.feature_bank import FeatureBank + + +def _detection(xyxy: tuple[float, float, float, float], conf: float = 0.9) -> sv.Detections: + return sv.Detections( + xyxy=np.array([xyxy], dtype=np.float32), + confidence=np.array([conf], dtype=np.float32), + ) + + +def _frame(seed: int = 0) -> np.ndarray: + rng = np.random.default_rng(seed) + return rng.integers(0, 255, (128, 128, 3), dtype=np.uint8) + + +def _norm(vec: np.ndarray) -> np.ndarray: + vec = vec.astype(np.float32) + return vec / np.linalg.norm(vec) + + +class _KeyedReIDEncoder: + """Deterministic embeddings keyed by detection top-left corner.""" + + def __init__(self, table: dict[tuple[int, int], np.ndarray] | None = None) -> None: + self.table = table or {} + self.calls = 0 + + def extract_features(self, detections: sv.Detections, frame: np.ndarray) -> np.ndarray: + self.calls += 1 + if len(detections) == 0: + return np.empty((0, 0), dtype=np.float32) + rows = [] + for box in detections.xyxy: + key = (round(float(box[0])), round(float(box[1]))) + rows.append(self.table.get(key, _norm(np.array([float(box[0]), float(box[1]), 1.0, 0.0])))) + return np.stack(rows) + + +def test_botsort_import_does_not_load_reid_model_stack() -> None: + """Importing BoT-SORT must not pull the heavy ``reid`` package (torch etc.).""" + result = subprocess.run( + [ + sys.executable, + "-c", + ( + "import sys; import trackers.core.botsort.tracker; " + "assert 'reid' not in sys.modules; " + "assert 'torch' not in sys.modules" + ), + ], + check=False, + capture_output=True, + text=True, + ) + assert result.returncode == 0, result.stderr + + +class TestFeatureBank: + """Unit tests for ``FeatureBank`` L2 + EMA behavior.""" + + def test_first_update_normalizes_embedding(self) -> None: + # BoT-SORT STrack.update_features: L2-normalize before storage. + bank = FeatureBank(alpha=0.9) + bank.update(np.array([3.0, 4.0], dtype=np.float32)) + feature = bank.feature + assert feature is not None + np.testing.assert_allclose(feature, [0.6, 0.8], atol=1e-6) + + def test_blends_on_unit_sphere(self) -> None: + # BoT-SORT: EMA on unit vectors, then L2-normalize the blend again. + bank = FeatureBank(alpha=0.75) + bank.update(np.array([1.0, 0.0], dtype=np.float32)) + bank.update(np.array([0.0, 1.0], dtype=np.float32)) + feature = bank.feature + assert feature is not None + # 0.75*[1,0] + 0.25*[0,1] = [0.75, 0.25], then / ||.|| + expected = np.array([0.75, 0.25], dtype=np.float32) + expected /= np.linalg.norm(expected) + np.testing.assert_allclose(feature, expected, atol=1e-6) + np.testing.assert_allclose(np.linalg.norm(feature), 1.0, atol=1e-6) + + def test_zero_embedding_is_accepted(self) -> None: + bank = FeatureBank() + bank.update(np.zeros(8, dtype=np.float32)) + feature = bank.feature + assert feature is not None + np.testing.assert_allclose(feature, 0.0) + + def test_non_finite_embedding_raises(self) -> None: + bank = FeatureBank() + with pytest.raises(ValueError, match="finite"): + bank.update(np.array([1.0, np.nan], dtype=np.float32)) + assert bank.feature is None + + def test_shape_change_raises(self) -> None: + bank = FeatureBank() + bank.update(np.array([1.0, 0.0], dtype=np.float32)) + before = bank.feature + assert before is not None + with pytest.raises(ValueError, match="shape"): + bank.update(np.array([1.0, 0.0, 0.0], dtype=np.float32)) + after = bank.feature + assert after is not None + np.testing.assert_allclose(after, before) + + +class TestAppearanceSimilarity: + """Unit tests for cosine ``appearance_similarity`` and embedding extraction.""" + + def test_identical_vectors_are_one(self) -> None: + similarity = appearance_similarity( + [np.array([1.0, 0.0], dtype=np.float32)], + np.array([[1.0, 0.0]], dtype=np.float32), + ) + np.testing.assert_allclose(similarity, [[1.0]], atol=1e-6) + + def test_orthogonal_vectors_are_zero(self) -> None: + similarity = appearance_similarity( + [np.array([1.0, 0.0], dtype=np.float32)], + np.array([[0.0, 1.0]], dtype=np.float32), + ) + np.testing.assert_allclose(similarity, [[0.0]], atol=1e-6) + + def test_normalizes_both_sides(self) -> None: + similarity = appearance_similarity( + [np.array([3.0, 4.0], dtype=np.float32)], + np.array([[6.0, 8.0]], dtype=np.float32), + ) + np.testing.assert_allclose(similarity, [[1.0]], atol=1e-6) + + def test_none_track_yields_zero_row(self) -> None: + similarity = appearance_similarity( + [None, np.array([1.0, 0.0], dtype=np.float32)], + np.array([[1.0, 0.0]], dtype=np.float32), + ) + np.testing.assert_allclose(similarity, [[0.0], [1.0]], atol=1e-6) + + def test_empty_inputs_return_empty_matrix(self) -> None: + assert appearance_similarity([], np.empty((0, 4), dtype=np.float32)).shape == (0, 0) + assert appearance_similarity( + [np.array([1.0, 0.0], dtype=np.float32)], + np.empty((0, 2), dtype=np.float32), + ).shape == (1, 0) + + def test_non_finite_detection_rows_raise(self) -> None: + with pytest.raises(ValueError, match="finite"): + appearance_similarity( + [np.array([1.0, 0.0], dtype=np.float32)], + np.array([[1.0, 0.0], [np.nan, 1.0]], dtype=np.float32), + ) + + def test_incompatible_track_dimensions_raise(self) -> None: + with pytest.raises(ValueError, match="dim"): + appearance_similarity( + [np.array([1.0, 0.0, 0.0], dtype=np.float32)], + np.array([[1.0, 0.0]], dtype=np.float32), + ) + + def test_extract_detection_embeddings_requires_one_row_per_box(self) -> None: + # Encoder must return embeddings.shape[0] == len(boxes). + class _WrongLengthEncoder: + def extract_features(self, detections: sv.Detections, frame: np.ndarray) -> np.ndarray: + return np.empty((0, 4), dtype=np.float32) + + with pytest.raises(ValueError, match="rows"): + extract_detection_embeddings( + _WrongLengthEncoder(), + _frame(), + np.array([[0.0, 0.0, 10.0, 10.0]], dtype=np.float32), + ) + + +class TestFuseBotsortReidAssociation: + """Unit tests for BoT-SORT IoU/appearance fusion gates.""" + + def test_appearance_can_win_when_proximity_passes(self) -> None: + # Association IoU 0.63 clears the proximity gate (needs IoU > 1 - 0.5 = 0.5), + # so a strong appearance score can beat it (0.63 → 0.9). + fused = fuse_botsort_reid_association( + np.array([[0.63]], dtype=np.float32), + np.array([[0.8]], dtype=np.float32), + proximity_threshold=0.5, + appearance_threshold=0.25, + ) + assert fused[0, 0] == pytest.approx(0.9) + + def test_low_proximity_ignores_appearance(self) -> None: + # Association IoU 0.36 fails the proximity gate (needs IoU > 1 - 0.5 = 0.5), + # so appearance is discarded even though it is strong (0.9). Score stays IoU-only. + iou_only = np.array([[0.36]], dtype=np.float32) + fused = fuse_botsort_reid_association( + iou_only, + np.array([[0.9]], dtype=np.float32), + proximity_threshold=0.5, + appearance_threshold=0.25, + ) + assert fused[0, 0] == pytest.approx(float(iou_only[0, 0])) + + def test_proximity_uses_standard_iou_not_giou(self) -> None: + # Association score uses a high GIoU-like value (0.80), but standard IoU is + # only 0.35 and fails the proximity gate, so appearance must not be used. + association_iou = np.array([[0.80]], dtype=np.float32) + fused = fuse_botsort_reid_association( + association_iou, + np.array([[0.95]], dtype=np.float32), + proximity_threshold=0.5, + appearance_threshold=0.25, + proximity_iou_similarity=np.array([[0.35]], dtype=np.float32), + ) + assert fused[0, 0] == pytest.approx(float(association_iou[0, 0])) + + +class TestBoTSORTTrackerReID: + """Integration-style tests for BoT-SORT tracker appearance association.""" + + def test_rejects_invalid_reid_ema_alpha(self) -> None: + with pytest.raises(ValueError, match="reid_ema_alpha"): + BoTSORTTracker(enable_cmc=False, reid_model=_KeyedReIDEncoder(), reid_ema_alpha=1.5) + + def test_requires_frame_when_reid_enabled(self) -> None: + tracker = BoTSORTTracker(enable_cmc=False, reid_model=_KeyedReIDEncoder()) + with pytest.raises(ValueError, match="requires frame"): + tracker.update(_detection((10.0, 10.0, 30.0, 30.0))) + + def test_feature_bank_initializes_on_spawn(self) -> None: + tracker = BoTSORTTracker(enable_cmc=False, reid_model=_KeyedReIDEncoder()) + tracker.update(_detection((10.0, 10.0, 30.0, 30.0)), frame=_frame()) + bank = tracker.tracks[0].feature_bank + assert bank is not None and bank.feature is not None + + def test_appearance_changes_assignment_vs_geometry_only(self) -> None: + identity = _norm(np.array([1.0, 0.0, 0.0, 0.0])) + impostor = _norm(np.array([0.0, 1.0, 0.0, 0.0])) + + class _PhaseEncoder: + phase = 1 + + def extract_features(self, detections: sv.Detections, frame: np.ndarray) -> np.ndarray: + rows = [] + for box in detections.xyxy: + key = (round(float(box[0])), round(float(box[1]))) + if self.phase == 1: + rows.append(identity) + elif key == (10, 10): + rows.append(impostor) + else: + rows.append(identity) + return np.stack(rows) + + encoder = _PhaseEncoder() + frame = _frame(1) + geo = BoTSORTTracker( + enable_cmc=False, + minimum_iou_threshold_first_assoc=0.01, + appearance_threshold=0.6, + proximity_threshold=0.99, + ) + geo.update(_detection((10.0, 10.0, 30.0, 30.0)), frame=frame) + + reid = BoTSORTTracker( + enable_cmc=False, + minimum_iou_threshold_first_assoc=0.01, + appearance_threshold=0.6, + proximity_threshold=0.99, + reid_model=encoder, + ) + reid.update(_detection((10.0, 10.0, 30.0, 30.0)), frame=frame) + track_id = int(reid.tracks[0].tracker_id) + + competitors = sv.Detections( + xyxy=np.array([[10.0, 10.0, 30.0, 30.0], [14.0, 14.0, 34.0, 34.0]], dtype=np.float32), + confidence=np.array([0.9, 0.9], dtype=np.float32), + ) + encoder.phase = 2 + geo_out = geo.update(competitors, frame=frame) + reid_out = reid.update(competitors, frame=frame) + + def _matched_xy(out: sv.Detections, tid: int) -> tuple[float, float]: + assert out.tracker_id is not None + box = out.xyxy[out.tracker_id == tid][0] + return float(box[0]), float(box[1]) + + assert _matched_xy(geo_out, int(geo.tracks[0].tracker_id)) == (10.0, 10.0) + assert _matched_xy(reid_out, track_id) == (14.0, 14.0) + + def test_low_confidence_stage_does_not_update_feature_bank(self) -> None: + model = _KeyedReIDEncoder({(10, 10): _norm(np.array([1.0, 0.0, 0.0, 0.0]))}) + tracker = BoTSORTTracker( + enable_cmc=False, + reid_model=model, + high_conf_det_threshold=0.8, + minimum_iou_threshold_second_assoc=0.01, + ) + tracker.update(_detection((10.0, 10.0, 30.0, 30.0)), frame=_frame(4)) + bank = tracker.tracks[0].feature_bank + assert bank is not None + before = bank.feature + assert before is not None + + calls_after_high = model.calls + tracker.update(_detection((12.0, 12.0, 32.0, 32.0), conf=0.5), frame=_frame(4)) + assert model.calls == calls_after_high + after = bank.feature + assert after is not None + np.testing.assert_allclose(before, after) + + @pytest.mark.integration + def test_real_reid_model_runs_over_frames(self) -> None: + """Smoke the ``trackers`` → ``reid`` boundary with a real encoder.""" + import reid + + reid_model = reid.ReIDModel.from_pretrained(architecture="osnet_x0_25", device="cpu") + tracker = BoTSORTTracker(enable_cmc=False, reid_model=reid_model) + + rng = np.random.default_rng(0) + box = np.array([30.0, 30.0, 70.0, 90.0], dtype=np.float32) + for _ in range(3): + frame = rng.integers(0, 255, (128, 128, 3), dtype=np.uint8) + detections = sv.Detections( + xyxy=box[None, :].copy(), + confidence=np.array([0.9], dtype=np.float32), + ) + result = tracker.update(detections, frame=frame) + assert result.tracker_id is not None + box = box + np.array([2.0, 1.0, 2.0, 1.0], dtype=np.float32) + + assert len(tracker.tracks) == 1 + bank = tracker.tracks[0].feature_bank + assert bank is not None and bank.feature is not None diff --git a/tests/scripts/test_track.py b/tests/scripts/test_track.py index be3867edb..7eac07eb6 100644 --- a/tests/scripts/test_track.py +++ b/tests/scripts/test_track.py @@ -6,6 +6,8 @@ from __future__ import annotations +import argparse +from pathlib import Path from typing import ClassVar import numpy as np @@ -13,10 +15,13 @@ import supervision as sv from trackers.scripts.track import ( + _apply_reid_tracker_params, _format_labels, _init_annotators, + _reid_requested, _resolve_class_filter, _resolve_track_id_filter, + add_track_subparser, ) @@ -194,3 +199,97 @@ def test_all_non_integer_returns_none(self, capsys: pytest.CaptureFixture) -> No result = _resolve_track_id_filter("abc,def") assert result is None assert "abc" in capsys.readouterr().err + + +class _FakeReIDModel: + last_kwargs: ClassVar[dict | None] = None + + @classmethod + def from_pretrained(cls, **kwargs: object) -> _FakeReIDModel: + cls.last_kwargs = dict(kwargs) + return cls() + + +class TestReidTrackCli: + """CLI wiring for optional BoT-SORT ReID model loading.""" + + def test_model_source_implies_enable(self) -> None: + args = argparse.Namespace( + tracker_reid_enable=False, + tracker_reid_model="osnet_x1_0_msmt17_combineall", + ) + assert _reid_requested(args) + + def test_requires_source_before_load(self) -> None: + args = argparse.Namespace( + tracker_reid_enable=True, + tracker_reid_model=None, + tracker_reid_device="cpu", + tracker_reid_architecture=None, + source=None, + ) + params, err = _apply_reid_tracker_params("botsort", args, {}) + assert err is not None and "--source" in err + assert params == {} + + def test_passes_architecture(self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + import sys + from types import ModuleType + + fake_reid = ModuleType("reid") + setattr(fake_reid, "ReIDModel", _FakeReIDModel) + monkeypatch.setitem(sys.modules, "reid", fake_reid) + + weights = tmp_path / "weights.pth" + weights.touch() + args = argparse.Namespace( + tracker_reid_enable=False, + tracker_reid_model=str(weights), + tracker_reid_device="cpu", + tracker_reid_architecture="osnet_x1_0", + source="video.mp4", + ) + params, err = _apply_reid_tracker_params("botsort", args, {}) + assert err is None + assert _FakeReIDModel.last_kwargs is not None + assert _FakeReIDModel.last_kwargs["architecture"] == "osnet_x1_0" + assert "reid_model" in params + + def test_rejects_non_botsort_tracker(self) -> None: + args = argparse.Namespace( + tracker_reid_enable=True, + tracker_reid_model=None, + tracker_reid_device="cpu", + tracker_reid_architecture=None, + source="video.mp4", + ) + _, error = _apply_reid_tracker_params("bytetrack", args, {}) + assert error is not None and "botsort" in error + + def test_architecture_requires_model(self, monkeypatch: pytest.MonkeyPatch) -> None: + import sys + from types import ModuleType + + fake_reid = ModuleType("reid") + setattr(fake_reid, "ReIDModel", object) + monkeypatch.setitem(sys.modules, "reid", fake_reid) + + args = argparse.Namespace( + tracker_reid_enable=True, + tracker_reid_model=None, + tracker_reid_device="cpu", + tracker_reid_architecture="osnet_x0_25", + source="video.mp4", + ) + params, error = _apply_reid_tracker_params("botsort", args, {}) + assert error is not None and "--tracker.reid.model" in error + assert params == {} + + def test_help_lists_reid_flags(self) -> None: + parser = argparse.ArgumentParser() + subparsers = parser.add_subparsers() + add_track_subparser(subparsers) + help_text = subparsers.choices["track"].format_help() + assert "--tracker.reid.enable" in help_text + assert "--tracker.reid.architecture" in help_text + assert "--tracker.appearance_threshold" in help_text diff --git a/uv.lock b/uv.lock index d0f51474a..fdb9024ae 100644 --- a/uv.lock +++ b/uv.lock @@ -158,6 +158,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/0c/37/fb6973edeb700f6e3d6ff222400602ab1830446c25c7b4676d8de93e65b8/backrefs-5.8-py39-none-any.whl", hash = "sha256:a66851e4533fb5b371aa0628e1fee1af05135616b86140c9d787a2ffdf4b8fdc", size = 380336, upload-time = "2025-02-25T16:53:29.858Z" }, ] +[[package]] +name = "beautifulsoup4" +version = "4.15.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "soupsieve" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/43/65/318323f98dbee45d42dff61d8f047181bc6f2268a9068cfad035a46be5af/beautifulsoup4-4.15.0.tar.gz", hash = "sha256:288e3ca7d54b06f2ac191970bc275c1939cb46d450b255bf6718b04aa37ab4f7", size = 632571, upload-time = "2026-06-07T16:44:20.453Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/88/c6/92fcd42f1ba33e1184263f25bfabf3d27c383410470f169e4b8163bf9c17/beautifulsoup4-4.15.0-py3-none-any.whl", hash = "sha256:d6f88de62e1d4e38ecb1077eb9724cd0eff29d2a08ca16a401e9b9e93f117cf9", size = 109924, upload-time = "2026-06-07T16:44:21.566Z" }, +] + [[package]] name = "bitsandbytes" version = "0.47.0" @@ -839,6 +852,22 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ab/6e/81d47999aebc1b155f81eca4477a616a70f238a2549848c38983f3c22a82/ftfy-6.3.1-py3-none-any.whl", hash = "sha256:7c70eb532015cd2f9adb53f101fb6c7945988d023a085d127d1573dc49dd0083", size = 44821, upload-time = "2024-10-26T00:50:33.425Z" }, ] +[[package]] +name = "gdown" +version = "6.1.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "beautifulsoup4" }, + { name = "filelock" }, + { name = "requests", extra = ["socks"] }, + { name = "tqdm" }, + { name = "typing-extensions", marker = "python_full_version < '3.12'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/6a/b5/a45f62f20664031bf74a6aeb6f8d8cd5910e411bf90d756bd6b09bdc6c35/gdown-6.1.0.tar.gz", hash = "sha256:361c6e04c6ca335df50b9d71f40bcfe9ab70fb26a1b0e890a427267781389553", size = 269670, upload-time = "2026-05-30T11:56:21.322Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7d/56/a99f0f159cce5b26d267317d436afee184f45fc7911938757d7cbbd2d10c/gdown-6.1.0-py3-none-any.whl", hash = "sha256:38a36a94275b8272f684db469bbd73b4d1f64cbbc1751bcb993a1b2be8f013c8", size = 19216, upload-time = "2026-05-30T11:56:20.016Z" }, +] + [[package]] name = "ghp-import" version = "2.1.0" @@ -2797,6 +2826,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/bd/24/12818598c362d7f300f18e74db45963dbcb85150324092410c8b49405e42/pyproject_hooks-1.2.0-py3-none-any.whl", hash = "sha256:9e5c6bfa8dcc30091c74b0cf803c81fdd29d94f01992a7707bc97babb1141913", size = 10216, upload-time = "2024-09-29T09:24:11.978Z" }, ] +[[package]] +name = "pysocks" +version = "1.7.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/bd/11/293dd436aea955d45fc4e8a35b6ae7270f5b8e00b53cf6c024c83b657a11/PySocks-1.7.1.tar.gz", hash = "sha256:3f8804571ebe159c380ac6de37643bb4685970655d3bba243530d6558b799aa0", size = 284429, upload-time = "2019-09-20T02:07:35.714Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8d/59/b4572118e098ac8e46e399a1dd0f2d85403ce8bbaad9ec79373ed6badaf9/PySocks-1.7.1-py3-none-any.whl", hash = "sha256:2725bd0a9925919b9b51739eea5f9e2bae91e83288108a9ad338b2e3a4435ee5", size = 16725, upload-time = "2019-09-20T02:06:22.938Z" }, +] + [[package]] name = "pytest" version = "9.1.1" @@ -3262,6 +3300,23 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/95/e4/a3b9480c78cf8ee86626cb06f8d931d74d775897d44201ccb813097ae697/regex-2026.1.15-cp314-cp314t-win_arm64.whl", hash = "sha256:ca89c5e596fc05b015f27561b3793dc2fa0917ea0d7507eebb448efd35274a70", size = 274837, upload-time = "2026-01-14T23:17:23.146Z" }, ] +[[package]] +name = "reid" +version = "0.1.0" +source = { git = "https://github.com/roboflow/re-ID.git?rev=main#d88d77cd7407d3a0158fbcff79eafa4e060234ea" } +dependencies = [ + { name = "gdown" }, + { name = "huggingface-hub" }, + { name = "numpy" }, + { name = "opencv-python" }, + { name = "pillow" }, + { name = "safetensors" }, + { name = "supervision" }, + { name = "timm" }, + { name = "torch" }, + { name = "torchvision" }, +] + [[package]] name = "requests" version = "2.34.2" @@ -3277,6 +3332,11 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a0/f4/c67b0b3f1b9245e8d266f0f112c500d50e5b4e83cb6f3b71b6528104182a/requests-2.34.2-py3-none-any.whl", hash = "sha256:2a0d60c172f83ac6ab31e4554906c0f3b3588d37b5cb939b1c061f4907e278e0", size = 73075, upload-time = "2026-05-14T19:25:26.443Z" }, ] +[package.optional-dependencies] +socks = [ + { name = "pysocks" }, +] + [[package]] name = "requests-file" version = "3.0.1" @@ -3770,6 +3830,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c1/d4/59e74daffcb57a07668852eeeb6035af9f32cbfd7a1d2511f17d2fe6a738/smmap-5.0.3-py3-none-any.whl", hash = "sha256:c106e05d5a61449cf6ba9a1e650227ecfb141590d2a98412103ff35d89fc7b2f", size = 24390, upload-time = "2026-03-09T03:43:24.361Z" }, ] +[[package]] +name = "soupsieve" +version = "2.9" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/80/f1/93422647dd7e461f23d254e6b2bfa687a85b53aeb4903fcdbb74474d4584/soupsieve-2.9.tar.gz", hash = "sha256:acee8417325c5653e1377dc31eccad59eb82cbc65942afe6174c53b3aaad63fc", size = 122122, upload-time = "2026-07-19T01:35:18.425Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7b/d6/3185ab5ad1280319b31986898f3206dd7227cd75e293d4dba2a5e6bf27a0/soupsieve-2.9-py3-none-any.whl", hash = "sha256:a2b2c76d67df2382d245409fd71e321a571717e58463efa32ace87dcadac2c12", size = 37387, upload-time = "2026-07-19T01:35:17.106Z" }, +] + [[package]] name = "sqlalchemy" version = "2.0.49" @@ -4142,6 +4211,9 @@ dependencies = [ detection = [ { name = "inference-models" }, ] +reid = [ + { name = "reid" }, +] tune = [ { name = "optuna" }, ] @@ -4180,12 +4252,13 @@ requires-dist = [ { name = "opencv-python", specifier = ">=4.8.0" }, { name = "optuna", marker = "extra == 'tune'", specifier = ">=3.0.0" }, { name = "pydeprecate", specifier = ">=0.7.0" }, + { name = "reid", marker = "extra == 'reid'", git = "https://github.com/roboflow/re-ID.git?rev=main" }, { name = "requests", specifier = ">=2.28.0" }, { name = "rich", specifier = ">=13.0.0" }, { name = "scipy", specifier = ">=1.13.1" }, { name = "supervision", specifier = ">=0.26.1" }, ] -provides-extras = ["detection", "tune"] +provides-extras = ["detection", "tune", "reid"] [package.metadata.requires-dev] build = [