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 (np.ndarray): Mean vector over the state distribution (8 dimensional).
        covariance (np.ndarray): Covariance of the state distribution (8x8 dimensional).
        measurements (np.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):
    """
    Create a track from an unassociated measurement.

    Args:
        measurement (np.ndarray): Bounding box coordinates (x, y, a, h) with center position (x, y), aspect ratio a,
            and height h.

    Returns:
        (np.ndarray): Mean vector (8-dimensional) of the new track. Unobserved velocities are initialized to 0 mean.
        (np.ndarray): Covariance matrix (8x8 dimensional) of the new track.

    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):
    """
    Run Kalman filter prediction step for multiple object states (Vectorized version).

    Args:
        mean (np.ndarray): The Nx8 dimensional mean matrix of the object states at the previous time step.
        covariance (np.ndarray): The Nx8x8 covariance matrix of the object states at the previous time step.

    Returns:
        (np.ndarray): Mean matrix of the predicted states with shape (N, 8).
        (np.ndarray): Covariance matrix of the predicted states with shape (N, 8, 8).

    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):
    """
    Run Kalman filter prediction step.

    Args:
        mean (np.ndarray): The 8-dimensional mean vector of the object state at the previous time step.
        covariance (np.ndarray): The 8x8-dimensional covariance matrix of the object state at the previous time step.

    Returns:
        (np.ndarray): Mean vector of the predicted state. Unobserved velocities are initialized to 0 mean.
        (np.ndarray): Covariance matrix of the predicted state.

    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):
    """
    Project state distribution to measurement space.

    Args:
        mean (np.ndarray): The state's mean vector (8 dimensional array).
        covariance (np.ndarray): The state's covariance matrix (8x8 dimensional).

    Returns:
        (np.ndarray): Projected mean of the given state estimate.
        (np.ndarray): Projected 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):
    """
    Run Kalman filter correction step.

    Args:
        mean (np.ndarray): The predicted state's mean vector (8 dimensional).
        covariance (np.ndarray): The state's covariance matrix (8x8 dimensional).
        measurement (np.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:
        (np.ndarray): Measurement-corrected state mean.
        (np.ndarray): Measurement-corrected state covariance.

    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):
    """
    Create track from unassociated measurement.

    Args:
        measurement (np.ndarray): Bounding box coordinates (x, y, w, h) with center position (x, y), width, and height.

    Returns:
        (np.ndarray): Mean vector (8 dimensional) of the new track. Unobserved velocities are initialized to 0 mean.
        (np.ndarray): Covariance matrix (8x8 dimensional) of the new track.

    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):
    """
    Run Kalman filter prediction step (Vectorized version).

    Args:
        mean (np.ndarray): The Nx8 dimensional mean matrix of the object states at the previous time step.
        covariance (np.ndarray): The Nx8x8 covariance matrix of the object states at the previous time step.

    Returns:
        (np.ndarray): Mean matrix of the predicted states with shape (N, 8).
        (np.ndarray): Covariance matrix of the predicted states with shape (N, 8, 8).

    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):
    """
    Run Kalman filter prediction step.

    Args:
        mean (np.ndarray): The 8-dimensional mean vector of the object state at the previous time step.
        covariance (np.ndarray): The 8x8-dimensional covariance matrix of the object state at the previous time step.

    Returns:
        (np.ndarray): Mean vector of the predicted state. Unobserved velocities are initialized to 0 mean.
        (np.ndarray): Covariance matrix of the predicted state.

    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):
    """
    Project state distribution to measurement space.

    Args:
        mean (np.ndarray): The state's mean vector (8 dimensional array).
        covariance (np.ndarray): The state's covariance matrix (8x8 dimensional).

    Returns:
        (np.ndarray): Projected mean of the given state estimate.
        (np.ndarray): Projected 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):
    """
    Run Kalman filter correction step.

    Args:
        mean (np.ndarray): The predicted state's mean vector (8 dimensional).
        covariance (np.ndarray): The state's covariance matrix (8x8 dimensional).
        measurement (np.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:
        (np.ndarray): Measurement-corrected state mean.
        (np.ndarray): Measurement-corrected state covariance.

    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)