ultralytics 8.1.29__py3-none-any.whl → 8.3.63__py3-none-any.whl

ultralytics/trackers/utils/kalman_filter.py

@@ -1,4 +1,4 @@
-# Ultralytics YOLO 🚀, AGPL-3.0 license
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
 import numpy as np
 import scipy.linalg
@@ -6,17 +6,49 @@ import scipy.linalg
 
 class KalmanFilterXYAH:
     """
-    For bytetrack. 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.
-
-    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).
+    A KalmanFilterXYAH class for tracking bounding boxes in image space using a Kalman filter.
+
+    Implements 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. Object motion follows a constant velocity model, and bounding box location (x, y, a, h) is
+    taken as a direct observation of the state space (linear observation model).
+
+    Attributes:
+        _motion_mat (np.ndarray): The motion matrix for the Kalman filter.
+        _update_mat (np.ndarray): The update matrix for the Kalman filter.
+        _std_weight_position (float): Standard deviation weight for position.
+        _std_weight_velocity (float): Standard deviation weight for velocity.
+
+    Methods:
+        initiate: Creates a track from an unassociated measurement.
+        predict: Runs the Kalman filter prediction step.
+        project: Projects the state distribution to measurement space.
+        multi_predict: Runs the Kalman filter prediction step (vectorized version).
+        update: Runs the Kalman filter correction step.
+        gating_distance: Computes the gating distance between state distribution and measurements.
+
+    Examples:
+        Initialize the Kalman filter and create a track from a measurement
+        >>> kf = KalmanFilterXYAH()
+        >>> measurement = np.array([100, 200, 1.5, 50])
+        >>> mean, covariance = kf.initiate(measurement)
+        >>> print(mean)
+        >>> print(covariance)
     """
 
     def __init__(self):
-        """Initialize Kalman filter model matrices with motion and observation uncertainty weights."""
+        """
+        Initialize Kalman filter model matrices with motion and observation uncertainty weights.
+
+        The Kalman filter is initialized with an 8-dimensional state space (x, y, a, h, vx, vy, va, vh), where (x, y)
+        represents the bounding box center position, 'a' is the aspect ratio, 'h' is the height, and their respective
+        velocities are (vx, vy, va, vh). The filter uses a constant velocity model for object motion and a linear
+        observation model for bounding box location.
+
+        Examples:
+            Initialize a Kalman filter for tracking:
+            >>> kf = KalmanFilterXYAH()
+        """
         ndim, dt = 4, 1.0
 
         # Create Kalman filter model matrices
@@ -32,15 +64,20 @@ class KalmanFilterXYAH:
 
     def initiate(self, measurement: np.ndarray) -> tuple:
         """
-        Create track from unassociated measurement.
+        Create a track from an unassociated measurement.
 
         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.
+            (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.
+
+        Examples:
+            >>> kf = KalmanFilterXYAH()
+            >>> measurement = np.array([100, 50, 1.5, 200])
+            >>> mean, covariance = kf.initiate(measurement)
         """
         mean_pos = measurement
         mean_vel = np.zeros_like(mean_pos)
@@ -64,12 +101,18 @@ class KalmanFilterXYAH:
         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 (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.
+
+        Examples:
+            >>> kf = KalmanFilterXYAH()
+            >>> mean = np.array([0, 0, 1, 1, 0, 0, 0, 0])
+            >>> covariance = np.eye(8)
+            >>> predicted_mean, predicted_covariance = kf.predict(mean, covariance)
         """
         std_pos = [
             self._std_weight_position * mean[3],
@@ -100,6 +143,12 @@ class KalmanFilterXYAH:
 
         Returns:
             (tuple[ndarray, ndarray]): Returns the projected mean and covariance matrix of the given state estimate.
+
+        Examples:
+            >>> kf = KalmanFilterXYAH()
+            >>> mean = np.array([0, 0, 1, 1, 0, 0, 0, 0])
+            >>> covariance = np.eye(8)
+            >>> projected_mean, projected_covariance = kf.project(mean, covariance)
         """
         std = [
             self._std_weight_position * mean[3],
@@ -115,15 +164,21 @@ class KalmanFilterXYAH:
 
     def multi_predict(self, mean: np.ndarray, covariance: np.ndarray) -> tuple:
         """
-        Run Kalman filter prediction step (Vectorized version).
+        Run Kalman filter prediction step for multiple object states (Vectorized version).
 
         Args:
             mean (ndarray): The Nx8 dimensional mean matrix of the object states at the previous time step.
             covariance (ndarray): The Nx8x8 covariance matrix 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.
+            (tuple[ndarray, ndarray]): Returns the mean matrix and covariance matrix of the predicted states.
+                The mean matrix has shape (N, 8) and the covariance matrix has shape (N, 8, 8). Unobserved velocities
+                are initialized to 0 mean.
+
+        Examples:
+            >>> mean = np.random.rand(10, 8)  # 10 object states
+            >>> covariance = np.random.rand(10, 8, 8)  # Covariance matrices for 10 object states
+            >>> predicted_mean, predicted_covariance = kalman_filter.multi_predict(mean, covariance)
         """
         std_pos = [
             self._std_weight_position * mean[:, 3],
@@ -160,6 +215,13 @@ class KalmanFilterXYAH:
 
         Returns:
             (tuple[ndarray, ndarray]): Returns the measurement-corrected state distribution.
+
+        Examples:
+            >>> kf = KalmanFilterXYAH()
+            >>> mean = np.array([0, 0, 1, 1, 0, 0, 0, 0])
+            >>> covariance = np.eye(8)
+            >>> measurement = np.array([1, 1, 1, 1])
+            >>> new_mean, new_covariance = kf.update(mean, covariance, measurement)
         """
         projected_mean, projected_cov = self.project(mean, covariance)
 
@@ -182,23 +244,31 @@ class KalmanFilterXYAH:
         metric: str = "maha",
     ) -> np.ndarray:
         """
-        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.
+        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.
 
         Args:
             mean (ndarray): Mean vector over the state distribution (8 dimensional).
             covariance (ndarray): Covariance of the state distribution (8x8 dimensional).
-            measurements (ndarray): An Nx4 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 (bool, optional): If True, distance computation is done with respect to the bounding box
-                center position only. Defaults to False.
-            metric (str, optional): The metric to use for calculating the distance. Options are 'gaussian' for the
-                squared Euclidean distance and 'maha' for the squared Mahalanobis distance. Defaults to 'maha'.
+            measurements (ndarray): An (N, 4) 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 (bool): If True, distance computation is done with respect to box center position only.
+            metric (str): The metric to use for calculating the distance. Options are 'gaussian' for the squared
+                Euclidean distance and 'maha' for the squared Mahalanobis distance.
 
         Returns:
             (np.ndarray): Returns an array of length N, where the i-th element contains the squared distance between
                 (mean, covariance) and `measurements[i]`.
+
+        Examples:
+            Compute gating distance using Mahalanobis metric:
+            >>> kf = KalmanFilterXYAH()
+            >>> mean = np.array([0, 0, 1, 1, 0, 0, 0, 0])
+            >>> covariance = np.eye(8)
+            >>> measurements = np.array([[1, 1, 1, 1], [2, 2, 1, 1]])
+            >>> distances = kf.gating_distance(mean, covariance, measurements, only_position=False, metric="maha")
         """
         mean, covariance = self.project(mean, covariance)
         if only_position:
@@ -218,13 +288,33 @@ class KalmanFilterXYAH:
 
 class KalmanFilterXYWH(KalmanFilterXYAH):
     """
-    For BoT-SORT. A simple Kalman filter for tracking bounding boxes in image space.
-
-    The 8-dimensional state space (x, y, w, h, vx, vy, vw, vh) contains the bounding box center position (x, y), width
-    w, height h, and their respective velocities.
+    A KalmanFilterXYWH class for tracking bounding boxes in image space using a Kalman filter.
 
-    Object motion follows a constant velocity model. The bounding box location (x, y, w, h) is taken as direct
+    Implements a Kalman filter for tracking bounding boxes with state space (x, y, w, h, vx, vy, vw, vh), where
+    (x, y) is the center position, w is the width, h is the height, and vx, vy, vw, vh are their respective velocities.
+    The object motion follows a constant velocity model, and the bounding box location (x, y, w, h) is taken as a direct
     observation of the state space (linear observation model).
+
+    Attributes:
+        _motion_mat (np.ndarray): The motion matrix for the Kalman filter.
+        _update_mat (np.ndarray): The update matrix for the Kalman filter.
+        _std_weight_position (float): Standard deviation weight for position.
+        _std_weight_velocity (float): Standard deviation weight for velocity.
+
+    Methods:
+        initiate: Creates a track from an unassociated measurement.
+        predict: Runs the Kalman filter prediction step.
+        project: Projects the state distribution to measurement space.
+        multi_predict: Runs the Kalman filter prediction step in a vectorized manner.
+        update: Runs the Kalman filter correction step.
+
+    Examples:
+        Create a Kalman filter and initialize a track
+        >>> kf = KalmanFilterXYWH()
+        >>> measurement = np.array([100, 50, 20, 40])
+        >>> mean, covariance = kf.initiate(measurement)
+        >>> print(mean)
+        >>> print(covariance)
     """
 
     def initiate(self, measurement: np.ndarray) -> tuple:
@@ -235,8 +325,24 @@ class KalmanFilterXYWH(KalmanFilterXYAH):
             measurement (ndarray): Bounding box coordinates (x, y, w, h) with center position (x, y), width, and height.
 
         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.
+            (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.
+
+        Examples:
+            >>> kf = KalmanFilterXYWH()
+            >>> measurement = np.array([100, 50, 20, 40])
+            >>> mean, covariance = kf.initiate(measurement)
+            >>> print(mean)
+            [100.  50.  20.  40.   0.   0.   0.   0.]
+            >>> print(covariance)
+            [[ 4.  0.  0.  0.  0.  0.  0.  0.]
+             [ 0.  4.  0.  0.  0.  0.  0.  0.]
+             [ 0.  0.  4.  0.  0.  0.  0.  0.]
+             [ 0.  0.  0.  4.  0.  0.  0.  0.]
+             [ 0.  0.  0.  0.  0.25  0.  0.  0.]
+             [ 0.  0.  0.  0.  0.  0.25  0.  0.]
+             [ 0.  0.  0.  0.  0.  0.  0.25  0.]
+             [ 0.  0.  0.  0.  0.  0.  0.  0.25]]
         """
         mean_pos = measurement
         mean_vel = np.zeros_like(mean_pos)
@@ -260,12 +366,18 @@ class KalmanFilterXYWH(KalmanFilterXYAH):
         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 (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.
+
+        Examples:
+            >>> kf = KalmanFilterXYWH()
+            >>> mean = np.array([0, 0, 1, 1, 0, 0, 0, 0])
+            >>> covariance = np.eye(8)
+            >>> predicted_mean, predicted_covariance = kf.predict(mean, covariance)
         """
         std_pos = [
             self._std_weight_position * mean[2],
@@ -296,6 +408,12 @@ class KalmanFilterXYWH(KalmanFilterXYAH):
 
         Returns:
             (tuple[ndarray, ndarray]): Returns the projected mean and covariance matrix of the given state estimate.
+
+        Examples:
+            >>> kf = KalmanFilterXYWH()
+            >>> mean = np.array([0, 0, 1, 1, 0, 0, 0, 0])
+            >>> covariance = np.eye(8)
+            >>> projected_mean, projected_cov = kf.project(mean, covariance)
         """
         std = [
             self._std_weight_position * mean[2],
@@ -320,6 +438,12 @@ class KalmanFilterXYWH(KalmanFilterXYAH):
         Returns:
             (tuple[ndarray, ndarray]): Returns the mean vector and covariance matrix of the predicted state. Unobserved
                 velocities are initialized to 0 mean.
+
+        Examples:
+            >>> mean = np.random.rand(5, 8)  # 5 objects with 8-dimensional state vectors
+            >>> covariance = np.random.rand(5, 8, 8)  # 5 objects with 8x8 covariance matrices
+            >>> kf = KalmanFilterXYWH()
+            >>> predicted_mean, predicted_covariance = kf.multi_predict(mean, covariance)
         """
         std_pos = [
             self._std_weight_position * mean[:, 2],
@@ -356,5 +480,12 @@ class KalmanFilterXYWH(KalmanFilterXYAH):
 
         Returns:
             (tuple[ndarray, ndarray]): Returns the measurement-corrected state distribution.
+
+        Examples:
+            >>> kf = KalmanFilterXYWH()
+            >>> mean = np.array([0, 0, 1, 1, 0, 0, 0, 0])
+            >>> covariance = np.eye(8)
+            >>> measurement = np.array([0.5, 0.5, 1.2, 1.2])
+            >>> new_mean, new_covariance = kf.update(mean, covariance, measurement)
         """
         return super().update(mean, covariance, measurement)

ultralytics/trackers/utils/matching.py

@@ -1,10 +1,10 @@
-# Ultralytics YOLO 🚀, AGPL-3.0 license
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
 import numpy as np
 import scipy
 from scipy.spatial.distance import cdist
 
-from ultralytics.utils.metrics import bbox_ioa, batch_probiou
+from ultralytics.utils.metrics import batch_probiou, bbox_ioa
 
 try:
     import lap  # for linear_assignment
@@ -13,26 +13,29 @@ try:
 except (ImportError, AssertionError, AttributeError):
     from ultralytics.utils.checks import check_requirements
 
-    check_requirements("lapx>=0.5.2")  # update to lap package from https://github.com/rathaROG/lapx
+    check_requirements("lap>=0.5.12")  # https://github.com/gatagat/lap
     import lap
 
 
 def linear_assignment(cost_matrix: np.ndarray, thresh: float, use_lap: bool = True) -> tuple:
     """
-    Perform linear assignment using scipy or lap.lapjv.
+    Perform linear assignment using either the scipy or lap.lapjv method.
 
     Args:
-        cost_matrix (np.ndarray): The matrix containing cost values for assignments.
+        cost_matrix (np.ndarray): The matrix containing cost values for assignments, with shape (N, M).
         thresh (float): Threshold for considering an assignment valid.
-        use_lap (bool, optional): Whether to use lap.lapjv. Defaults to True.
+        use_lap (bool): Use lap.lapjv for the assignment. If False, scipy.optimize.linear_sum_assignment is used.
 
     Returns:
-        Tuple with:
-            - matched indices
-            - unmatched indices from 'a'
-            - unmatched indices from 'b'
+        matched_indices (np.ndarray): Array of matched indices of shape (K, 2), where K is the number of matches.
+        unmatched_a (np.ndarray): Array of unmatched indices from the first set, with shape (L,).
+        unmatched_b (np.ndarray): Array of unmatched indices from the second set, with shape (M,).
+
+    Examples:
+        >>> cost_matrix = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
+        >>> thresh = 5.0
+        >>> matched_indices, unmatched_a, unmatched_b = linear_assignment(cost_matrix, thresh, use_lap=True)
     """
-
     if cost_matrix.size == 0:
         return np.empty((0, 2), dtype=int), tuple(range(cost_matrix.shape[0])), tuple(range(cost_matrix.shape[1]))
 
@@ -68,8 +71,13 @@ def iou_distance(atracks: list, btracks: list) -> np.ndarray:
 
     Returns:
         (np.ndarray): Cost matrix computed based on IoU.
-    """
 
+    Examples:
+        Compute IoU distance between two sets of tracks
+        >>> atracks = [np.array([0, 0, 10, 10]), np.array([20, 20, 30, 30])]
+        >>> btracks = [np.array([5, 5, 15, 15]), np.array([25, 25, 35, 35])]
+        >>> cost_matrix = iou_distance(atracks, btracks)
+    """
     if atracks and isinstance(atracks[0], np.ndarray) or btracks and isinstance(btracks[0], np.ndarray):
         atlbrs = atracks
         btlbrs = btracks
@@ -98,14 +106,20 @@ def embedding_distance(tracks: list, detections: list, metric: str = "cosine") -
     Compute distance between tracks and detections based on embeddings.
 
     Args:
-        tracks (list[STrack]): List of tracks.
-        detections (list[BaseTrack]): List of detections.
-        metric (str, optional): Metric for distance computation. Defaults to 'cosine'.
+        tracks (list[STrack]): List of tracks, where each track contains embedding features.
+        detections (list[BaseTrack]): List of detections, where each detection contains embedding features.
+        metric (str): Metric for distance computation. Supported metrics include 'cosine', 'euclidean', etc.
 
     Returns:
-        (np.ndarray): Cost matrix computed based on embeddings.
+        (np.ndarray): Cost matrix computed based on embeddings with shape (N, M), where N is the number of tracks
+            and M is the number of detections.
+
+    Examples:
+        Compute the embedding distance between tracks and detections using cosine metric
+        >>> tracks = [STrack(...), STrack(...)]  # List of track objects with embedding features
+        >>> detections = [BaseTrack(...), BaseTrack(...)]  # List of detection objects with embedding features
+        >>> cost_matrix = embedding_distance(tracks, detections, metric="cosine")
     """
-
     cost_matrix = np.zeros((len(tracks), len(detections)), dtype=np.float32)
     if cost_matrix.size == 0:
         return cost_matrix
@@ -122,13 +136,18 @@ def fuse_score(cost_matrix: np.ndarray, detections: list) -> np.ndarray:
     Fuses cost matrix with detection scores to produce a single similarity matrix.
 
     Args:
-        cost_matrix (np.ndarray): The matrix containing cost values for assignments.
-        detections (list[BaseTrack]): List of detections with scores.
+        cost_matrix (np.ndarray): The matrix containing cost values for assignments, with shape (N, M).
+        detections (list[BaseTrack]): List of detections, each containing a score attribute.
 
     Returns:
-        (np.ndarray): Fused similarity matrix.
-    """
+        (np.ndarray): Fused similarity matrix with shape (N, M).
 
+    Examples:
+        Fuse a cost matrix with detection scores
+        >>> cost_matrix = np.random.rand(5, 10)  # 5 tracks and 10 detections
+        >>> detections = [BaseTrack(score=np.random.rand()) for _ in range(10)]
+        >>> fused_matrix = fuse_score(cost_matrix, detections)
+    """
     if cost_matrix.size == 0:
         return cost_matrix
     iou_sim = 1 - cost_matrix