Reference for ultralytics/trackers/utils/kalman_filter.py

def gating_distance(
    self,
    mean: np.ndarray,
    covariance: np.ndarray,
    measurements: np.ndarray,
    only_position: bool = False,
    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.

    Args:
        mean (ndarray): Mean vector over the state distribution (8 dimensional).
        covariance (ndarray): Covariance of the state distribution (8x8 dimensional).
        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:
        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)
        return np.sum(z * z, axis=0)  # square maha
    else:
        raise ValueError("Invalid distance metric")

def initiate(self, measurement: np.ndarray) -> tuple:
    """
    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.

    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)
    mean = np.r_[mean_pos, mean_vel]

    std = [
        2 * self._std_weight_position * measurement[3],
        2 * self._std_weight_position * measurement[3],
        1e-2,
        2 * self._std_weight_position * measurement[3],
        10 * self._std_weight_velocity * measurement[3],
        10 * self._std_weight_velocity * measurement[3],
        1e-5,
        10 * self._std_weight_velocity * measurement[3],
    ]
    covariance = np.diag(np.square(std))
    return mean, covariance

def multi_predict(self, mean: np.ndarray, covariance: np.ndarray) -> tuple:
    """
    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 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],
        self._std_weight_position * mean[:, 3],
        1e-2 * np.ones_like(mean[:, 3]),
        self._std_weight_position * mean[:, 3],
    ]
    std_vel = [
        self._std_weight_velocity * mean[:, 3],
        self._std_weight_velocity * mean[:, 3],
        1e-5 * np.ones_like(mean[:, 3]),
        self._std_weight_velocity * mean[:, 3],
    ]
    sqr = np.square(np.r_[std_pos, std_vel]).T

    motion_cov = [np.diag(sqr[i]) for i in range(len(mean))]
    motion_cov = np.asarray(motion_cov)

    mean = np.dot(mean, self._motion_mat.T)
    left = np.dot(self._motion_mat, covariance).transpose((1, 0, 2))
    covariance = np.dot(left, self._motion_mat.T) + motion_cov

    return mean, covariance

def predict(self, mean: np.ndarray, covariance: np.ndarray) -> tuple:
    """
    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.

    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],
        self._std_weight_position * mean[3],
        1e-2,
        self._std_weight_position * mean[3],
    ]
    std_vel = [
        self._std_weight_velocity * mean[3],
        self._std_weight_velocity * mean[3],
        1e-5,
        self._std_weight_velocity * mean[3],
    ]
    motion_cov = np.diag(np.square(np.r_[std_pos, std_vel]))

    mean = np.dot(mean, self._motion_mat.T)
    covariance = np.linalg.multi_dot((self._motion_mat, covariance, self._motion_mat.T)) + motion_cov

    return mean, covariance

def project(self, mean: np.ndarray, covariance: np.ndarray) -> tuple:
    """
    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).

    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],
        self._std_weight_position * mean[3],
        1e-1,
        self._std_weight_position * mean[3],
    ]
    innovation_cov = np.diag(np.square(std))

    mean = np.dot(self._update_mat, mean)
    covariance = np.linalg.multi_dot((self._update_mat, covariance, self._update_mat.T))
    return mean, covariance + innovation_cov

def update(self, mean: np.ndarray, covariance: np.ndarray, measurement: np.ndarray) -> tuple:
    """
    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.

    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)

    chol_factor, lower = scipy.linalg.cho_factor(projected_cov, lower=True, check_finite=False)
    kalman_gain = scipy.linalg.cho_solve(
        (chol_factor, lower), np.dot(covariance, self._update_mat.T).T, check_finite=False
    ).T
    innovation = measurement - projected_mean

    new_mean = mean + np.dot(innovation, kalman_gain.T)
    new_covariance = covariance - np.linalg.multi_dot((kalman_gain, projected_cov, kalman_gain.T))
    return new_mean, new_covariance

def initiate(self, measurement: np.ndarray) -> tuple:
    """
    Create track from unassociated measurement.

    Args:
        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.

    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)
    mean = np.r_[mean_pos, mean_vel]

    std = [
        2 * self._std_weight_position * measurement[2],
        2 * self._std_weight_position * measurement[3],
        2 * self._std_weight_position * measurement[2],
        2 * self._std_weight_position * measurement[3],
        10 * self._std_weight_velocity * measurement[2],
        10 * self._std_weight_velocity * measurement[3],
        10 * self._std_weight_velocity * measurement[2],
        10 * self._std_weight_velocity * measurement[3],
    ]
    covariance = np.diag(np.square(std))
    return mean, covariance

def multi_predict(self, mean, covariance) -> tuple:
    """
    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 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.

    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],
        self._std_weight_position * mean[:, 3],
        self._std_weight_position * mean[:, 2],
        self._std_weight_position * mean[:, 3],
    ]
    std_vel = [
        self._std_weight_velocity * mean[:, 2],
        self._std_weight_velocity * mean[:, 3],
        self._std_weight_velocity * mean[:, 2],
        self._std_weight_velocity * mean[:, 3],
    ]
    sqr = np.square(np.r_[std_pos, std_vel]).T

    motion_cov = [np.diag(sqr[i]) for i in range(len(mean))]
    motion_cov = np.asarray(motion_cov)

    mean = np.dot(mean, self._motion_mat.T)
    left = np.dot(self._motion_mat, covariance).transpose((1, 0, 2))
    covariance = np.dot(left, self._motion_mat.T) + motion_cov

    return mean, covariance

def predict(self, mean, covariance) -> tuple:
    """
    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.

    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],
        self._std_weight_position * mean[3],
        self._std_weight_position * mean[2],
        self._std_weight_position * mean[3],
    ]
    std_vel = [
        self._std_weight_velocity * mean[2],
        self._std_weight_velocity * mean[3],
        self._std_weight_velocity * mean[2],
        self._std_weight_velocity * mean[3],
    ]
    motion_cov = np.diag(np.square(np.r_[std_pos, std_vel]))

    mean = np.dot(mean, self._motion_mat.T)
    covariance = np.linalg.multi_dot((self._motion_mat, covariance, self._motion_mat.T)) + motion_cov

    return mean, covariance

def project(self, mean, covariance) -> tuple:
    """
    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).

    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],
        self._std_weight_position * mean[3],
        self._std_weight_position * mean[2],
        self._std_weight_position * mean[3],
    ]
    innovation_cov = np.diag(np.square(std))

    mean = np.dot(self._update_mat, mean)
    covariance = np.linalg.multi_dot((self._update_mat, covariance, self._update_mat.T))
    return mean, covariance + innovation_cov

def update(self, mean, covariance, measurement) -> tuple:
    """
    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, w, h), where (x, y) is the center
            position, w the width, and h the height of the bounding box.

    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)