EyeTrackVR/EyeTrackApp/pye3dcustom/kalman.py
2022-07-12 12:30:31 -07:00

59 lines
2.1 KiB
Python

"""
(*)~---------------------------------------------------------------------------
Pupil - eye tracking platform
Copyright (C) 2012-2019 Pupil Labs
Distributed under the terms of the GNU
Lesser General Public License (LGPL v3.0).
See COPYING and COPYING.LESSER for license details.
---------------------------------------------------------------------------~(*)
"""
import cv2
import numpy as np
class KalmanFilter(object):
def __init__(self):
self.filter = cv2.KalmanFilter(7, 3, 0, cv2.CV_32F)
self.filter.measurementMatrix = np.asarray(
[[1, 0, 0, 0, 0, 0, 0], [0, 1, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 1]],
dtype=np.float32,
)
self.filter.processNoiseCov = 1e-4 * np.eye(7, dtype=np.float32)
self.filter.measurementNoiseCov = 1e-5 * np.eye(3, dtype=np.float32)
self.filter.measurementNoiseCov[2][2] = 0.1
self.filter.statePost = np.asarray([0, 0, 0, 0, 0, 0, 2.0], dtype=np.float32)
self.filter.errorCovPost = np.eye(7, dtype=np.float32)
self.last_call = -1
def predict(self, t):
if self.last_call != -1 and t > self.last_call:
dt = t - self.last_call
self.filter.transitionMatrix = np.asarray(
[
[1, 0, dt, 0, 0.5 * dt * dt, 0, 0],
[0, 1, 0, dt, 0, 0.5 * dt * dt, 0],
[0, 0, 1, 0, dt, 0, 0],
[0, 0, 0, 1, 0, dt, 0],
[0, 0, 0, 0, 1, 0, 0],
[0, 0, 0, 0, 0, 1, 0],
[0, 0, 0, 0, 0, 0, 1],
],
dtype=np.float32,
)
prediction = self.filter.predict()
phi, theta, pupil_radius = (
prediction[0][0],
prediction[1][0],
prediction[6][0],
)
else:
phi, theta, pupil_radius = -np.pi / 2, np.pi / 2, 0
self.last_call = t
return phi, theta, pupil_radius
def correct(self, phi, theta, radius):
self.filter.correct(np.asarray([phi, theta, radius], dtype=np.float32))