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, frame_rate=30)
detections = ultralytics_track(detections, tracker)

# Select a snippet that contains an ID loss
first_frame = 253  # some frame
last_frame = first_frame + 50
detections_snippet = detections[first_frame:last_frame]
flow_reid_obj = FlowReID()
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 pandas as pd
from tno.quantum.problems.mot import (
    DetectionSequence,
    FlowReID,
)
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)

# 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="example_video.avi",
    output_filename="after.avi",
    fps=50,
)