133 lines
4.8 KiB
Python
133 lines
4.8 KiB
Python
"""Real-time gaze filters: the existing fixed median+EMA baseline, and two
|
|
ML-driven variants (hard state-switch, continuous confidence-weighted) built
|
|
on the classifier trained by train_filter_model.py.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import statistics
|
|
from collections import deque
|
|
from pathlib import Path
|
|
|
|
import joblib
|
|
|
|
from gaze_features import FEATURE_COLUMNS, OnlineFeatureTracker
|
|
|
|
FIXATION_ALPHA = 0.15
|
|
PURSUIT_ALPHA = 0.35
|
|
SACCADE_ALPHA = 0.9
|
|
SOFT_MIN_ALPHA = 0.15
|
|
SOFT_MAX_ALPHA = 0.9
|
|
MEDIAN_WINDOW = 5
|
|
|
|
|
|
class BaselineFilter:
|
|
"""Fixed median + EMA + dead-zone filter, matching visualize_surface_gaze.py."""
|
|
|
|
def __init__(self, smoothing: float = 0.18, median_window: int = 5, dead_zone: float = 3.0) -> None:
|
|
self.smoothing = smoothing
|
|
self.dead_zone = dead_zone
|
|
self.x_history: deque[float] = deque(maxlen=median_window)
|
|
self.y_history: deque[float] = deque(maxlen=median_window)
|
|
self.filtered: tuple[float, float] | None = None
|
|
|
|
def reset(self) -> None:
|
|
self.x_history.clear()
|
|
self.y_history.clear()
|
|
self.filtered = None
|
|
|
|
def update(self, timestamp: float, x: float, y: float, confidence: float) -> tuple[float, float]:
|
|
self.x_history.append(x)
|
|
self.y_history.append(y)
|
|
median_x = statistics.median(self.x_history)
|
|
median_y = statistics.median(self.y_history)
|
|
if self.filtered is None:
|
|
self.filtered = (median_x, median_y)
|
|
else:
|
|
prev_x, prev_y = self.filtered
|
|
fx = prev_x + self.smoothing * (median_x - prev_x)
|
|
fy = prev_y + self.smoothing * (median_y - prev_y)
|
|
if abs(fx - prev_x) < self.dead_zone:
|
|
fx = prev_x
|
|
if abs(fy - prev_y) < self.dead_zone:
|
|
fy = prev_y
|
|
self.filtered = (fx, fy)
|
|
return self.filtered
|
|
|
|
|
|
class MLAdaptiveFilter:
|
|
"""Classifier-driven filter. mode="hard" switches behavior by predicted class;
|
|
mode="soft" blends smoothing strength continuously with P(saccade)."""
|
|
|
|
def __init__(
|
|
self,
|
|
model_path: str | Path,
|
|
cm_per_px: float,
|
|
viewing_distance_cm: float,
|
|
mode: str = "soft",
|
|
) -> None:
|
|
if mode not in ("hard", "soft"):
|
|
raise ValueError("mode must be 'hard' or 'soft'")
|
|
self.mode = mode
|
|
self.model = joblib.load(model_path)
|
|
meta_path = Path(model_path).with_name(Path(model_path).stem + "_meta.json")
|
|
self.meta = json.loads(meta_path.read_text(encoding="utf-8")) if meta_path.exists() else {}
|
|
self.saccade_index = list(self.model.classes_).index("saccade")
|
|
|
|
self.tracker = OnlineFeatureTracker(cm_per_px, viewing_distance_cm)
|
|
self.x_history: deque[float] = deque(maxlen=MEDIAN_WINDOW)
|
|
self.y_history: deque[float] = deque(maxlen=MEDIAN_WINDOW)
|
|
self.filtered: tuple[float, float] | None = None
|
|
self.last_label: str | None = None
|
|
self.last_prob_saccade: float | None = None
|
|
|
|
def reset(self) -> None:
|
|
self.tracker.reset()
|
|
self.x_history.clear()
|
|
self.y_history.clear()
|
|
self.filtered = None
|
|
self.last_label = None
|
|
self.last_prob_saccade = None
|
|
|
|
def update(self, timestamp: float, x: float, y: float, confidence: float) -> tuple[float, float]:
|
|
features = self.tracker.update(timestamp, x, y, confidence)
|
|
if features is None:
|
|
# Not enough history yet for a prediction; pass raw through as median-of-one.
|
|
self.x_history.append(x)
|
|
self.y_history.append(y)
|
|
self.filtered = (x, y)
|
|
return self.filtered
|
|
|
|
row = [[features[c] for c in FEATURE_COLUMNS]]
|
|
proba = self.model.predict_proba(row)[0]
|
|
prob_saccade = float(proba[self.saccade_index])
|
|
label = self.model.classes_[proba.argmax()]
|
|
self.last_label = label
|
|
self.last_prob_saccade = prob_saccade
|
|
|
|
if self.mode == "hard":
|
|
if label == "saccade":
|
|
self.x_history.clear()
|
|
self.y_history.clear()
|
|
self.x_history.append(x)
|
|
self.y_history.append(y)
|
|
self.filtered = (x, y)
|
|
return self.filtered
|
|
alpha = FIXATION_ALPHA if label == "fixation" else PURSUIT_ALPHA
|
|
else:
|
|
alpha = SOFT_MIN_ALPHA + prob_saccade * (SOFT_MAX_ALPHA - SOFT_MIN_ALPHA)
|
|
|
|
self.x_history.append(x)
|
|
self.y_history.append(y)
|
|
median_x = statistics.median(self.x_history)
|
|
median_y = statistics.median(self.y_history)
|
|
if self.filtered is None:
|
|
self.filtered = (median_x, median_y)
|
|
else:
|
|
prev_x, prev_y = self.filtered
|
|
fx = prev_x + alpha * (median_x - prev_x)
|
|
fy = prev_y + alpha * (median_y - prev_y)
|
|
self.filtered = (fx, fy)
|
|
return self.filtered
|