mirror of
https://github.com/EyeTrackVR/EyeTrackVR.git
synced 2025-11-04 14:39:42 +08:00
59 lines
2.1 KiB
Python
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))
|