Heat-Vision / src /tracking /kalman.py
TulkinRB's picture
Add stuff
0bdfe9d
import numpy as np
class VariableSpeedKalmanFilter:
H = np.eye(8)
def __init__(self, x_0, y_0, w_0, h_0, vx_0=0, vy_0=0, vh_0=0, vw_0=0, init_std=1, framerate=30):
self.x = np.array([[x_0, y_0, w_0, h_0, vx_0, vy_0, vw_0, vh_0]]).T
self.dt = 1 / framerate
self.P = np.eye(8) * init_std**2
self.F = np.eye(8) + (np.diag([1, 1, 1, 1], 4) * self.dt)
self.sigma_pos = 0.05 * self.dt * 30
self.sigma_vel = 0.00625 * self.dt * 30
self.sigma_measure = 0.05 * self.dt * 30
self.next_x = self.x.copy()
self.next_P = self.P.copy()
self.previous_box = np.array([x_0, y_0, w_0, h_0]).reshape((4, 1))
self.previous_var_box = np.array([
self.sigma_measure * self.next_x[2], self.sigma_measure * self.next_x[3],
self.sigma_measure * self.next_x[2], self.sigma_measure * self.next_x[3]
])**2
self.time_since_last_measurement = 0
def predict(self):
self.x = self.next_x.copy()
self.P = self.next_P.copy()
self.next_x = self.F @ self.x
Q_diag = np.array([
self.sigma_pos * self.x[2], self.sigma_pos * self.x[3], self.sigma_pos * self.x[2],
self.sigma_pos * self.x[3], self.sigma_vel * self.x[2], self.sigma_vel * self.x[3],
self.sigma_vel * self.x[2], self.sigma_vel * self.x[3]
])**2
Q = np.diag(Q_diag.flatten())
self.next_P = self.F @ self.P @ self.F.T + Q
self.time_since_last_measurement += self.dt
return self.next_x
def update(self, box):
box = box.reshape(4, 1)
velocities = (box - self.previous_box) / self.time_since_last_measurement
z = np.vstack((box, velocities))
y = z - (self.H @ self.next_x)
var_box = np.array([
self.sigma_measure * self.next_x[2], self.sigma_measure * self.next_x[3],
self.sigma_measure * self.next_x[2], self.sigma_measure * self.next_x[3]
])**2
var_vel = (var_box + self.previous_var_box) / (self.time_since_last_measurement ** 2)
cov_box_vel = var_box / self.time_since_last_measurement
R = np.block([
[np.diag(var_box.flatten()), np.zeros((4, 4))],
[np.zeros((4, 4)), np.diag(var_vel.flatten())],
])
S = self.H @ self.next_P @ self.H.T + R
K = self.next_P @ self.H.T @ np.linalg.inv(S)
self.x = self.next_x + K @ y
I = np.eye(8)
self.P = (I - K @ self.H) @ self.next_P
self.next_x = self.x.copy()
self.next_P = self.P.copy()
self.previous_box = box.copy()
self.time_since_last_measurement = 0
return self.x