Refactor/typing tracker (#2131)
* refactor(tracker): improve typing and docstrings * refactor(tracker): improve typing for track objects * refactor(tracker): improve typing in matching and filters * Apply suggestions from review * Refactor: enhance type annotations, assertions, and docstrings across `byte_tracker` for clarity and type safety. * Refactor: replace "Parameters" with "Args" in docstrings for consistency with style guide --------- Co-authored-by: Jirka Borovec <6035284+Borda@users.noreply.github.com> Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
This commit is contained in:
Daniel
GitHub
parent
00c96c03ee
commit
83dfc87ef5
|
|
@ -223,10 +223,6 @@ module = [
|
|||
"supervision.key_points.core",
|
||||
"supervision.key_points.skeletons",
|
||||
"supervision.metrics.utils.utils",
|
||||
"supervision.tracker.byte_tracker.core",
|
||||
"supervision.tracker.byte_tracker.kalman_filter",
|
||||
"supervision.tracker.byte_tracker.matching",
|
||||
"supervision.tracker.byte_tracker.single_object_track",
|
||||
]
|
||||
ignore_errors = true
|
||||
|
||||
|
|
|
|||
|
|
@ -1,4 +1,9 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from typing import cast
|
||||
|
||||
import numpy as np
|
||||
import numpy.typing as npt
|
||||
|
||||
from supervision.detection.core import Detections
|
||||
from supervision.detection.utils.iou_and_nms import box_iou_batch
|
||||
|
|
@ -16,20 +21,20 @@ class ByteTrack:
|
|||
<source src="https://media.roboflow.com/supervision/video-examples/how-to/track-objects/annotate-video-with-traces.mp4" type="video/mp4">
|
||||
</video>
|
||||
|
||||
Parameters:
|
||||
track_activation_threshold (float): Detection confidence threshold
|
||||
Args:
|
||||
track_activation_threshold: Detection confidence threshold
|
||||
for track activation. Increasing track_activation_threshold improves accuracy
|
||||
and stability but might miss true detections. Decreasing it increases
|
||||
completeness but risks introducing noise and instability.
|
||||
lost_track_buffer (int): Number of frames to buffer when a track is lost.
|
||||
lost_track_buffer: Number of frames to buffer when a track is lost.
|
||||
Increasing lost_track_buffer enhances occlusion handling, significantly
|
||||
reducing the likelihood of track fragmentation or disappearance caused
|
||||
by brief detection gaps.
|
||||
minimum_matching_threshold (float): Threshold for matching tracks with detections.
|
||||
minimum_matching_threshold: Threshold for matching tracks with detections.
|
||||
Increasing minimum_matching_threshold improves accuracy but risks fragmentation.
|
||||
Decreasing it improves completeness but risks false positives and drift.
|
||||
frame_rate (int): The frame rate of the video.
|
||||
minimum_consecutive_frames (int): Number of consecutive frames that an object must
|
||||
frame_rate: The frame rate of the video.
|
||||
minimum_consecutive_frames: Number of consecutive frames that an object must
|
||||
be tracked before it is considered a 'valid' track.
|
||||
Increasing minimum_consecutive_frames prevents the creation of accidental tracks from
|
||||
false detection or double detection, but risks missing shorter tracks.
|
||||
|
|
@ -68,7 +73,7 @@ class ByteTrack:
|
|||
detection results.
|
||||
|
||||
Args:
|
||||
detections (Detections): The detections to pass through the tracker.
|
||||
detections: The detections to pass through the tracker.
|
||||
|
||||
Example:
|
||||
```python
|
||||
|
|
@ -101,6 +106,9 @@ class ByteTrack:
|
|||
)
|
||||
```
|
||||
"""
|
||||
if detections.confidence is None:
|
||||
raise ValueError("Detections confidence must be provided for tracking.")
|
||||
|
||||
tensors = np.hstack(
|
||||
(
|
||||
detections.xyxy,
|
||||
|
|
@ -115,7 +123,7 @@ class ByteTrack:
|
|||
|
||||
ious = box_iou_batch(detection_bounding_boxes, track_bounding_boxes)
|
||||
|
||||
iou_costs = 1 - ious
|
||||
iou_costs: npt.NDArray[np.float32] = 1 - ious
|
||||
|
||||
matches, _, _ = matching.linear_assignment(iou_costs, 0.5)
|
||||
detections.tracker_id = np.full(len(detections), -1, dtype=int)
|
||||
|
|
@ -124,7 +132,8 @@ class ByteTrack:
|
|||
tracks[i_track].external_track_id
|
||||
)
|
||||
|
||||
return detections[detections.tracker_id != -1]
|
||||
filtered = detections[detections.tracker_id != -1]
|
||||
return cast(Detections, filtered)
|
||||
|
||||
else:
|
||||
detections = Detections.empty()
|
||||
|
|
@ -148,15 +157,15 @@ class ByteTrack:
|
|||
self.lost_tracks = []
|
||||
self.removed_tracks = []
|
||||
|
||||
def update_with_tensors(self, tensors: np.ndarray) -> list[STrack]:
|
||||
def update_with_tensors(self, tensors: npt.NDArray[np.float32]) -> list[STrack]:
|
||||
"""
|
||||
Updates the tracker with the provided tensors and returns the updated tracks.
|
||||
|
||||
Parameters:
|
||||
Args:
|
||||
tensors: The new tensors to update with.
|
||||
|
||||
Returns:
|
||||
List[STrack]: Updated tracks.
|
||||
Updated tracks.
|
||||
"""
|
||||
self.frame_id += 1
|
||||
activated_starcks = []
|
||||
|
|
@ -195,7 +204,7 @@ class ByteTrack:
|
|||
|
||||
""" Add newly detected tracklets to tracked_stracks"""
|
||||
unconfirmed = []
|
||||
tracked_stracks = [] # type: list[STrack]
|
||||
tracked_stracks: list[STrack] = []
|
||||
|
||||
for track in self.tracked_tracks:
|
||||
if not track.is_activated:
|
||||
|
|
@ -319,9 +328,9 @@ def joint_tracks(
|
|||
Joins two lists of tracks, ensuring that the resulting list does not
|
||||
contain tracks with duplicate internal_track_id values.
|
||||
|
||||
Parameters:
|
||||
track_list_a: First list of tracks (with internal_track_id attribute).
|
||||
track_list_b: Second list of tracks (with internal_track_id attribute).
|
||||
Args:
|
||||
track_list_a: First list of tracks.
|
||||
track_list_b: Second list of tracks.
|
||||
|
||||
Returns:
|
||||
Combined list of tracks from track_list_a and track_list_b
|
||||
|
|
@ -338,15 +347,14 @@ def joint_tracks(
|
|||
return result
|
||||
|
||||
|
||||
def sub_tracks(track_list_a: list[STrack], track_list_b: list[STrack]) -> list[int]:
|
||||
def sub_tracks(track_list_a: list[STrack], track_list_b: list[STrack]) -> list[STrack]:
|
||||
"""
|
||||
Returns a list of tracks from track_list_a after removing any tracks
|
||||
that share the same internal_track_id with tracks in track_list_b.
|
||||
|
||||
Parameters:
|
||||
track_list_a: List of tracks (with internal_track_id attribute).
|
||||
track_list_b: List of tracks (with internal_track_id attribute) to
|
||||
be subtracted from track_list_a.
|
||||
Args:
|
||||
track_list_a: List of tracks.
|
||||
track_list_b: List of tracks to be subtracted from track_list_a.
|
||||
Returns:
|
||||
List of remaining tracks from track_list_a after subtraction.
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -1,4 +1,7 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import numpy as np
|
||||
import numpy.typing as npt
|
||||
import scipy.linalg
|
||||
|
||||
|
||||
|
|
@ -6,40 +9,37 @@ class KalmanFilter:
|
|||
"""
|
||||
A simple Kalman filter for tracking bounding boxes in image space.
|
||||
|
||||
The 8-dimensional state space
|
||||
|
||||
x, y, a, h, vx, vy, va, vh
|
||||
|
||||
contains the bounding box center position (x, y), aspect ratio a, height h,
|
||||
and their respective velocities.
|
||||
The 8-dimensional state space is (x, y, a, h, vx, vy, va, vh), where
|
||||
(x, y) is the bounding box center, a is the aspect ratio (w/h), h is
|
||||
the height, and their respective velocities.
|
||||
|
||||
Object motion follows a constant velocity model. The bounding box location
|
||||
(x, y, a, h) is taken as direct observation of the state space (linear
|
||||
observation model).
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
def __init__(self) -> None:
|
||||
ndim, dt = 4, 1.0
|
||||
|
||||
self._motion_mat = np.eye(2 * ndim, 2 * ndim)
|
||||
self._motion_mat: npt.NDArray[np.float64] = np.eye(2 * ndim, 2 * ndim)
|
||||
for i in range(ndim):
|
||||
self._motion_mat[i, ndim + i] = dt
|
||||
self._update_mat = np.eye(ndim, 2 * ndim)
|
||||
self._std_weight_position = 1.0 / 20
|
||||
self._std_weight_velocity = 1.0 / 160
|
||||
self._update_mat: npt.NDArray[np.float64] = np.eye(ndim, 2 * ndim)
|
||||
|
||||
def initiate(self, measurement: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
|
||||
self._std_weight_position: float = 1.0 / 20
|
||||
self._std_weight_velocity: float = 1.0 / 160
|
||||
|
||||
def initiate(
|
||||
self, measurement: npt.NDArray[np.float32]
|
||||
) -> tuple[npt.NDArray[np.float32], npt.NDArray[np.float32]]:
|
||||
"""
|
||||
Create track from an unassociated measurement.
|
||||
Create track from unassociated measurement.
|
||||
|
||||
Args:
|
||||
measurement (ndarray): Bounding box coordinates (x, y, a, h) with
|
||||
center position (x, y), aspect ratio a, and height h.
|
||||
measurement: The initial measurement vector.
|
||||
|
||||
Returns:
|
||||
Tuple[ndarray, ndarray]: Returns the mean vector (8 dimensional) and
|
||||
covariance matrix (8x8 dimensional) of the new track.
|
||||
Unobserved velocities are initialized to 0 mean.
|
||||
The mean vector and covariance matrix of the new track.
|
||||
"""
|
||||
mean_pos = measurement
|
||||
mean_vel = np.zeros_like(mean_pos)
|
||||
|
|
@ -59,21 +59,17 @@ class KalmanFilter:
|
|||
return mean, covariance
|
||||
|
||||
def predict(
|
||||
self, mean: np.ndarray, covariance: np.ndarray
|
||||
) -> tuple[np.ndarray, np.ndarray]:
|
||||
self, mean: npt.NDArray[np.float32], covariance: npt.NDArray[np.float32]
|
||||
) -> tuple[npt.NDArray[np.float32], npt.NDArray[np.float32]]:
|
||||
"""
|
||||
Run Kalman filter prediction step.
|
||||
|
||||
Args:
|
||||
mean (ndarray): The 8 dimensional mean vector of the object
|
||||
state at the previous time step.
|
||||
covariance (ndarray): The 8x8 dimensional covariance matrix of
|
||||
the object state at the previous time step.
|
||||
mean: The object state mean at the previous time step.
|
||||
covariance: The object state covariance at the previous time step.
|
||||
|
||||
Returns:
|
||||
Tuple[ndarray, ndarray]: Returns the mean vector and
|
||||
covariance matrix of the predicted state.
|
||||
Unobserved velocities are initialized to 0 mean.
|
||||
The mean vector and covariance matrix of the predicted state.
|
||||
"""
|
||||
std_pos = [
|
||||
self._std_weight_position * mean[3],
|
||||
|
|
@ -98,18 +94,17 @@ class KalmanFilter:
|
|||
return mean, covariance
|
||||
|
||||
def project(
|
||||
self, mean: np.ndarray, covariance: np.ndarray
|
||||
) -> tuple[np.ndarray, np.ndarray]:
|
||||
self, mean: npt.NDArray[np.float32], covariance: npt.NDArray[np.float32]
|
||||
) -> tuple[npt.NDArray[np.float32], npt.NDArray[np.float32]]:
|
||||
"""
|
||||
Project state distribution to measurement space.
|
||||
|
||||
Args:
|
||||
mean (ndarray): The state's mean vector (8 dimensional array).
|
||||
covariance (ndarray): The state's covariance matrix (8x8 dimensional).
|
||||
mean: The state's mean vector.
|
||||
covariance: The state's covariance matrix.
|
||||
|
||||
Returns:
|
||||
Tuple[ndarray, ndarray]: Returns the projected mean and
|
||||
covariance matrix of the given state estimate.
|
||||
The projected mean and covariance matrix of the given state estimate.
|
||||
"""
|
||||
std = [
|
||||
self._std_weight_position * mean[3],
|
||||
|
|
@ -126,21 +121,16 @@ class KalmanFilter:
|
|||
return mean, covariance + innovation_cov
|
||||
|
||||
def multi_predict(
|
||||
self, mean: np.ndarray, covariance: np.ndarray
|
||||
) -> tuple[np.ndarray, np.ndarray]:
|
||||
self, mean: npt.NDArray[np.float32], covariance: npt.NDArray[np.float32]
|
||||
) -> tuple[npt.NDArray[np.float32], npt.NDArray[np.float32]]:
|
||||
"""
|
||||
Run Kalman filter prediction step (Vectorized version).
|
||||
|
||||
Args:
|
||||
mean (ndarray): The Nx8 dimensional mean matrix
|
||||
of the object states at the previous time step.
|
||||
covariance (ndarray): The Nx8x8 dimensional covariance matrices
|
||||
of the object states at the previous time step.
|
||||
mean: The object state means at the previous time step.
|
||||
covariance: The object state covariances at the previous time step.
|
||||
|
||||
Returns:
|
||||
Tuple[ndarray, ndarray]: Returns the mean vector and
|
||||
covariance matrix of the predicted state.
|
||||
Unobserved velocities are initialized to 0 mean.
|
||||
The mean vector and covariance matrix of the predicted state.
|
||||
"""
|
||||
std_pos = [
|
||||
self._std_weight_position * mean[:, 3],
|
||||
|
|
@ -168,21 +158,21 @@ class KalmanFilter:
|
|||
return mean, covariance
|
||||
|
||||
def update(
|
||||
self, mean: np.ndarray, covariance: np.ndarray, measurement: np.ndarray
|
||||
) -> tuple[np.ndarray, np.ndarray]:
|
||||
self,
|
||||
mean: npt.NDArray[np.float32],
|
||||
covariance: npt.NDArray[np.float32],
|
||||
measurement: npt.NDArray[np.float32],
|
||||
) -> tuple[npt.NDArray[np.float32], npt.NDArray[np.float32]]:
|
||||
"""
|
||||
Run Kalman filter correction step.
|
||||
|
||||
Args:
|
||||
mean (ndarray): The predicted state's mean vector (8 dimensional).
|
||||
covariance (ndarray): The state's covariance matrix (8x8 dimensional).
|
||||
measurement (ndarray): The 4-dimensional measurement vector (x, y, a, h),
|
||||
where (x, y) is the center position, a the aspect ratio,
|
||||
and h the height of the bounding box.
|
||||
mean: The predicted state's mean vector.
|
||||
covariance: The state's covariance matrix.
|
||||
measurement: The measurement vector.
|
||||
|
||||
Returns:
|
||||
Tuple[ndarray, ndarray]: Returns the measurement-corrected
|
||||
state distribution.
|
||||
The measurement-corrected state distribution.
|
||||
"""
|
||||
projected_mean, projected_cov = self.project(mean, covariance)
|
||||
|
||||
|
|
|
|||
|
|
@ -3,17 +3,18 @@ from __future__ import annotations
|
|||
from typing import TYPE_CHECKING
|
||||
|
||||
import numpy as np
|
||||
import numpy.typing as npt
|
||||
from scipy.optimize import linear_sum_assignment
|
||||
|
||||
from supervision.detection.utils.iou_and_nms import box_iou_batch
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from supervision.tracker.byte_tracker.core import STrack
|
||||
from supervision.tracker.byte_tracker.single_object_track import STrack
|
||||
|
||||
|
||||
def indices_to_matches(
|
||||
cost_matrix: np.ndarray, indices: np.ndarray, thresh: float
|
||||
) -> tuple[np.ndarray, tuple, tuple]:
|
||||
cost_matrix: npt.NDArray[np.float32], indices: npt.NDArray[np.int_], thresh: float
|
||||
) -> tuple[npt.NDArray[np.int_], tuple[int, ...], tuple[int, ...]]:
|
||||
matched_cost = cost_matrix[tuple(zip(*indices))]
|
||||
matched_mask = matched_cost <= thresh
|
||||
|
||||
|
|
@ -24,8 +25,8 @@ def indices_to_matches(
|
|||
|
||||
|
||||
def linear_assignment(
|
||||
cost_matrix: np.ndarray, thresh: float
|
||||
) -> tuple[np.ndarray, tuple[int], tuple[int, int]]:
|
||||
cost_matrix: npt.NDArray[np.float32], thresh: float
|
||||
) -> tuple[npt.NDArray[np.int_], tuple[int, ...], tuple[int, ...]]:
|
||||
if cost_matrix.size == 0:
|
||||
return (
|
||||
np.empty((0, 2), dtype=int),
|
||||
|
|
@ -40,7 +41,10 @@ def linear_assignment(
|
|||
return indices_to_matches(cost_matrix, indices, thresh)
|
||||
|
||||
|
||||
def iou_distance(atracks: list[STrack], btracks: list[STrack]) -> np.ndarray:
|
||||
def iou_distance(
|
||||
atracks: list[STrack] | list[npt.NDArray[np.float32]],
|
||||
btracks: list[STrack] | list[npt.NDArray[np.float32]],
|
||||
) -> npt.NDArray[np.float32]:
|
||||
if (len(atracks) > 0 and isinstance(atracks[0], np.ndarray)) or (
|
||||
len(btracks) > 0 and isinstance(btracks[0], np.ndarray)
|
||||
):
|
||||
|
|
@ -58,7 +62,9 @@ def iou_distance(atracks: list[STrack], btracks: list[STrack]) -> np.ndarray:
|
|||
return cost_matrix
|
||||
|
||||
|
||||
def fuse_score(cost_matrix: np.ndarray, stracks: list[STrack]) -> np.ndarray:
|
||||
def fuse_score(
|
||||
cost_matrix: npt.NDArray[np.float32], stracks: list[STrack]
|
||||
) -> npt.NDArray[np.float32]:
|
||||
if cost_matrix.size == 0:
|
||||
return cost_matrix
|
||||
iou_sim = 1 - cost_matrix
|
||||
|
|
|
|||
|
|
@ -20,7 +20,7 @@ class STrack:
|
|||
def __init__(
|
||||
self,
|
||||
tlwh: npt.NDArray[np.float32],
|
||||
score: npt.NDArray[np.float32],
|
||||
score: float,
|
||||
minimum_consecutive_frames: int,
|
||||
shared_kalman: KalmanFilter,
|
||||
internal_id_counter: IdCounter,
|
||||
|
|
@ -32,12 +32,13 @@ class STrack:
|
|||
self.frame_id = 0
|
||||
|
||||
self._tlwh = np.asarray(tlwh, dtype=np.float32)
|
||||
self.kalman_filter = None
|
||||
self.kalman_filter: KalmanFilter | None = None
|
||||
self.shared_kalman = shared_kalman
|
||||
self.mean, self.covariance = None, None
|
||||
self.mean: npt.NDArray[np.float32] | None = None
|
||||
self.covariance: npt.NDArray[np.float32] | None = None
|
||||
self.is_activated = False
|
||||
|
||||
self.score = score
|
||||
self.score: float = score
|
||||
self.tracklet_len = 0
|
||||
|
||||
self.minimum_consecutive_frames = minimum_consecutive_frames
|
||||
|
|
@ -48,6 +49,9 @@ class STrack:
|
|||
self.external_track_id = self.external_id_counter.NO_ID
|
||||
|
||||
def predict(self) -> None:
|
||||
assert self.mean is not None
|
||||
assert self.covariance is not None
|
||||
assert self.kalman_filter is not None
|
||||
mean_state = self.mean.copy()
|
||||
if self.state != TrackState.Tracked:
|
||||
mean_state[7] = 0
|
||||
|
|
@ -61,6 +65,8 @@ class STrack:
|
|||
multi_mean = []
|
||||
multi_covariance = []
|
||||
for i, st in enumerate(stracks):
|
||||
assert st.mean is not None
|
||||
assert st.covariance is not None
|
||||
multi_mean.append(st.mean.copy())
|
||||
multi_covariance.append(st.covariance)
|
||||
if st.state != TrackState.Tracked:
|
||||
|
|
@ -93,6 +99,9 @@ class STrack:
|
|||
self.start_frame = frame_id
|
||||
|
||||
def re_activate(self, new_track: STrack, frame_id: int) -> None:
|
||||
assert self.kalman_filter is not None
|
||||
assert self.mean is not None
|
||||
assert self.covariance is not None
|
||||
self.mean, self.covariance = self.kalman_filter.update(
|
||||
self.mean, self.covariance, self.tlwh_to_xyah(new_track.tlwh)
|
||||
)
|
||||
|
|
@ -104,12 +113,15 @@ class STrack:
|
|||
|
||||
def update(self, new_track: STrack, frame_id: int) -> None:
|
||||
"""
|
||||
Update a matched track
|
||||
:type new_track: STrack
|
||||
:type frame_id: int
|
||||
:type update_feature: bool
|
||||
:return:
|
||||
Update a matched track.
|
||||
|
||||
Args:
|
||||
new_track: The new track data.
|
||||
frame_id: The current frame ID.
|
||||
"""
|
||||
assert self.kalman_filter is not None
|
||||
assert self.mean is not None
|
||||
assert self.covariance is not None
|
||||
self.frame_id = frame_id
|
||||
self.tracklet_len += 1
|
||||
|
||||
|
|
@ -147,7 +159,7 @@ class STrack:
|
|||
return ret
|
||||
|
||||
@staticmethod
|
||||
def tlwh_to_xyah(tlwh) -> npt.NDArray[np.float32]:
|
||||
def tlwh_to_xyah(tlwh: npt.NDArray[np.float32]) -> npt.NDArray[np.float32]:
|
||||
"""Convert bounding box to format `(center x, center y, aspect ratio,
|
||||
height)`, where the aspect ratio is `width / height`.
|
||||
"""
|
||||
|
|
@ -160,13 +172,13 @@ class STrack:
|
|||
return self.tlwh_to_xyah(self.tlwh)
|
||||
|
||||
@staticmethod
|
||||
def tlbr_to_tlwh(tlbr) -> npt.NDArray[np.float32]:
|
||||
def tlbr_to_tlwh(tlbr: npt.NDArray[np.float32]) -> npt.NDArray[np.float32]:
|
||||
ret = np.asarray(tlbr).copy()
|
||||
ret[2:] -= ret[:2]
|
||||
return ret
|
||||
|
||||
@staticmethod
|
||||
def tlwh_to_tlbr(tlwh) -> npt.NDArray[np.float32]:
|
||||
def tlwh_to_tlbr(tlwh: npt.NDArray[np.float32]) -> npt.NDArray[np.float32]:
|
||||
ret = np.asarray(tlwh).copy()
|
||||
ret[2:] += ret[:2]
|
||||
return ret
|
||||
|
|
|
|||
|
|
@ -1,14 +1,33 @@
|
|||
from __future__ import annotations
|
||||
|
||||
|
||||
class IdCounter:
|
||||
def __init__(self, start_id: int = 0):
|
||||
"""
|
||||
Initialize the ID counter.
|
||||
|
||||
Args:
|
||||
start_id: The starting integer for the counter.
|
||||
|
||||
Raises:
|
||||
ValueError: If start_id is less than or equal to -1.
|
||||
"""
|
||||
self.start_id = start_id
|
||||
if self.start_id <= self.NO_ID:
|
||||
raise ValueError(f"start_id must be greater than {self.NO_ID}")
|
||||
self.reset()
|
||||
|
||||
def reset(self) -> None:
|
||||
"""Reset the counter to the initial start_id."""
|
||||
self._id = self.start_id
|
||||
|
||||
def new_id(self) -> int:
|
||||
"""
|
||||
Get the current ID and increment the counter.
|
||||
|
||||
Returns:
|
||||
The newly assigned ID.
|
||||
"""
|
||||
returned_id = self._id
|
||||
self._id += 1
|
||||
return returned_id
|
||||
|
|
|
|||
Loading…
Reference in New Issue