diff --git a/pyproject.toml b/pyproject.toml index 35298577..55dc2681 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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 diff --git a/src/supervision/tracker/byte_tracker/core.py b/src/supervision/tracker/byte_tracker/core.py index 52bdef0e..7c024887 100644 --- a/src/supervision/tracker/byte_tracker/core.py +++ b/src/supervision/tracker/byte_tracker/core.py @@ -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: - 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. """ diff --git a/src/supervision/tracker/byte_tracker/kalman_filter.py b/src/supervision/tracker/byte_tracker/kalman_filter.py index 5b678ec5..fbaf779c 100644 --- a/src/supervision/tracker/byte_tracker/kalman_filter.py +++ b/src/supervision/tracker/byte_tracker/kalman_filter.py @@ -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) diff --git a/src/supervision/tracker/byte_tracker/matching.py b/src/supervision/tracker/byte_tracker/matching.py index bdd8196e..9d2c7dc6 100644 --- a/src/supervision/tracker/byte_tracker/matching.py +++ b/src/supervision/tracker/byte_tracker/matching.py @@ -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 diff --git a/src/supervision/tracker/byte_tracker/single_object_track.py b/src/supervision/tracker/byte_tracker/single_object_track.py index c7573300..ce460390 100644 --- a/src/supervision/tracker/byte_tracker/single_object_track.py +++ b/src/supervision/tracker/byte_tracker/single_object_track.py @@ -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 diff --git a/src/supervision/tracker/byte_tracker/utils.py b/src/supervision/tracker/byte_tracker/utils.py index cd2a1036..3404f809 100644 --- a/src/supervision/tracker/byte_tracker/utils.py +++ b/src/supervision/tracker/byte_tracker/utils.py @@ -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