diff --git a/supervision/tracker/byte_tracker/core.py b/supervision/tracker/byte_tracker/core.py index f9bfe452..d117682f 100644 --- a/supervision/tracker/byte_tracker/core.py +++ b/supervision/tracker/byte_tracker/core.py @@ -230,10 +230,18 @@ class ByteTrack: ) detections = Detections.empty() if len(tracks) > 0: - detections.xyxy = np.array([track.tlbr for track in tracks], dtype=np.float32) - detections.class_id = np.array([int(t.class_ids) for t in tracks], dtype=int) - detections.tracker_id = np.array([int(t.track_id) for t in tracks], dtype=int) - detections.confidence = np.array([t.score for t in tracks], dtype=np.float32) + detections.xyxy = np.array( + [track.tlbr for track in tracks], dtype=np.float32 + ) + detections.class_id = np.array( + [int(t.class_ids) for t in tracks], dtype=int + ) + detections.tracker_id = np.array( + [int(t.track_id) for t in tracks], dtype=int + ) + detections.confidence = np.array( + [t.score for t in tracks], dtype=np.float32 + ) return detections @@ -390,7 +398,9 @@ class ByteTrack: return output_stracks -def joint_tracks(track_list_a: List[STrack], track_list_b: List[STrack]) -> List[STrack]: +def joint_tracks( + track_list_a: List[STrack], track_list_b: List[STrack] +) -> List[STrack]: """ Joins two lists of tracks, ensuring that the resulting list does not contain tracks with duplicate track_id values. @@ -446,7 +456,11 @@ def remove_duplicate_tracks(tracks_a: List, tracks_b: List) -> Tuple[List, List] else: duplicates_a.add(track_index_a) - result_a = [track for index, track in enumerate(tracks_a) if index not in duplicates_a] - result_b = [track for index, track in enumerate(tracks_b) if index not in duplicates_b] + result_a = [ + track for index, track in enumerate(tracks_a) if index not in duplicates_a + ] + result_b = [ + track for index, track in enumerate(tracks_b) if index not in duplicates_b + ] return result_a, result_b diff --git a/supervision/tracker/byte_tracker/kalman_filter.py b/supervision/tracker/byte_tracker/kalman_filter.py index c2a2ecac..9077dc32 100644 --- a/supervision/tracker/byte_tracker/kalman_filter.py +++ b/supervision/tracker/byte_tracker/kalman_filter.py @@ -1,3 +1,5 @@ +from typing import Tuple + import numpy as np import scipy.linalg @@ -16,40 +18,27 @@ class KalmanFilter: 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): ndim, dt = 4, 1.0 - # Create Kalman filter model matrices. self._motion_mat = 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) - - # Motion and observation uncertainty are chosen relative to the current - # state estimate. These weights control the amount of uncertainty in - # the model. This is a bit hacky. self._std_weight_position = 1.0 / 20 self._std_weight_velocity = 1.0 / 160 - def initiate(self, measurement): - """Create track from unassociated measurement. + def initiate(self, measurement: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: + """ + Create track from an unassociated measurement. - Parameters - ---------- - measurement : ndarray - Bounding box coordinates (x, y, a, h) with center position (x, y), - aspect ratio a, and height h. - - Returns - ------- - (ndarray, ndarray) - Returns the mean vector (8 dimensional) and covariance matrix (8x8 - dimensional) of the new track. Unobserved velocities are initialized - to 0 mean. + Args: + measurement (ndarray): Bounding box coordinates (x, y, a, h) with center position (x, y), aspect ratio a, and height h. + 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. """ mean_pos = measurement mean_vel = np.zeros_like(mean_pos) @@ -68,24 +57,16 @@ class KalmanFilter: covariance = np.diag(np.square(std)) return mean, covariance - def predict(self, mean, covariance): - """Run Kalman filter prediction step. + def predict(self, mean: np.ndarray, covariance: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: + """ + Run Kalman filter prediction step. - Parameters - ---------- - 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. - - Returns - ------- - (ndarray, ndarray) - Returns the mean vector and covariance matrix of the predicted - state. Unobserved velocities are initialized to 0 mean. + 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. + Returns: + Tuple[ndarray, ndarray]: Returns the mean vector and covariance matrix of the predicted state. Unobserved velocities are initialized to 0 mean. """ std_pos = [ self._std_weight_position * mean[3], @@ -101,7 +82,6 @@ class KalmanFilter: ] motion_cov = np.diag(np.square(np.r_[std_pos, std_vel])) - # mean = np.dot(self._motion_mat, mean) mean = np.dot(mean, self._motion_mat.T) covariance = ( np.linalg.multi_dot((self._motion_mat, covariance, self._motion_mat.T)) @@ -110,22 +90,16 @@ class KalmanFilter: return mean, covariance - def project(self, mean, covariance): - """Project state distribution to measurement space. + def project(self, mean: np.ndarray, covariance: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: + """ + Project state distribution to measurement space. - Parameters - ---------- - mean : ndarray - The state's mean vector (8 dimensional array). - covariance : ndarray - The state's covariance matrix (8x8 dimensional). - - Returns - ------- - (ndarray, ndarray) - Returns the projected mean and covariance matrix of the given state - estimate. + Args: + mean (ndarray): The state's mean vector (8 dimensional array). + covariance (ndarray): The state's covariance matrix (8x8 dimensional). + Returns: + Tuple[ndarray, ndarray]: Returns the projected mean and covariance matrix of the given state estimate. """ std = [ self._std_weight_position * mean[3], @@ -141,21 +115,16 @@ class KalmanFilter: ) return mean, covariance + innovation_cov - def multi_predict(self, mean, covariance) -> [np.ndarray, np.ndarray]: - """Run Kalman filter prediction step (Vectorized version). - Parameters - ---------- - mean : ndarray - The Nx8 dimensional mean matrix of the object states at the previous - time step. - covariance : ndarray - The Nx8x8 dimensional covariance matrics of the object states at the - previous time step. - Returns - ------- - (ndarray, ndarray) - Returns the mean vector and covariance matrix of the predicted - state. Unobserved velocities are initialized to 0 mean. + def multi_predict(self, mean: np.ndarray, covariance: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: + """ + 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. + + Returns: + Tuple[ndarray, ndarray]: Returns the mean vector and covariance matrix of the predicted state. Unobserved velocities are initialized to 0 mean. """ std_pos = [ self._std_weight_position * mean[:, 3], @@ -182,25 +151,19 @@ class KalmanFilter: return mean, covariance - def update(self, mean, covariance, measurement): - """Run Kalman filter correction step. + def update( + self, mean: np.ndarray, covariance: np.ndarray, measurement: np.ndarray + ) -> Tuple[np.ndarray, np.ndarray]: + """ + Run Kalman filter correction step. - Parameters - ---------- - 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. - - Returns - ------- - (ndarray, ndarray) - Returns the measurement-corrected state distribution. + 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. + Returns: + Tuple[ndarray, ndarray]: Returns the measurement-corrected state distribution. """ projected_mean, projected_cov = self.project(mean, covariance) @@ -219,48 +182,3 @@ class KalmanFilter: (kalman_gain, projected_cov, kalman_gain.T) ) return new_mean, new_covariance - - def gating_distance( - self, mean, covariance, measurements, only_position=False, metric="maha" - ): - """Compute gating distance between state distribution and measurements. - A suitable distance threshold can be obtained from `chi2inv95`. If - `only_position` is False, the chi-square distribution has 4 degrees of - freedom, otherwise 2. - Parameters - ---------- - mean : ndarray - Mean vector over the state distribution (8 dimensional). - covariance : ndarray - Covariance of the state distribution (8x8 dimensional). - measurements : ndarray - An Nx4 dimensional matrix of N measurements, each in - format (x, y, a, h) where (x, y) is the bounding box center - position, a the aspect ratio, and h the height. - only_position : Optional[bool] - If True, distance computation is done with respect to the bounding - box center position only. - Returns - ------- - ndarray - Returns an array of length N, where the i-th element contains the - squared Mahalanobis distance between (mean, covariance) and - `measurements[i]`. - """ - mean, covariance = self.project(mean, covariance) - if only_position: - mean, covariance = mean[:2], covariance[:2, :2] - measurements = measurements[:, :2] - - d = measurements - mean - if metric == "gaussian": - return np.sum(d * d, axis=1) - elif metric == "maha": - cholesky_factor = np.linalg.cholesky(covariance) - z = scipy.linalg.solve_triangular( - cholesky_factor, d.T, lower=True, check_finite=False, overwrite_b=True - ) - squared_maha = np.sum(z * z, axis=0) - return squared_maha - else: - raise ValueError("invalid distance metric")