Multi object tracking

This page contains advanced examples for the tno.quantum.problems.mot package.

Requirements

Install the following dependencies to run the examples below:

pip install tno.quantum.problems.mot
pip install ultralytics

Example: getting detections with ultralytics

Starting from a video file, first generate a detection file compatible with ultralytics. This can be done, for example, with a pre-trained YOLO model, as shown below

import csv
from ultralytics import YOLO

model = YOLO("path/to/model")
video_path = "path/to/video"
detections_csv_path = "path/to/output.csv"

# Run inference frame-by-frame
results = model(
    video_path,
    stream=True,  # important for video
    conf=0.25,
)

with open(detections_csv_path, mode="w", newline="") as f:
    writer = csv.writer(f)
    for frame_idx, r in enumerate(results):
        if r.boxes is None:
            continue

        for box in r.boxes:
            cls_id = int(box.cls[0])
            conf = float(box.conf[0])
            x, y, w, h = box.xywhn[0].tolist()

            writer.writerow(
                [
                    frame_idx,
                    cls_id,
                    f"{x:.6f}",
                    f"{y:.6f}",
                    f"{w:.6f}",
                    f"{h:.6f}",
                    f"{conf:.6f}",
                ]
            )

Example: re-id module

The following example shows how to run the tracking algorithm of ultralytics and the re-id module on top to fix an ID switch.

import importlib.resources
import pandas as pd

from ultralytics.trackers import BOTSORT
from ultralytics.utils import YAML, IterableSimpleNamespace

from tno.quantum.problems.mot import (
    DetectionSequence,
    ultralytics_track,
    FlowReID,
)
from tno.quantum.problems.mot.dataset import (
    make_gif_with_detections,
)

# Load detections from CSV file
detections_csv_path = "path/to/detections.csv"
detections = pd.read_csv(detections_csv_path, header=None)
detections = DetectionSequence.from_dataframe(detections)

# Run ultralytics tracking algorithm
tracker_path = str(
    importlib.resources.files("ultralytics.cfg") / "trackers" / "botsort.yaml"
)
cfg = IterableSimpleNamespace(**YAML.load(tracker_path))
tracker = BOTSORT(cfg)
detections = ultralytics_track(detections, tracker)

# Select a snippet that contains an ID loss
flow_reid_obj = FlowReID()
first_frame, last_frame = flow_reid_obj.pick_sequences(detections)[0]
detections_snippet = detections[first_frame:last_frame]
reid_detections_snippet = flow_reid_obj.run_classic(detections_snippet)

# Create visualization
video_path = "path/to/video"
make_gif_with_detections(
    detections_snippet,
    input_filename=video_path,
    output_filename="before.gif",
    fps=20,
    first_frame=first_frame,
)
make_gif_with_detections(
    reid_detections_snippet,
    input_filename=video_path,
    output_filename="after.gif",
    fps=20,
    first_frame=first_frame,
)

This results in the following video snippets.

Instead of selecting a frame sequence manually, use the pick_sequences method to find frame sequences where an ID loss or ID switch is suspected.

import importlib.resources

import pandas as pd
from ultralytics.trackers import BOTSORT
from ultralytics.utils import YAML, IterableSimpleNamespace

from tno.quantum.problems.mot import (
    DetectionSequence,
    FlowReID,
    ultralytics_track,
)
from tno.quantum.problems.mot.dataset import (
    make_video_with_detections,
)

# Load detections from CSV file
detections_csv_path = "path/to/detections.csv"
detections = pd.read_csv(detections_csv_path, header=None)
detections = DetectionSequence.from_dataframe(detections)

# Run ultralytics tracking algorithm; pick_sequences needs track IDs
tracker_path = str(
    importlib.resources.files("ultralytics.cfg") / "trackers" / "botsort.yaml"
)
cfg = IterableSimpleNamespace(**YAML.load(tracker_path))
detections = ultralytics_track(detections, BOTSORT(cfg))

# Pick sequences
flow_reid_obj = FlowReID()
sequences = flow_reid_obj.pick_sequences(detections)
print(sequences)
print(f"Found {len(sequences)} frame sequences for re-id.")

# To run re-id continuously on each of the detected sequences, use
# run_continuous_reid.
reid_detections = flow_reid_obj.run_continuous_reid(detections)

# Save the video after re-id
make_video_with_detections(
    reid_detections,
    input_filename="path/to/video",
    output_filename="after.avi",
    fps=50,
)