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

298 lines
11 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 logging
import typing as T
import numpy as np
from .abstract import TwoSphereModelAbstract, SphereCenterEstimates
from ..camera import CameraModel
from ..constants import _EYE_RADIUS_DEFAULT, DEFAULT_SPHERE_CENTER
from ..geometry.intersections import nearest_point_on_sphere_to_line
from ..geometry.primitives import Circle, Line
from ..geometry.projections import (
project_line_into_image_plane,
project_point_into_image_plane,
unproject_ellipse,
)
from ..geometry.utilities import normalize
from ..observation import BasicStorage, Observation, ObservationStorage
from ..refraction import Refractionizer
logger = logging.getLogger(__name__)
class TwoSphereModel(TwoSphereModelAbstract):
def __init__(
self,
camera: CameraModel,
storage_cls: T.Type[ObservationStorage] = None,
storage_kwargs: T.Dict = None,
):
if storage_cls:
kwargs = storage_kwargs if storage_kwargs is not None else {}
self.storage = storage_cls(**kwargs)
else:
self.storage = BasicStorage()
self.camera = camera
self.refractionizer = Refractionizer()
self._set_default_model_params()
@property
def sphere_center(self) -> np.ndarray:
return self._sphere_center
@sphere_center.setter
def sphere_center(self, coordinates: np.ndarray):
self._sphere_center = coordinates
@property
def corrected_sphere_center(self) -> np.ndarray:
return self._corrected_sphere_center
@corrected_sphere_center.setter
def corrected_sphere_center(self, coordinates: np.ndarray):
self._corrected_sphere_center = coordinates
@property
def projected_sphere_center(self) -> np.ndarray:
return self._projected_sphere_center
@projected_sphere_center.setter
def projected_sphere_center(self, projected_sphere_center: np.ndarray):
self._projected_sphere_center = projected_sphere_center
def _set_default_model_params(self):
# Overwrite in subclasses that do not allow setting these attributes
self._sphere_center = np.asarray(DEFAULT_SPHERE_CENTER)
self._corrected_sphere_center = self.refractionizer.correct_sphere_center(
np.asarray([[*self.sphere_center]])
)[0]
self.rms_residual = np.nan
def add_observation(self, observation: Observation):
self.storage.add(observation)
@property
def n_observations(self) -> int:
return self.storage.count()
def set_sphere_center(self, new_sphere_center):
self.sphere_center = new_sphere_center
self.corrected_sphere_center = self.refractionizer.correct_sphere_center(
np.asarray([[*self.sphere_center]])
)[0]
def estimate_sphere_center(
self,
from_2d=None,
prior_3d=None,
prior_strength=0.0,
calculate_rms_residual=False,
):
self.projected_sphere_center = (
from_2d if from_2d is not None else self.estimate_sphere_center_2d()
)
sphere_center, rms_residual = self.estimate_sphere_center_3d(
self.projected_sphere_center,
prior_3d,
prior_strength,
calculate_rms_residual=calculate_rms_residual,
)
self.set_sphere_center(sphere_center)
self.rms_residual = rms_residual if rms_residual is not None else float("nan")
return SphereCenterEstimates(
self.projected_sphere_center, sphere_center, rms_residual
)
def estimate_sphere_center_2d(self):
observations = self.storage.observations
# slightly faster than np.array
aux_2d = np.concatenate([obs.aux_2d for obs in observations])
aux_2d.shape = -1, 2, 3
# Estimate projected sphere center by nearest intersection of 2d gaze lines
sum_aux_2d = aux_2d.sum(axis=0)
projected_sphere_center = np.linalg.pinv(sum_aux_2d[:2, :2]) @ sum_aux_2d[:2, 2]
return projected_sphere_center
def estimate_sphere_center_3d(
self,
sphere_center_2d,
prior_3d=None,
prior_strength=0.0,
calculate_rms_residual=False,
) -> T.Tuple[np.array, T.Optional[float]]:
observations, aux_3d, gaze_2d = self._prep_data()
sum_aux_3d, disamb_indices, aux_3d_disamb = self._disambiguate_dierkes_lines(
aux_3d, gaze_2d, sphere_center_2d
)
sphere_center = self._calc_sphere_center(sum_aux_3d, prior_3d, prior_strength)
rms_residual = (
self._calc_rms_residual(
observations, disamb_indices, sphere_center, aux_3d_disamb
)
if calculate_rms_residual
else None
)
return sphere_center, rms_residual
def _prep_data(self):
observations = self.storage.observations
aux_3d = np.concatenate([obs.aux_3d for obs in observations])
aux_3d.shape = -1, 2, 3, 4
gaze_2d = np.concatenate([obs.gaze_2d_line for obs in observations])
gaze_2d.shape = -1, 4
return observations, aux_3d, gaze_2d
def _disambiguate_dierkes_lines(self, aux_3d, gaze_2d, sphere_center_2d):
# Disambiguate Dierkes lines
# We want gaze_2d to points towards the sphere center. gaze_2d was collected
# from Dierkes[0]. If it points into the correct direction, we know that
# Dierkes[0] is the correct one to use, otherwise we need to use Dierkes[1]. We
# can check that with the sign of the dot product.
gaze_2d_origins = gaze_2d[:, :2]
gaze_2d_directions = gaze_2d[:, 2:]
gaze_2d_towards_center = gaze_2d_origins - sphere_center_2d
dot_products = np.sum(gaze_2d_towards_center * gaze_2d_directions, axis=1)
disambiguation_indices = np.where(dot_products < 0, 1, 0)
obs_idc = np.arange(disambiguation_indices.shape[0])
aux_3d_disambiguated = aux_3d[obs_idc, disambiguation_indices, :, :]
# Estimate sphere center by nearest intersection of Dierkes lines
sum_aux_3d = aux_3d_disambiguated.sum(axis=0)
return sum_aux_3d, disambiguation_indices, aux_3d_disambiguated
def _calc_sphere_center(self, sum_aux_3d, prior_3d=None, prior_strength=0.0):
matrix = sum_aux_3d[:3, :3]
try:
if prior_3d is None:
return np.linalg.pinv(matrix) @ sum_aux_3d[:3, 3]
else:
return np.linalg.pinv(matrix + prior_strength * np.eye(3)) @ (
sum_aux_3d[:3, 3] + prior_strength * prior_3d
)
except np.linalg.LinAlgError:
# happens if lines are parallel, very rare
return DEFAULT_SPHERE_CENTER
def _calc_rms_residual(
self, observations, disamb_indices, sphere_center, aux_3d_disamb
):
# Here we use eq. (10) in https://docplayer.net/21072949-Least-squares-intersection-of-lines.html.
origins_dierkes_lines = np.array(
[
obs.get_Dierkes_line(idx).origin
for obs, idx in zip(observations, disamb_indices)
]
)
origins_dierkes_lines.shape = -1, 3, 1
deltas = origins_dierkes_lines - sphere_center[:, np.newaxis]
tmp = np.einsum("ijk,ikl->ijl", aux_3d_disamb[:, :3, :3], deltas)
squared_residuals = np.einsum(
"ikj,ijk->i", np.transpose(deltas, (0, 2, 1)), tmp
)
rms_residual = np.clip(squared_residuals, 0.0, None)
rms_residual = np.mean(np.sqrt(rms_residual))
return rms_residual
# GAZE PREDICTION
def _extract_unproject_disambiguate(self, pupil_datum):
ellipse = self._extract_ellipse(pupil_datum)
circle_3d_pair = unproject_ellipse(ellipse, self.camera.focal_length)
if circle_3d_pair:
circle_3d = self._disambiguate_circle_3d_pair(circle_3d_pair)
else:
circle_3d = Circle([0.0, 0.0, 0.0], [0.0, 0.0, -1.0], 0.0)
return circle_3d
def _disambiguate_circle_3d_pair(self, circle_3d_pair):
circle_center_2d = project_point_into_image_plane(
circle_3d_pair[0].center, self.camera.focal_length
)
circle_normal_2d = normalize(
project_line_into_image_plane(
Line(circle_3d_pair[0].center, circle_3d_pair[0].normal),
self.camera.focal_length,
).direction
)
sphere_center_2d = project_point_into_image_plane(
self.sphere_center, self.camera.focal_length
)
if np.dot(circle_center_2d - sphere_center_2d, circle_normal_2d) >= 0:
return circle_3d_pair[0]
else:
return circle_3d_pair[1]
def predict_pupil_circle(
self, observation: Observation, use_unprojection: bool = False
) -> Circle:
if observation.invalid:
return Circle.null()
circle_3d = self._disambiguate_circle_3d_pair(observation.circle_3d_pair)
unprojection_depth = np.linalg.norm(circle_3d.center)
direction = circle_3d.center / unprojection_depth
nearest_point_on_sphere = nearest_point_on_sphere_to_line(
self.sphere_center, _EYE_RADIUS_DEFAULT, [0.0, 0.0, 0.0], direction
)
if use_unprojection:
gaze_vector = circle_3d.normal
else:
gaze_vector = normalize(nearest_point_on_sphere - self.sphere_center)
radius = np.linalg.norm(nearest_point_on_sphere) / unprojection_depth
pupil_circle = Circle(nearest_point_on_sphere, gaze_vector, radius)
return pupil_circle
def apply_refraction_correction(self, pupil_circle):
input_features = np.asarray(
[[*self.sphere_center, *pupil_circle.normal, pupil_circle.radius]]
)
refraction_corrected_params = self.refractionizer.correct_pupil_circle(
input_features
)[0]
refraction_corrected_gaze_vector = normalize(refraction_corrected_params[:3])
refraction_corrected_radius = refraction_corrected_params[-1]
refraction_corrected_pupil_center = (
self.corrected_sphere_center
+ _EYE_RADIUS_DEFAULT * refraction_corrected_gaze_vector
)
refraction_corrected_pupil_circle = Circle(
refraction_corrected_pupil_center,
refraction_corrected_gaze_vector,
refraction_corrected_radius,
)
return refraction_corrected_pupil_circle
def mean_observation_circularity(self):
observation_circularities = [
observation.ellipse.circularity()
for observation in self.storage.observations
]
return np.mean(observation_circularities)
def cleanup(self):
pass