501 lines
15 KiB
Python
501 lines
15 KiB
Python
from dataclasses import dataclass
|
|
from typing import Generator, Optional, Protocol, Tuple, TypeAlias, TypedDict, Union
|
|
|
|
import numpy as np
|
|
from jaxtyping import Int, Float, Num, jaxtyped
|
|
from scipy.optimize import linear_sum_assignment
|
|
|
|
NDArray = np.ndarray
|
|
|
|
|
|
@dataclass
|
|
class LinearMotionNoInputModel:
|
|
F: Num[NDArray, "n n"]
|
|
Q: Num[NDArray, "n n"]
|
|
|
|
|
|
@dataclass
|
|
class LinearMeasurementModel:
|
|
H: Num[NDArray, "m n"]
|
|
R: Num[NDArray, "m m"]
|
|
|
|
|
|
Measurement = Num[NDArray, "m"]
|
|
|
|
|
|
@dataclass
|
|
class GaussianState:
|
|
x: Num[NDArray, "n"]
|
|
P: Num[NDArray, "n n"]
|
|
|
|
|
|
@dataclass(kw_only=True)
|
|
class CvModelGaussianState(GaussianState):
|
|
"""
|
|
constant velocity model with no input.
|
|
|
|
Note that the state vector is `[x, y, v_x, v_y]`.
|
|
|
|
This class can only verify the shape of the state vector and covariance matrix,
|
|
but not the content/order of the state vector.
|
|
"""
|
|
|
|
x: Num[NDArray, "4"]
|
|
P: Num[NDArray, "4 4"]
|
|
|
|
@staticmethod
|
|
def from_gaussian(state: GaussianState) -> "CvModelGaussianState":
|
|
assert state.x.shape == (4,), "state must have 4 elements"
|
|
assert state.P.shape == (4, 4), "covariance must be 4x4"
|
|
return CvModelGaussianState(x=state.x, P=state.P)
|
|
|
|
@property
|
|
def position(self) -> Num[NDArray, "2"]:
|
|
return self.x[:2]
|
|
|
|
@property
|
|
def velocity(self) -> Num[NDArray, "2"]:
|
|
return self.x[2:]
|
|
|
|
|
|
def predict(
|
|
state: GaussianState,
|
|
motion_model: LinearMotionNoInputModel,
|
|
) -> GaussianState:
|
|
x = state.x
|
|
P = state.P
|
|
F = motion_model.F
|
|
Q = motion_model.Q
|
|
assert x.shape[0] == F.shape[0], "state and transition model are not compatible"
|
|
assert F.shape[0] == F.shape[1], "transition model is not square"
|
|
assert (
|
|
F.shape[0] == Q.shape[0]
|
|
), "transition model and noise model are not compatible"
|
|
x_priori = F @ x
|
|
P_priori = F @ P @ F.T + Q
|
|
return GaussianState(x=x_priori, P=P_priori)
|
|
|
|
|
|
@dataclass
|
|
class PosterioriResult:
|
|
# updated state
|
|
state: GaussianState
|
|
innovation: NDArray
|
|
r"""
|
|
y. Innovation refers to the difference between the observed measurement and the predicted measurement. Also known as the residual.
|
|
|
|
.. math::
|
|
|
|
y = z - H x_{\text{priori}}
|
|
"""
|
|
innovation_covariance: NDArray
|
|
r"""
|
|
S. Innovation covariance refers to the covariance of the innovation (or residual) vector.
|
|
|
|
.. math::
|
|
|
|
S = H P H^T + R
|
|
"""
|
|
posteriori_measurement: NDArray
|
|
r"""
|
|
z_posteriori. The updated measurement prediction.
|
|
|
|
.. math::
|
|
|
|
z_{\text{posteriori}} = H x_{\text{posteriori}}
|
|
"""
|
|
mahalanobis_distance: NDArray
|
|
r"""
|
|
The Mahalanobis distance is a measure of the distance between a point P and a distribution D, introduced by P. Mahalanobis in 1936.
|
|
|
|
.. math::
|
|
|
|
\sqrt{y^T S^{-1} y}
|
|
"""
|
|
squared_mahalanobis_distance: NDArray
|
|
"""
|
|
If you are using the distance for statistical tests, such as identifying
|
|
outliers, the squared Mahalanobis distance is often used because it corresponds
|
|
to the chi-squared distribution when the underlying distribution is multivariate
|
|
normal.
|
|
"""
|
|
|
|
|
|
def predict_measurement(
|
|
state: GaussianState,
|
|
measure_model: LinearMeasurementModel,
|
|
) -> Measurement:
|
|
x = state.x
|
|
H = measure_model.H
|
|
return H @ x # type: ignore
|
|
|
|
|
|
def update(
|
|
measurement: Measurement,
|
|
state: GaussianState,
|
|
measure_model: LinearMeasurementModel,
|
|
) -> PosterioriResult:
|
|
x = state.x
|
|
P = state.P
|
|
H = measure_model.H
|
|
R = measure_model.R
|
|
assert x.shape[0] == H.shape[1], "state and measurement model are not compatible"
|
|
assert H.shape[0] == R.shape[0], "measurement model is not square"
|
|
assert H.shape[0] == R.shape[1], "measurement model is not square"
|
|
z = measurement
|
|
inv = np.linalg.inv
|
|
# innovation
|
|
# the priori measurement residual
|
|
y = z - H @ x
|
|
# innovation covariance
|
|
S = H @ P @ H.T + R
|
|
# Kalman gain
|
|
K = P @ H.T @ inv(S)
|
|
# posteriori state
|
|
x_posteriori = x + K @ y
|
|
# dummy identity matrix
|
|
I = np.eye(P.shape[0])
|
|
# posteriori covariance
|
|
I_KH = I - K @ H
|
|
P_posteriori = I_KH @ P @ I_KH.T + K @ R @ K.T
|
|
posteriori_state = GaussianState(x=x_posteriori, P=P_posteriori)
|
|
posteriori_measurement = H @ x_posteriori
|
|
s_m = y.T @ inv(S) @ y
|
|
return PosterioriResult(
|
|
state=posteriori_state,
|
|
innovation=y,
|
|
innovation_covariance=S,
|
|
posteriori_measurement=posteriori_measurement,
|
|
mahalanobis_distance=np.sqrt(s_m),
|
|
squared_mahalanobis_distance=s_m,
|
|
)
|
|
|
|
|
|
def cv_model(
|
|
v_x: float,
|
|
v_y: float,
|
|
dt: float,
|
|
q: float,
|
|
r: float,
|
|
) -> Tuple[
|
|
LinearMotionNoInputModel,
|
|
LinearMeasurementModel,
|
|
GaussianState,
|
|
]:
|
|
"""
|
|
Create a constant velocity model with no input
|
|
|
|
Args:
|
|
v_x: initial velocity in x direction
|
|
v_y: initial velocity in y direction
|
|
dt: time interval
|
|
q: process noise
|
|
r: measurement noise
|
|
|
|
Returns:
|
|
motion_model: motion model
|
|
measure_model: measurement model
|
|
state: initial state
|
|
"""
|
|
# yapf: disable
|
|
F = np.array([[1, 0, dt, 0],
|
|
[0, 1, 0, dt],
|
|
[0, 0, 1, 0],
|
|
[0, 0, 0, 1]])
|
|
H = np.array([[1, 0, 0, 0],
|
|
[0, 1, 0, 0]])
|
|
# yapf: enable
|
|
Q = q * np.eye(4)
|
|
R = r * np.eye(2)
|
|
P = np.eye(4)
|
|
motion_model = LinearMotionNoInputModel(F=F, Q=Q)
|
|
measure_model = LinearMeasurementModel(H=H, R=R)
|
|
state = GaussianState(x=np.array([0, 0, v_x, v_y]), P=P)
|
|
return motion_model, measure_model, state
|
|
|
|
|
|
def outer_distance(x: NDArray, y: NDArray) -> NDArray:
|
|
"""
|
|
Here's equivalent python code:
|
|
|
|
```python
|
|
res = jnp.empty((x.shape[0], y.shape[0]))
|
|
for i in range(x.shape[0]):
|
|
for j in range(y.shape[0]):
|
|
# res[i, j] = jnp.linalg.norm(x[i] - y[j])
|
|
res = res.at[i, j].set(jnp.linalg.norm(x[i] - y[j]))
|
|
return res
|
|
```
|
|
|
|
See Also
|
|
--------
|
|
`outer product <https://en.wikipedia.org/wiki/Outer_product>`_
|
|
"""
|
|
|
|
x_expanded = x[:, None, :]
|
|
y_expanded = y[None, :, :]
|
|
diff = y_expanded - x_expanded
|
|
return np.linalg.norm(diff, axis=-1)
|
|
|
|
|
|
@dataclass
|
|
class Tracking:
|
|
id: int
|
|
state: GaussianState
|
|
survived_time_steps: int
|
|
missed_time_steps: int
|
|
|
|
|
|
@dataclass
|
|
class TrackerParams:
|
|
dt: float = 1.0
|
|
cov_threshold: float = 4.0
|
|
tentative_mahalanobis_threshold: float = 10.0
|
|
confirm_mahalanobis_threshold: float = 10.0
|
|
forming_tracks_euclidean_threshold: float = 25.0
|
|
survival_steps_threshold: int = 3
|
|
|
|
|
|
class Tracker:
|
|
"""
|
|
A simple GNN tracker
|
|
"""
|
|
|
|
_last_measurements: NDArray = np.empty((0, 2), dtype=np.float32)
|
|
_tentative_tracks: list[Tracking] = []
|
|
_confirmed_tracks: list[Tracking] = []
|
|
_last_id: int = 0
|
|
|
|
def __init__(self):
|
|
self._last_measurements = np.empty((0, 2), dtype=np.float32)
|
|
self._tentative_tracks = []
|
|
self._confirmed_tracks = []
|
|
|
|
@staticmethod
|
|
def _predict(tracks: list[Tracking], dt: float = 1.0):
|
|
return [
|
|
Tracking(
|
|
id=track.id,
|
|
state=predict(track.state, Tracker.motion_model(dt=dt)),
|
|
survived_time_steps=track.survived_time_steps,
|
|
missed_time_steps=track.missed_time_steps,
|
|
)
|
|
for track in tracks
|
|
]
|
|
|
|
@staticmethod
|
|
def _data_associate_and_update(
|
|
measurements: NDArray, tracks: list[Tracking], distance_threshold: float = 3
|
|
) -> NDArray:
|
|
"""
|
|
Match tracks with measurements and update the tracks
|
|
|
|
Parameters
|
|
----------
|
|
[in] measurements: Float["a 2"]
|
|
[in,out] tracks: Tracking["b"]
|
|
|
|
Returns
|
|
----------
|
|
return
|
|
Float["... 2"] the unmatched measurements
|
|
|
|
Effect
|
|
----------
|
|
find the best match by minimum Mahalanobis distance, please note that I assume the state has been predicted
|
|
"""
|
|
if len(tracks) == 0:
|
|
return measurements
|
|
|
|
def _update(measurement: NDArray, tracking: Tracking):
|
|
return update(measurement, tracking.state, Tracker.measurement_model())
|
|
|
|
def outer_posteriori(
|
|
measurements: NDArray, tracks: list[Tracking]
|
|
) -> list[list[PosterioriResult]]:
|
|
"""
|
|
calculate the outer posteriori for each measurement and track
|
|
|
|
Parameters
|
|
----------
|
|
[in] measurements: Float["a 2"]
|
|
[in] tracks: Tracking["b"]
|
|
|
|
Returns
|
|
----------
|
|
PosterioriResult["a b"]
|
|
"""
|
|
return [
|
|
[_update(measurement, tracking) for measurement in measurements]
|
|
for tracking in tracks
|
|
]
|
|
|
|
def posteriori_to_mahalanobis(
|
|
posteriori: list[list[PosterioriResult]],
|
|
) -> NDArray:
|
|
"""
|
|
Parameters
|
|
----------
|
|
[in] posteriori: PosterioriResult["a b"]
|
|
|
|
Returns
|
|
----------
|
|
Float["a b"]
|
|
"""
|
|
return np.array(
|
|
[[r_m.mahalanobis_distance for r_m in p_t] for p_t in posteriori],
|
|
dtype=np.float32,
|
|
)
|
|
|
|
posteriors = outer_posteriori(measurements, tracks)
|
|
distances = posteriori_to_mahalanobis(posteriors)
|
|
row, col = linear_sum_assignment(np.array(distances))
|
|
row = np.array(row)
|
|
col = np.array(col)
|
|
|
|
def to_be_deleted() -> Generator[Tuple[int, int], None, None]:
|
|
for i, j in zip(row, col):
|
|
post: PosterioriResult = posteriors[i][j]
|
|
if post.mahalanobis_distance > distance_threshold:
|
|
yield i, j
|
|
|
|
for i, j in to_be_deleted():
|
|
row = row[row != i]
|
|
col = col[col != j]
|
|
|
|
for i, j in zip(row, col):
|
|
track: Tracking = tracks[i]
|
|
post: PosterioriResult = posteriors[i][j]
|
|
track.state = post.state
|
|
track.survived_time_steps += 1
|
|
tracks[i] = track
|
|
|
|
for i, track in enumerate(tracks):
|
|
if i not in row:
|
|
# reset the survived time steps once missed
|
|
track.missed_time_steps += 1
|
|
tracks[i] = track
|
|
# remove measurements that have been matched
|
|
left_measurements = np.delete(measurements, col, axis=0)
|
|
return left_measurements
|
|
|
|
def _tracks_from_past_measurements(
|
|
self, measurements: NDArray, dt: float = 1.0, distance_threshold: float = 3.0
|
|
):
|
|
"""
|
|
consume the last measurements and create tentative tracks from them
|
|
|
|
Note
|
|
----
|
|
mutate self._tentative_tracks and self._last_measurements
|
|
"""
|
|
if self._last_measurements.shape[0] == 0:
|
|
self._last_measurements = measurements
|
|
return
|
|
distances = outer_distance(self._last_measurements, measurements)
|
|
row, col = linear_sum_assignment(distances)
|
|
row = np.array(row)
|
|
col = np.array(col)
|
|
|
|
def to_be_deleted() -> Generator[Tuple[int, int], None, None]:
|
|
for i, j in zip(row, col):
|
|
euclidean_distance = distances[i, j]
|
|
if euclidean_distance > distance_threshold:
|
|
yield i, j
|
|
|
|
for i, j in to_be_deleted():
|
|
row = row[row != i]
|
|
col = col[col != j]
|
|
|
|
for i, j in zip(row, col):
|
|
coord = measurements[j]
|
|
vel = (coord - self._last_measurements[i]) / dt
|
|
s = np.concatenate([coord, vel])
|
|
state = GaussianState(x=s, P=np.eye(4))
|
|
track = Tracking(
|
|
id=self._last_id,
|
|
state=state,
|
|
survived_time_steps=0,
|
|
missed_time_steps=0,
|
|
)
|
|
self._last_id += 1
|
|
self._tentative_tracks.append(track)
|
|
# update the last measurements with the unmatched measurements
|
|
self._last_measurements = np.delete(measurements, col, axis=0)
|
|
|
|
def _transfer_tentative_to_confirmed(self, survival_steps_threshold: int = 3):
|
|
"""
|
|
transfer tentative tracks to confirmed tracks
|
|
|
|
Note
|
|
----
|
|
mutate self._tentative_tracks and self._confirmed_tracks in place
|
|
"""
|
|
for i, track in enumerate(self._tentative_tracks):
|
|
if track.survived_time_steps > survival_steps_threshold:
|
|
self._confirmed_tracks.append(track)
|
|
self._tentative_tracks.pop(i)
|
|
|
|
@staticmethod
|
|
def _track_cov_deleter(tracks: list[Tracking], cov_threshold: float = 4.0):
|
|
"""
|
|
delete tracks with covariance trace greater than threshold
|
|
|
|
Parameters
|
|
----------
|
|
[in,out] tracks: list[Tracking]
|
|
cov_threshold: float
|
|
the threshold of the covariance trace
|
|
|
|
Note
|
|
----
|
|
mutate tracks in place
|
|
"""
|
|
for i, track in enumerate(tracks):
|
|
# https://numpy.org/doc/stable/reference/generated/numpy.trace.html
|
|
if np.trace(track.state.P) > cov_threshold:
|
|
tracks.pop(i)
|
|
|
|
def next_measurements(self, measurements: NDArray, params: TrackerParams):
|
|
self._confirmed_tracks = self._predict(self._confirmed_tracks, params.dt)
|
|
self._tentative_tracks = self._predict(self._tentative_tracks, params.dt)
|
|
left_ = self._data_associate_and_update(
|
|
measurements, self._confirmed_tracks, params.confirm_mahalanobis_threshold
|
|
)
|
|
left = self._data_associate_and_update(
|
|
left_, self._tentative_tracks, params.tentative_mahalanobis_threshold
|
|
)
|
|
self._transfer_tentative_to_confirmed(params.survival_steps_threshold)
|
|
self._tracks_from_past_measurements(
|
|
left, params.dt, params.forming_tracks_euclidean_threshold
|
|
)
|
|
self._track_cov_deleter(self._tentative_tracks, params.cov_threshold)
|
|
self._track_cov_deleter(self._confirmed_tracks, params.cov_threshold)
|
|
|
|
@property
|
|
def confirmed_tracks(self):
|
|
return self._confirmed_tracks
|
|
|
|
@staticmethod
|
|
def motion_model(dt: float = 1, q: float = 0.05) -> LinearMotionNoInputModel:
|
|
"""
|
|
a constant velocity motion model
|
|
"""
|
|
# yapf: disable
|
|
F = np.array([[1, 0, dt, 0],
|
|
[0, 1, 0, dt],
|
|
[0, 0, 1, 0],
|
|
[0, 0, 0, 1]])
|
|
# yapf: enable
|
|
Q = q * np.eye(4)
|
|
return LinearMotionNoInputModel(F=F, Q=Q)
|
|
|
|
@staticmethod
|
|
def measurement_model(r: float = 0.75) -> LinearMeasurementModel:
|
|
# yapf: disable
|
|
H = np.array([[1, 0, 0, 0],
|
|
[0, 1, 0, 0]])
|
|
# yapf: enable
|
|
R = r * np.eye(2)
|
|
return LinearMeasurementModel(H=H, R=R)
|