EyeTrackVR/EyeTrackApp-SingleCam/pye3dcustom/refraction.py
2022-07-12 12:12:38 -07:00

142 lines
5.1 KiB
Python

import itertools
from pathlib import Path
from .cpp.refraction_correction import apply_correction_pipeline
import numpy as np
import msgpack
LOAD_DIR = Path(__file__).parent / "refraction_models"
LOAD_VERSION = 1
class ModelDeserializationError(Exception):
pass
class Refractionizer:
def __init__(self, degree=3, type_="default", custom_load_dir=None):
self.pipeline_radius_as_list = self.load_config_from_msgpack(
"radius", type_, degree, custom_load_dir
)
self.pipeline_gaze_vector_as_list = self.load_config_from_msgpack(
"gaze_vector", type_, degree, custom_load_dir
)
self.pipeline_sphere_center_as_list = self.load_config_from_msgpack(
"sphere_center", type_, degree, custom_load_dir
)
self.pipeline_pupil_circle_as_list = self.load_config_from_msgpack(
"pupil_circle", type_, degree, custom_load_dir
)
@staticmethod
def load_config_from_msgpack(feature, type_, degree, custom_load_dir=None):
load_dir = Path(custom_load_dir or LOAD_DIR).resolve()
name = f"{type_}_refraction_model_{feature}_degree_{degree}.msgpack"
path = load_dir / name
with path.open("rb") as file:
config_model = msgpack.unpack(file)
Refractionizer._validate_loaded_model_config(config_model)
try:
return list(
itertools.chain(
Refractionizer._polynomial_features_from_config(config_model),
Refractionizer._standard_scaler_from_config(config_model),
Refractionizer._linear_regression_from_config(config_model),
)
)
except KeyError as err:
raise ModelDeserializationError from err
@staticmethod
def _validate_loaded_model_config(config_model):
if not isinstance(config_model, dict) or "version" not in config_model:
raise ModelDeserializationError("Unrecognized format")
if config_model["version"] != LOAD_VERSION:
raise ModelDeserializationError(
f"Unexpected version `{config_model['version']}` "
f"(expected `{LOAD_VERSION}``)"
)
@staticmethod
def _polynomial_features_from_config(config_model):
yield np.array(config_model["steps"]["PolynomialFeatures"]["powers"])
@staticmethod
def _standard_scaler_from_config(config_model):
config_scaler = config_model["steps"]["StandardScaler"]
yield np.array(config_scaler["mean"])
yield np.array(config_scaler["var"])
@staticmethod
def _linear_regression_from_config(config_model):
config_lin_reg = config_model["steps"]["LinearRegression"]
yield np.array(config_lin_reg["coef"])
yield np.array(config_lin_reg["intercept"])
@staticmethod
def _apply_correction_pipeline(X, pipeline_arrays):
return apply_correction_pipeline(np.asarray(X).T, *pipeline_arrays)
def correct_radius(self, X):
return self._apply_correction_pipeline(X, self.pipeline_radius_as_list)
def correct_gaze_vector(self, X):
return self._apply_correction_pipeline(X, self.pipeline_gaze_vector_as_list)
def correct_sphere_center(self, X):
return self._apply_correction_pipeline(X, self.pipeline_sphere_center_as_list)
def correct_pupil_circle(self, X):
return self._apply_correction_pipeline(X, self.pipeline_pupil_circle_as_list)
class SklearnRefractionizer(Refractionizer):
def __init__(self, degree=3, type_="default", custom_load_dir=None):
self.correct_radius = self.load_predict_fn_from_joblib_pickle(
"radius", type_, degree, custom_load_dir
)
self.correct_gaze_vector = self.load_predict_fn_from_joblib_pickle(
"gaze_vector", type_, degree, custom_load_dir
)
self.correct_sphere_center = self.load_predict_fn_from_joblib_pickle(
"sphere_center", type_, degree, custom_load_dir
)
self.correct_pupil_circle = self.load_predict_fn_from_joblib_pickle(
"pupil_circle", type_, degree, custom_load_dir
)
@staticmethod
def load_predict_fn_from_joblib_pickle(
feature, type_, degree, custom_load_dir=None
):
import joblib
load_dir = Path(custom_load_dir or LOAD_DIR).resolve()
name = f"{type_}_refraction_model_{feature}_degree_{degree}.save"
path = load_dir / name
try:
pipeline = joblib.load(path)
except FileNotFoundError as err:
raise
except Exception as exc:
raise ModelDeserializationError(
f"Failed to load pickled model from {path}"
) from exc
return pipeline.predict
if __name__ == "__main__":
refractionizer = Refractionizer()
print(refractionizer.correct_sphere_center([[0.0, 0.0, 35.0]]))
print(refractionizer.correct_radius([[0.0, 0.0, 35.0, 0.0, 0.0, -1.0, 2.0]]))
print(refractionizer.correct_gaze_vector([[0.0, 0.0, 35.0, 0.0, 0.0, -1.0, 2.0]]))
print(refractionizer.correct_pupil_circle([[0.0, 0.0, 35.0, 0.0, 0.0, -1.0, 2.0]]))