inference_pipeline.py

import numpy as np
import onnxruntime as ort
from scipy.interpolate import Rbf
from collections import deque


# =============================================================================
# GAZE ANGLE SMOOTHER  —  APPLY BEFORE calibration.predict(), NOT AFTER
#
# This is the single most important piece of the inference path, and its
# absence was the root cause of the "z-scoring made everything worse" episode.
#
# WHY IT MATTERS
# The RBF maps (pitch, yaw) -> screen. Because the calibration z-scores each
# axis by its own spread, and the vertical (pitch) axis has a very small
# spread, the RBF has a STEEP gradient along pitch: a tiny change in pitch
# swings the prediction a long way vertically. That steepness is necessary
# (otherwise the RBF cannot resolve vertical position at all) — but it means
# the RBF AMPLIFIES whatever noise is on its input.
#
# Feeding it a single raw frame therefore amplifies that frame's noise.
# Smoothing the OUTPUT afterwards (One Euro on screen coords) cannot undo
# this: by then the noise has already been multiplied through a nonlinear map.
# The noise must be removed BEFORE the query.
#
# Measured effect in simulation (z-scored RBF, realistic noise):
#     no input smoothing  -> err_x 113-348px, err_y ~330px   (unusable)
#     10-frame smoothing  ->  err_x  ~38px,   err_y ~179px   (usable)
#
# A rolling MEDIAN (not mean) is used so a single blink-tail or landmark
# glitch frame cannot drag the estimate.
#
# WINDOW is a latency/accuracy trade: at ~30fps, 10 frames is ~0.33s, which
# is roughly one fixation. Larger windows keep improving accuracy slightly but
# make the live dot feel laggy.
# =============================================================================

class GazeAngleSmoother:
    def __init__(self, window=10):
        self.window = window
        self._pitch = deque(maxlen=window)
        self._yaw   = deque(maxlen=window)
        self._ear   = deque(maxlen=window)

    def __call__(self, pitch, yaw, ear=0.0):
        self._pitch.append(float(pitch))
        self._yaw.append(float(yaw))
        self._ear.append(float(ear))
        return (float(np.median(self._pitch)),
                float(np.median(self._yaw)),
                float(np.median(self._ear)))

    def reset(self):
        self._pitch.clear()
        self._yaw.clear()
        self._ear.clear()


# =============================================================================
# ONNX MODEL (FIXED FOR BINOCULAR v4)
# =============================================================================

class GazeONNXModel:
    def __init__(self, onnx_path: str):
        self.session = ort.InferenceSession(
            onnx_path, providers=["CPUExecutionProvider"]
        )

        input_names = [inp.name for inp in self.session.get_inputs()]
        self.is_binocular = "eye_patch_binocular" in input_names

        print(f"ONNX loaded: {onnx_path}")
        print(f"Mode: {'binocular (v4)' if self.is_binocular else 'monocular'}")

        if self.is_binocular:
            dummy_patch = np.zeros((1, 2, 36, 60), dtype=np.float32)
            dummy_pose  = np.zeros((1, 3), dtype=np.float32)
            self.session.run(
                ["gaze"],
                {"eye_patch_binocular": dummy_patch, "head_pose": dummy_pose}
            )
        else:
            dummy_patch = np.zeros((1, 1, 36, 60), dtype=np.float32)
            dummy_pose  = np.zeros((1, 3), dtype=np.float32)
            self.session.run(
                ["gaze"],
                {"eye_patch": dummy_patch, "head_pose": dummy_pose}
            )

    def predict(self, left_patch, right_patch, head_pose):
        def norm(img):
            return (img.astype(np.float32) / 255.0 - 0.5) / 0.5

        pose = head_pose.astype(np.float32)[np.newaxis, :]

        if self.is_binocular:
            patch = np.stack([norm(left_patch), norm(right_patch)], axis=0)
            patch = patch[np.newaxis, :, :, :]

            out = self.session.run(
                ["gaze"],
                {"eye_patch_binocular": patch, "head_pose": pose}
            )[0][0]

        else:
            patch = norm(left_patch)[np.newaxis, np.newaxis, :, :]

            out = self.session.run(
                ["gaze"],
                {"eye_patch": patch, "head_pose": pose}
            )[0][0]

        return float(out[0]), float(out[1])  # pitch, yaw


# =============================================================================
# RBF CALIBRATION (with extrapolation clamp)
# =============================================================================

class RBFGazeCalibration:
    def __init__(self, screen_w=1493, screen_h=933,
                 function="multiquadric", smooth=0.1):
        self.screen_w = screen_w
        self.screen_h = screen_h
        self.function = function
        self.smooth   = smooth
        self._rbf_x   = None
        self._rbf_y   = None
        self._calibrated = False

        # Range of (pitch, yaw) actually seen during calibration. Multiquadric
        # RBF surfaces are only well-behaved INSIDE the convex hull of their
        # training points. Outside it, the surface is dominated by whichever
        # basis functions happen to be nearest, and the gradient direction is
        # not guaranteed to continue the trend. Clamping pitch/yaw to a small
        # margin beyond the calibrated range, before ever calling the RBF,
        # keeps every query inside (or just at the edge of) the well-behaved
        # interpolation region instead of letting it wander into extrapolation
        # territory.
        self._pitch_range = None
        self._yaw_range   = None

        # Output-side safety net, computed from where your calibration dots
        # actually were, not the full screen. Loosened from the original
        # 6%/50px to 10%/70px — the tighter clamp was creating a visible
        # dead zone near the bottom/right edges where predictions couldn't
        # reach even when the underlying gaze estimate was pointing there.
        # Combined with pushing the calibration grid itself to 0.02/0.98
        # (see calibrate.py), this should shrink that gap meaningfully
        # without reopening the reversal-error risk the clamp exists for.
        self._screen_x_range = None
        self._screen_y_range = None

        self._pitch_center = 0.0
        self._yaw_center    = 0.0
        self._pitch_scale   = 1.0
        self._yaw_scale     = 1.0
        self._vert_center   = 0.0
        self._vert_scale    = 1.0
        self._vert_feature  = "pitch"   # or "ear" — chosen by measurement


    def calibrate(self, pitch_yaw, screen_points, ear=None, screen_size=None):
        """
        pitch_yaw : (n, 2) array of [pitch, yaw] per calibration point
        screen_points : (n, 2) array of [sx, sy]
        ear : (n,) optional array of eye-aperture (EAR) per point.
        screen_size : (W, H) TRUE screen size in pixels. MUST be passed.

        WHY screen_size IS NOT OPTIONAL IN PRACTICE:
        This class used to fall back to hardcoded defaults of 1493x933 (a
        stale leftover from an old monitor), and InsightUXPipeline never
        passed anything, so EVERY prediction was silently hard-clamped to
        1493x933 in predict():
            sx = max(0, min(sx, self.screen_w))
        On a 1920x1200 screen that made the rightmost 427px (22% of width)
        and the bottom 267px (22% of height) literally unreachable — the
        gaze estimate was fine, the clamp threw it away. Symptom: the dot
        stops dead at an invisible line near the right and bottom edges.
        The true size is now recorded here and persisted in calibration.pkl.

        VERTICAL FEATURE SELECTION.

        Measurement showed the CNN's pitch output correlates with screen-Y at
        only r ~= 0.41 (it explains ~16% of vertical variance) while yaw
        correlates with screen-X at r ~= 0.99. The model is close to blind
        vertically, and no mapping can invent information its input lacks.

        Eye aperture is a physically independent vertical cue: looking down
        lowers the eyelid, so EAR shrinks. It is already computed for blink
        detection and costs nothing extra.

        Rather than assume EAR is better, this MEASURES both candidates
        against the true screen-Y and uses whichever actually correlates
        more strongly. If EAR turns out to be useless, nothing changes and
        we fall back to pitch — so this cannot make things worse.
        """
        if screen_size is not None:
            self.screen_w = int(screen_size[0])
            self.screen_h = int(screen_size[1])
        print(f"  screen size recorded: {self.screen_w} x {self.screen_h}")

        pitch = pitch_yaw[:, 0]
        yaw   = pitch_yaw[:, 1]
        sy    = screen_points[:, 1]

        def _r(a, b):
            a = np.asarray(a, float); b = np.asarray(b, float)
            if a.std() < 1e-12 or b.std() < 1e-12:
                return 0.0
            return float(np.corrcoef(a, b)[0, 1])

        r_pitch_y = _r(pitch, sy)
        r_ear_y   = _r(ear, sy) if ear is not None else 0.0

        # Require EAR to be meaningfully better, not just noise-better.
        use_ear = (ear is not None) and (abs(r_ear_y) > abs(r_pitch_y) + 0.10)

        if use_ear:
            vert = np.asarray(ear, dtype=float)
            self._vert_feature = "ear"
            print(f"  vertical feature: EYE APERTURE (r={r_ear_y:+.3f}) "
                  f"— beats CNN pitch (r={r_pitch_y:+.3f})")
        else:
            vert = pitch
            self._vert_feature = "pitch"
            if ear is not None:
                print(f"  vertical feature: CNN PITCH (r={r_pitch_y:+.3f}) "
                      f"— eye aperture was not better (r={r_ear_y:+.3f})")

        self._pitch_center = float(np.mean(pitch))
        self._yaw_center   = float(np.mean(yaw))
        self._pitch_scale  = float(np.std(pitch)) or 1e-3
        self._yaw_scale    = float(np.std(yaw)) or 1e-3
        self._vert_center  = float(np.mean(vert))
        self._vert_scale   = float(np.std(vert)) or 1e-3

        pitch_n = (pitch - self._pitch_center) / self._pitch_scale
        yaw_n   = (yaw   - self._yaw_center)   / self._yaw_scale
        vert_n  = (vert  - self._vert_center)  / self._vert_scale

        # Horizontal mapping is healthy — leave it on (pitch, yaw).
        self._rbf_x = Rbf(pitch_n, yaw_n, screen_points[:, 0],
                          function=self.function, smooth=self.smooth)
        # Vertical mapping uses whichever vertical cue actually tracks Y.
        self._rbf_y = Rbf(vert_n, yaw_n, screen_points[:, 1],
                          function=self.function, smooth=self.smooth)

        # NOTE — a leave-one-out "gain correction" (fitting a linear stretch
        # to counteract edge under-reach) used to live here. It was REMOVED
        # after simulation showed it roughly DOUBLES vertical error
        # (306px -> 519px). The reason is straightforward in hindsight: when
        # an axis is noise-dominated, the leave-one-out residuals it fits on
        # are mostly noise, and a linear stretch fitted to noise just
        # amplifies that noise. Do not reintroduce it without re-running that
        # experiment.

        margin_p = 0.10 * (pitch.max() - pitch.min() + 1e-6)
        margin_y = 0.10 * (yaw.max()   - yaw.min()   + 1e-6)
        self._pitch_range = (float(pitch.min() - margin_p), float(pitch.max() + margin_p))
        self._yaw_range   = (float(yaw.min()   - margin_y), float(yaw.max()   + margin_y))

        sx_vals = screen_points[:, 0]
        sy_vals = screen_points[:, 1]
        margin_sx = max(70.0, 0.10 * (sx_vals.max() - sx_vals.min() + 1e-6))
        margin_sy = max(70.0, 0.10 * (sy_vals.max() - sy_vals.min() + 1e-6))
        self._screen_x_range = (float(sx_vals.min() - margin_sx), float(sx_vals.max() + margin_sx))
        self._screen_y_range = (float(sy_vals.min() - margin_sy), float(sy_vals.max() + margin_sy))

        self._calibrated = True
        print(f"RBF calibrated on {len(pitch_yaw)} points")
        print(f"  pitch clamp range: {self._pitch_range[0]:.4f} to {self._pitch_range[1]:.4f}")
        print(f"  yaw clamp range:   {self._yaw_range[0]:.4f} to {self._yaw_range[1]:.4f}")
        print(f"  screen x clamp:    {self._screen_x_range[0]:.0f} to {self._screen_x_range[1]:.0f}")
        print(f"  screen y clamp:    {self._screen_y_range[0]:.0f} to {self._screen_y_range[1]:.0f}")

    def predict(self, pitch, yaw, ear=0.0):
        if not self._calibrated:
            raise RuntimeError("Calibration not done")

        if self._pitch_range is not None:
            pitch = float(np.clip(pitch, self._pitch_range[0], self._pitch_range[1]))
        if self._yaw_range is not None:
            yaw = float(np.clip(yaw, self._yaw_range[0], self._yaw_range[1]))

        pitch_n = (pitch - self._pitch_center) / self._pitch_scale
        yaw_n   = (yaw   - self._yaw_center)   / self._yaw_scale

        # Vertical uses whichever cue calibration measured as stronger.
        vert = float(ear) if self._vert_feature == "ear" else float(pitch)
        vert_n = (vert - self._vert_center) / self._vert_scale

        sx = float(self._rbf_x(pitch_n, yaw_n))
        sy = float(self._rbf_y(vert_n, yaw_n))

        if self._screen_x_range is not None:
            sx = float(np.clip(sx, self._screen_x_range[0], self._screen_x_range[1]))
        if self._screen_y_range is not None:
            sy = float(np.clip(sy, self._screen_y_range[0], self._screen_y_range[1]))

        sx = max(0, min(sx, self.screen_w))
        sy = max(0, min(sy, self.screen_h))

        return sx, sy

    def save(self, path):
        import pickle
        with open(path, "wb") as f:
            pickle.dump({
                "x": self._rbf_x,
                "y": self._rbf_y,
                "pitch_range": self._pitch_range,
                "yaw_range": self._yaw_range,
                "screen_x_range": self._screen_x_range,
                "screen_y_range": self._screen_y_range,
                "pitch_center": self._pitch_center,
                "yaw_center": self._yaw_center,
                "pitch_scale": self._pitch_scale,
                "yaw_scale": self._yaw_scale,
                "vert_center": self._vert_center,
                "vert_scale": self._vert_scale,
                "vert_feature": self._vert_feature,
                "screen_w": self.screen_w,
                "screen_h": self.screen_h,
            }, f)

    def load(self, path):
        import pickle
        with open(path, "rb") as f:
            data = pickle.load(f)
        self._rbf_x = data["x"]
        self._rbf_y = data["y"]
        self._pitch_range = data.get("pitch_range")
        self._yaw_range   = data.get("yaw_range")
        self._screen_x_range = data.get("screen_x_range")
        self._screen_y_range = data.get("screen_y_range")
        # .get() with fallback so an OLDER calibration.pkl (saved before
        # this normalization fix) still loads without crashing — it just
        # won't have the fix applied until you recalibrate.
        self._pitch_center = data.get("pitch_center", 0.0)
        self._yaw_center   = data.get("yaw_center", 0.0)
        self._pitch_scale  = data.get("pitch_scale", 1.0)
        self._yaw_scale    = data.get("yaw_scale", 1.0)
        self._vert_center  = data.get("vert_center", 0.0)
        self._vert_scale   = data.get("vert_scale", 1.0)
        self._vert_feature = data.get("vert_feature", "pitch")
        # Restore the screen size this calibration was actually built for.
        # An OLD calibration.pkl won't have it — warn loudly rather than
        # silently clamping to the stale 1493x933 defaults again.
        if "screen_w" in data and "screen_h" in data:
            self.screen_w = int(data["screen_w"])
            self.screen_h = int(data["screen_h"])
        else:
            print("[RBFGazeCalibration] WARNING: this calibration.pkl predates the "
                  "screen-size fix and has no screen dimensions stored. Predictions "
                  "will be clamped to the stale default "
                  f"{self.screen_w}x{self.screen_h}. Re-run calibrate.py.")
        self._calibrated = True


# =============================================================================
# FULL PIPELINE
# =============================================================================

class InsightUXPipeline:
    def __init__(self, onnx_path, calibration_path=None):
        self.gaze_model = GazeONNXModel(onnx_path)
        self.calibration = RBFGazeCalibration()

        if calibration_path:
            self.calibration.load(calibration_path)

    def predict_gaze_vector(self, left_patch, head_pose, right_patch=None):
        if right_patch is None:
            right_patch = left_patch

        pitch, yaw = self.gaze_model.predict(left_patch, right_patch, head_pose)

        gvx = -np.cos(pitch) * np.sin(yaw)
        gvy = -np.sin(pitch)

        return float(gvx), float(gvy), float(pitch), float(yaw)

    def predict_screen(self, left_patch, head_pose, right_patch=None, ear=0.0):
        _, _, pitch, yaw = self.predict_gaze_vector(left_patch, head_pose, right_patch)
        return self.calibration.predict(pitch, yaw, ear)