calibrate.py

import cv2
import numpy as np
import time
import os
import random
import pyautogui
import sys

# Windows consoles default to a non-UTF-8 codepage (cp1252). A single
# unicode character in a print() (the "⚠" warning glyph below, for example)
# then crashes the whole calibration run with UnicodeEncodeError — this is
# not hypothetical, it happened mid-session. errors="replace" means an
# unencodable glyph degrades to "?" instead of killing the process.
if hasattr(sys.stdout, "reconfigure"):
    sys.stdout.reconfigure(encoding="utf-8", errors="replace")
if hasattr(sys.stderr, "reconfigure"):
    sys.stderr.reconfigure(encoding="utf-8", errors="replace")

import warnings
warnings.filterwarnings("ignore", category=DeprecationWarning)

# Bundled read-only assets (models/checkpoints) vs. per-user writable data
# (calibration.pkl, baseline_pose.pkl) need different roots once frozen by
# PyInstaller (onedir nests bundled data under _internal/, alongside a
# persistent folder holding the actual InsightUX.exe).
if getattr(sys, "frozen", False):
    RESOURCE_DIR = sys._MEIPASS
    DATA_DIR     = os.path.dirname(sys.executable)
else:
    RESOURCE_DIR = os.path.dirname(os.path.abspath(__file__))
    DATA_DIR     = RESOURCE_DIR

# browser_session.py's start_calibration() sets this env var to the active
# InsightUX profile's own folder before launching this file as a subprocess,
# so calibration.pkl/baseline_pose.pkl land there instead of the global
# DATA_DIR above — keeping one user's calibration from overwriting another's.
# Unset (e.g. running `python calibrate.py` directly for development) means
# exactly today's behavior: writes go to the global DATA_DIR.
_USER_DATA_DIR = os.environ.get("INSIGHTUX_USER_DATA_DIR")
if _USER_DATA_DIR:
    DATA_DIR = _USER_DATA_DIR
    os.makedirs(DATA_DIR, exist_ok=True)

from preprocessing.preprocessing_pipeline import (
    create_face_mesh,
    estimate_camera_matrix,
    estimate_head_pose,
    compute_iris_radius,
    compute_ear,
    step1_normalize,
    step2_illumination,
    LEFT_EYE_INDICES,
    LEFT_EAR_INDICES,
    LEFT_IRIS_INDICES,
    RIGHT_EYE_INDICES,
    RIGHT_EAR_INDICES,
    RIGHT_IRIS_INDICES,
)

from inference_pipeline import InsightUXPipeline


# =============================================================================
# FIX A — POSE NORMALIZATION
# solvePnP returns degrees (typically ±30 deg at a laptop, now that the
# FACE_3D_MODEL Y/Z sign bug in preprocessing_pipeline.py is fixed).
# The CNN stream_b was trained expecting values near [-1, +1].
# Dividing by 30 maps ±30 deg -> ±1, matching training scale.
# Keep POSE_NORM_SCALE identical in main_webcam_pipeline.py.
# =============================================================================
POSE_NORM_SCALE = 30.0

def normalize_pose(head_pose):
    return np.array([
        head_pose.pitch / POSE_NORM_SCALE,
        head_pose.yaw   / POSE_NORM_SCALE,
        head_pose.roll  / POSE_NORM_SCALE,
    ], dtype=np.float32)


# =============================================================================
# FIX B — HEAD-PITCH COMPENSATION
# =============================================================================
HEAD_PITCH_COMPENSATION = 0.0  # reverted — 0.35 made accuracy worse, not better

def compensate_pitch(raw_pitch, head_pitch_deg):
    return raw_pitch - np.radians(head_pitch_deg) * HEAD_PITCH_COMPENSATION


# =============================================================================
# FIX C — HEAD-YAW COMPENSATION
# =============================================================================
HEAD_YAW_COMPENSATION = 0.0  # disabled

def compensate_yaw(raw_yaw, head_yaw_deg):
    return raw_yaw - np.radians(head_yaw_deg) * HEAD_YAW_COMPENSATION


# =============================================================================
# LIVE QUALITY THRESHOLDS (used during actual point collection, not just
# the pre-calibration gate)
# =============================================================================
LIGHT_MIN_BRIGHTNESS = 60
LIGHT_MAX_BRIGHTNESS = 200
BLINK_EAR_THRESHOLD  = 0.20
MIN_SAMPLES_OK        = 15   # below this, warn the point may be unreliable


# =============================================================================
# PRE-CALIBRATION CHECKS
# =============================================================================

def check_lighting(frame):
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    brightness = np.mean(gray)
    if brightness < LIGHT_MIN_BRIGHTNESS:
        return False, "Too dark - increase lighting"
    elif brightness > LIGHT_MAX_BRIGHTNESS:
        return False, "Too bright - reduce lighting"
    return True, None


def check_face_center(lms, frame_shape):
    H, W = frame_shape[:2]
    xs = [lm.x * W for lm in lms]
    ys = [lm.y * H for lm in lms]
    cx = np.mean(xs)
    cy = np.mean(ys)
    if cx < W * 0.35:
        return False, "Move face RIGHT in frame"
    elif cx > W * 0.65:
        return False, "Move face LEFT in frame"
    if cy < H * 0.35:
        return False, "Sit closer or lower your camera"
    elif cy > H * 0.65:
        return False, "Sit further or raise your camera"
    return True, None


def get_full_feedback(frame, lms, frame_shape, yaw, roll):
    msgs = []
    ok = True
    H, W = frame_shape[:2]

    nose_y     = lms[1].y   * H
    forehead_y = lms[10].y  * H
    chin_y     = lms[152].y * H
    midpoint_y = (forehead_y + chin_y) / 2.0
    face_h     = chin_y - forehead_y

    if face_h > 1:
        pitch_ratio = (nose_y - midpoint_y) / face_h
        if pitch_ratio < -0.10:
            msgs.append("Lift your head UP")
            ok = False
        elif pitch_ratio > 0.15:
            msgs.append("Tilt your head DOWN slightly")
            ok = False

    if yaw < -15:
        msgs.append("Turn face slightly RIGHT")
        ok = False
    elif yaw > 15:
        msgs.append("Turn face slightly LEFT")
        ok = False

    if roll < -10:
        msgs.append("Tilt head slightly RIGHT")
        ok = False
    elif roll > 10:
        msgs.append("Tilt head slightly LEFT")
        ok = False

    light_ok, light_msg = check_lighting(frame)
    if not light_ok:
        msgs.append(light_msg)
        ok = False

    center_ok, center_msg = check_face_center(lms, frame_shape)
    if not center_ok:
        msgs.append(center_msg)
        ok = False

    if ok:
        msgs.append("Perfect! Hold still...")

    return ok, msgs


# =============================================================================
# FACE ORIENTATION GATE
# =============================================================================

def face_orientation_gate(face_mesh, cap, cam_matrix_ref):
    print("\nChecking face setup before calibration...")
    print("Position your face straight, centred, at normal laptop distance.")

    HOLD_SECONDS   = 2.0
    good_since     = None
    cam_matrix_loc = cam_matrix_ref[0]

    cv2.namedWindow("Face Check", cv2.WINDOW_NORMAL)
    cv2.resizeWindow("Face Check", 640, 420)

    while True:
        ret, frame = cap.read()
        if not ret:
            continue

        if cam_matrix_loc is None:
            cam_matrix_loc    = estimate_camera_matrix(frame.shape)
            cam_matrix_ref[0] = cam_matrix_loc

        rgb     = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        results = face_mesh.process(rgb)
        display = frame.copy()
        H, W    = display.shape[:2]

        if not results.multi_face_landmarks:
            good_since = None
            cv2.putText(display, "No face detected - look at camera",
                        (20, 40), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
        else:
            lms       = results.multi_face_landmarks[0].landmark
            head_pose = estimate_head_pose(lms, frame.shape, cam_matrix_loc)

            if head_pose is None:
                good_since = None
                cv2.putText(display, "Pose failed - move slightly",
                            (20, 40), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
            else:
                is_good, msgs = get_full_feedback(
                    frame, lms, frame.shape, head_pose.yaw, head_pose.roll
                )

                if is_good:
                    if good_since is None:
                        good_since = time.time()
                    elapsed   = time.time() - good_since
                    remaining = max(0, HOLD_SECONDS - elapsed)

                    bar_w = int((elapsed / HOLD_SECONDS) * (W - 40))
                    bar_w = min(bar_w, W - 40)
                    cv2.rectangle(display, (20, H-50), (W-20, H-25), (40, 40, 40), -1)
                    cv2.rectangle(display, (20, H-50), (20+bar_w, H-25), (0, 220, 0), -1)
                    cv2.putText(display, f"Hold still... {remaining:.1f}s",
                                (20, 40), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)

                    if elapsed >= HOLD_SECONDS:
                        cv2.putText(display, "Starting calibration!",
                                    (20, 80), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
                        cv2.imshow("Face Check", display)
                        cv2.waitKey(800)
                        cv2.destroyWindow("Face Check")
                        return
                else:
                    good_since = None
                    for i, msg in enumerate(msgs):
                        cv2.putText(display, msg, (20, 40 + i * 35),
                                    cv2.FONT_HERSHEY_SIMPLEX, 0.75, (0, 165, 255), 2)

                cv2.putText(display,
                            f"Yaw:{head_pose.yaw:+.1f}  Roll:{head_pose.roll:+.1f}",
                            (20, H - 65), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (160, 160, 160), 1)

        cv2.imshow("Face Check", display)
        if cv2.waitKey(1) & 0xFF == 27:
            cv2.destroyWindow("Face Check")
            return


# =============================================================================
# FINE-TUNING TARGET ANGLES
# =============================================================================

def screen_to_gaze_angles(sx, sy, screen_w, screen_h, k_h=0.6, k_v=0.4):
    """
    Convert a screen point into the gaze angles the eyes must adopt to look
    at it. These are the TARGETS the fine-tune trains against, so getting the
    geometry right matters.

    k_h / k_v are tan(half-angle) horizontally and vertically:
        k = (half screen dimension) / (viewing distance)

    The defaults (0.6 / 0.4) are the ORIGINAL fabricated values and are wrong
    on two counts:
      * magnitude — roughly 1.7x too large for a laptop at arm's length
      * ratio     — 0.6/0.4 = 1.50, but a 16:9 screen demands 1.78 (its aspect
                    ratio), regardless of screen size or viewing distance.
    Training hard against a wrong ratio squashes one axis relative to the
    other: yaw discrimination collapsed (adjacent columns landed 0.011 apart,
    r fell 0.995 -> 0.760) while pitch was fine.

    calibrate.py now computes k_h / k_v from the real screen dimensions and
    the viewing distance measured by solvePnP, and passes them in.
    """
    nx    = (sx - screen_w  / 2) / (screen_w  / 2)
    ny    = (sy - screen_h / 2) / (screen_h / 2)
    pitch = float(np.arctan(ny * k_v))
    yaw   = float(np.arctan(nx * k_h))
    return pitch, yaw


def finetune_on_calibration_v4(
    onnx_path,
    left_patches, right_patches, head_poses, screen_points,
    screen_w, screen_h,
    ckpt_path="checkpoints/best_model_v4.pt",
    out_onnx_path=None,
    steps=300, lr=5e-4,
    k_h=0.6, k_v=0.4,
):
    try:
        import torch
        import sys
        sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
        from models.model_v4 import GazeCNNv4
    except ImportError as e:
        print(f"[Fine-tune] Skipped: {e}")
        return False

    if not os.path.exists(ckpt_path):
        print(f"[Fine-tune] Skipped: checkpoint not found at {ckpt_path}")
        return False

    if out_onnx_path is None:
        out_onnx_path = onnx_path

    device = torch.device("cpu")
    model  = GazeCNNv4().to(device)
    ckpt   = torch.load(ckpt_path, map_location=device)
    model.load_state_dict(ckpt["model_state_dict"])

    for p in model.parameters():
        p.requires_grad = False
    for p in model.fc_a.parameters():     p.requires_grad = True
    for p in model.stream_b.parameters(): p.requires_grad = True
    for p in model.fusion.parameters():   p.requires_grad = True

    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"[Fine-tune] Trainable: {trainable:,} parameters")

    optimizer = torch.optim.Adam(
        filter(lambda p: p.requires_grad, model.parameters()),
        lr=lr, weight_decay=1e-4
    )

    def angular_loss(pred, target):
        def to_vec(a):
            p, y = a[:, 0], a[:, 1]
            v = torch.stack([torch.cos(p)*torch.sin(y),
                             torch.sin(p),
                             torch.cos(p)*torch.cos(y)], dim=1)
            return v / (v.norm(dim=1, keepdim=True) + 1e-8)
        cos_sim = (to_vec(pred) * to_vec(target)).sum(dim=1)
        return (1 - cos_sim).mean()

    patches_l, poses_t, gazes = [], [], []
    for l, r, pose, sp in zip(left_patches, right_patches, head_poses, screen_points):
        def norm(x): return (x.astype(np.float32) / 255.0 - 0.5) / 0.5
        patches_l.append(np.stack([norm(l), norm(r)], axis=0))
        poses_t.append(pose)
        gazes.append(screen_to_gaze_angles(sp[0], sp[1], screen_w, screen_h, k_h, k_v))

    patches_t = torch.tensor(np.array(patches_l), dtype=torch.float32)
    poses_t   = torch.tensor(np.array(poses_t),   dtype=torch.float32)
    gazes_t   = torch.tensor(np.array(gazes),     dtype=torch.float32)

    model.eval()
    with torch.no_grad():
        feat_a_raw = model.pool(model.backbone(patches_t)).flatten(1)

    model.fc_a.train()
    model.stream_b.train()
    model.fusion.train()

    print(f"[Fine-tune] {steps} steps on {len(patches_l)} samples (backbone cached)...")
    for step in range(steps):
        optimizer.zero_grad()
        feat_a = model.fc_a(feat_a_raw)
        feat_b = model.stream_b(poses_t)
        fused  = torch.cat([feat_a, feat_b], dim=1)
        pred   = model.fusion(fused)
        loss   = angular_loss(pred, gazes_t)
        loss.backward()
        optimizer.step()
        if (step + 1) % 50 == 0:
            print(f"  Step {step+1}/{steps} | Loss: {loss.item():.5f}")

    model.eval()
    dp    = torch.zeros(1, 2, 36, 60)
    dpose = torch.zeros(1, 3)
    torch.onnx.export(
        model, (dp, dpose), out_onnx_path,
        input_names  = ["eye_patch_binocular", "head_pose"],
        output_names = ["gaze"],
        dynamic_axes = {"eye_patch_binocular": {0: "batch"},
                        "head_pose": {0: "batch"}, "gaze": {0: "batch"}},
        opset_version=14,
        dynamo=False,
    )
    print(f"[Fine-tune] Adapted ONNX saved: {out_onnx_path}")
    return True


# =============================================================================
# OUTLIER FILTERING
# =============================================================================

def filter_unreliable_points(gaze_vectors, screen_points, head_pitches_deg,
                              point_dispersion, point_frames, point_ear, factor=3.0):
    """
    Drop calibration points that were INTERNALLY NOISY, i.e. the model's
    prediction wobbled wildly while you held your gaze on that one dot.

    This replaces the previous filter, which removed points whose gaze angle
    was far from the median of all points. That was backwards: the corner
    points are SUPPOSED to be far from the median — that's what makes them
    corners. That filter's failure mode was to preferentially delete the
    screen extremes, which are exactly the regions we're trying to fix.

    A point is unreliable if its own frame-to-frame dispersion is much larger
    than the typical point's. That's an honest reliability signal and it is
    completely independent of WHERE on screen the point sits, so corners are
    no longer penalized for being corners.
    """
    disp = np.array(point_dispersion)              # (n, 2) -> [mad_pitch, mad_yaw]
    combined = disp[:, 0] + disp[:, 1]
    med_disp = float(np.median(combined))
    if med_disp <= 1e-9:
        keep = np.ones(len(combined), dtype=bool)
    else:
        keep = combined <= (factor * med_disp)

    n_removed = int((~keep).sum())
    if n_removed:
        removed = [i + 1 for i in np.where(~keep)[0]]
        print(f"[Reliability filter] Removed {n_removed} point(s) with excessive "
              f"frame-to-frame jitter: {removed}")

    idx = np.where(keep)[0].tolist()
    return (
        [gaze_vectors[i]     for i in idx],
        [screen_points[i]    for i in idx],
        [head_pitches_deg[i] for i in idx],
        [point_dispersion[i] for i in idx],
        [point_frames[i]     for i in idx],
        [point_ear[i]        for i in idx],
    )


# =============================================================================
# CONFIG
# =============================================================================

ONNX_PATH = os.path.join(RESOURCE_DIR, "models", "gaze_cnn_v4.onnx")
CKPT_PATH = os.path.join(RESOURCE_DIR, "checkpoints", "best_model_v4.pt")

# Fine-tuning (below) used to save its adapted model back over ONNX_PATH
# itself — i.e. it overwrote the one bundled model every profile's live
# inference loads, so one person calibrating silently changed what every
# other profile's eye-tracking used. USER_ONNX_OUT_PATH is that same
# browser_session.py-supplied env var as DATA_DIR above: when set, the
# fine-tuned model is saved per-profile instead. Unset (standalone/dev
# runs) reproduces the exact previous behavior — out_onnx_path=ONNX_PATH.
USER_ONNX_OUT_PATH = os.environ.get("INSIGHTUX_USER_ONNX_OUT") or ONNX_PATH


def _current_onnx_path():
    """Prefer this profile's own previously fine-tuned model, if a prior
    calibration run produced one, over the stock bundled model — so the
    live preview during point collection matches what real tracking will
    actually use. No-op (always ONNX_PATH) when USER_ONNX_OUT_PATH isn't
    set to a distinct per-user path, i.e. standalone/dev runs."""
    if USER_ONNX_OUT_PATH != ONNX_PATH and os.path.exists(USER_ONNX_OUT_PATH):
        return USER_ONNX_OUT_PATH
    return ONNX_PATH

# --- Fine-tune target geometry -------------------------------------------
# USE_MEASURED_GEOMETRY=False keeps the original hand-picked constants
# (k_h=0.6, k_v=0.4). Those are geometrically wrong in the abstract, but the
# setup using them is currently WORKING, and a working system beats a
# theoretically tidier one. Do not change this at the same time as anything
# else, or we won't know what caused what.
#
# Set True to A/B test targets derived from your real screen size and the
# viewing distance solvePnP measures. If you do, set SCREEN_DIAGONAL_INCHES
# to your actual laptop diagonal first.
USE_MEASURED_GEOMETRY  = False
SCREEN_DIAGONAL_INCHES = 15.6

SCREEN_W, SCREEN_H = pyautogui.size()

points = [
    (0.02, 0.02), (0.35, 0.02), (0.65, 0.02), (0.98, 0.02),
    (0.02, 0.35), (0.35, 0.35), (0.65, 0.35), (0.98, 0.35),
    (0.02, 0.65), (0.35, 0.65), (0.65, 0.65), (0.98, 0.65),
    (0.02, 0.98), (0.35, 0.98), (0.65, 0.98), (0.98, 0.98),
]

def get_duration(py):
    # Top AND bottom rows get more time. Bottom rows specifically need it:
    # looking down partially occludes the iris under the eyelid, which
    # means more frames get blink-skipped there than in the middle rows,
    # so the same 3s window yields fewer usable samples unless we give it
    # more time to compensate.
    return 5 if (py <= 0.4 or py >= 0.6) else 3

random.Random(7).shuffle(points)

# How many genuinely different frames to keep per calibration point. These
# feed BOTH the fine-tune and the RBF refit. Previously this was 5, but the
# frames were then discarded and replaced with 5 copies of a single frame —
# so the fine-tune saw 16 unique images while believing it had 80.
FRAMES_KEPT_PER_POINT = 12

# FINE-TUNING: back ON.
#
# It was briefly disabled on the theory that it was corrupting the model.
# That theory was WRONG and the data refuted it: with fine-tuning on, mean
# validation error was 169-204px; with it off, 275px. Worse.
#
# The instability seen earlier (model output flipping sign between runs) is
# better explained by the duplication bug: the fine-tune was being fed the
# SAME frame 5x per point, i.e. 16 unique images, which overfits trivially.
# That bug is now fixed — it receives FRAMES_KEPT_PER_POINT genuinely
# different frames per point (16 x 12 = ~192 real samples), so fine-tuning
# is now doing what it was always supposed to do.
#
# The fabricated targets in screen_to_gaze_angles() (arctan(ny*0.4) etc.)
# remain a real weakness — those constants encode no true geometry. But
# empirically the fine-tune still helps, so it stays on until something
# better replaces it.
FINETUNE_ENABLED = True


def main():
    # =============================================================================
    # MAIN CALIBRATION
    # =============================================================================

    pipeline  = InsightUXPipeline(_current_onnx_path())
    face_mesh = create_face_mesh(static_image_mode=False)
    cap       = cv2.VideoCapture(0)

    cam_matrix = None

    gaze_vectors     = []   # per-point median [pitch, yaw] (pre-fine-tune)
    screen_points    = []   # per-point [sx, sy]
    head_pitches_deg = []   # per-point median head pitch (deg)
    point_dispersion = []   # per-point [mad_pitch, mad_yaw]  <- the noise floor
    point_frames     = []   # per-point dict of REAL diverse frames (see below)
    point_ear        = []   # per-point median eye-aperture (candidate vertical cue)

    # Session-wide diagnostics, so we can warn about systemic issues at the end
    # rather than just per-point.
    session_brightness   = []
    session_distance_mm  = []   # viewing distance from solvePnP — real geometry, not a guess
    session_blink_skips  = 0
    session_light_skips  = 0
    points_with_low_yield = []

    cam_matrix_ref = [None]
    face_orientation_gate(face_mesh, cap, cam_matrix_ref)
    if cam_matrix_ref[0] is not None:
        cam_matrix = cam_matrix_ref[0]

    cv2.namedWindow("Calibration", cv2.WINDOW_NORMAL)
    cv2.setWindowProperty("Calibration", cv2.WND_PROP_FULLSCREEN, cv2.WINDOW_FULLSCREEN)

    total = len(points)
    print(f"Calibration started - {total} points")

    for idx, (px, py) in enumerate(points):
        sx, sy   = int(px * SCREEN_W), int(py * SCREEN_H)
        duration = get_duration(py)

        settle_s = 1.5 if idx == 0 else 0.4
        settle_start = time.time()
        while time.time() - settle_start < settle_s:
            ret, frame = cap.read()
            if not ret:
                continue
            screen = np.zeros((SCREEN_H, SCREEN_W, 3), dtype=np.uint8)
            cv2.circle(screen, (sx, sy), 18, (80, 80, 80), -1)
            msg = "Get ready..." if idx == 0 else "Settling..."
            cv2.putText(screen, msg, (50, 50),
                        cv2.FONT_HERSHEY_SIMPLEX, 1, (180, 180, 180), 2)
            cv2.imshow("Calibration", screen)
            cv2.waitKey(1)

        samples    = []
        ear_list   = []
        l_list     = []
        r_list     = []
        p_list     = []
        hp_list    = []

        point_blink_skips = 0
        point_light_skips = 0

        start = time.time()

        while time.time() - start < duration:
            ret, frame = cap.read()
            if not ret:
                continue

            if cam_matrix is None:
                cam_matrix = estimate_camera_matrix(frame.shape)

            # LIVE LIGHTING CHECK — skip this frame and tell the person, instead
            # of silently feeding a too-dark/too-bright frame into the model.
            light_ok, light_msg = check_lighting(frame)
            gray_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            session_brightness.append(float(np.mean(gray_frame)))
            if not light_ok:
                point_light_skips += 1
                screen  = np.zeros((SCREEN_H, SCREEN_W, 3), dtype=np.uint8)
                cv2.circle(screen, (sx, sy), 18, (0, 140, 255), 2)
                cv2.putText(screen, f"Point {idx+1}/{total} - {light_msg}",
                            (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 140, 255), 2)
                cv2.imshow("Calibration", screen)
                cv2.waitKey(1)
                continue

            rgb     = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            results = face_mesh.process(rgb)

            if not results.multi_face_landmarks:
                continue

            lms       = results.multi_face_landmarks[0].landmark
            head_pose = estimate_head_pose(lms, frame.shape, cam_matrix)
            if head_pose is None:
                continue

            # LIVE BLINK CHECK — average EAR across both eyes. Below threshold
            # means eyes are closed/closing; that frame's eye patches carry no
            # real gaze signal and would just add noise to this point's median.
            ear_l = compute_ear(lms, LEFT_EAR_INDICES, frame.shape)
            ear_r = compute_ear(lms, RIGHT_EAR_INDICES, frame.shape)
            avg_ear = (ear_l + ear_r) / 2.0
            if avg_ear < BLINK_EAR_THRESHOLD:
                point_blink_skips += 1
                screen  = np.zeros((SCREEN_H, SCREEN_W, 3), dtype=np.uint8)
                cv2.circle(screen, (sx, sy), 18, (180, 180, 0), 2)
                cv2.putText(screen, f"Point {idx+1}/{total} - Blink detected, skipping frame",
                            (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (180, 180, 0), 2)
                cv2.imshow("Calibration", screen)
                cv2.waitKey(1)
                continue

            # FIX A: normalize pose to [-1, +1] before CNN
            pose_vec = normalize_pose(head_pose)

            def process_eye(eye_idx, ear_idx, iris_idx):
                s1 = step1_normalize(frame, lms, head_pose, eye_idx, ear_idx, iris_idx)
                if not s1.is_open:
                    return None
                ir = compute_iris_radius(lms, iris_idx, frame.shape)
                s2 = step2_illumination(s1, ir)
                return s2.blended if s2.is_usable else None

            l = process_eye(LEFT_EYE_INDICES,  LEFT_EAR_INDICES,  LEFT_IRIS_INDICES)
            r = process_eye(RIGHT_EYE_INDICES, RIGHT_EAR_INDICES, RIGHT_IRIS_INDICES)

            if l is None and r is None:
                continue
            if l is None: l = r
            if r is None: r = l

            _, _, raw_pitch, raw_yaw = pipeline.predict_gaze_vector(l, pose_vec, r)

            pitch = compensate_pitch(raw_pitch, head_pose.pitch)
            yaw   = compensate_yaw(raw_yaw, head_pose.yaw)

            samples.append([pitch, yaw])
            ear_list.append(avg_ear)
            # solvePnP's tvec is in the same units as FACE_3D_MODEL (mm), so
            # tvec[2] is a genuine measurement of how far your face is from the
            # camera. We were computing this every frame and throwing it away.
            try:
                session_distance_mm.append(float(head_pose.tvec[2]))
            except Exception:
                pass
            l_list.append(l)
            r_list.append(r)
            p_list.append(pose_vec)
            hp_list.append(head_pose.pitch)

            elapsed = time.time() - start
            screen  = np.zeros((SCREEN_H, SCREEN_W, 3), dtype=np.uint8)
            angle   = int(360 * elapsed / duration)
            cv2.ellipse(screen, (sx, sy), (28, 28), -90, 0, angle, (0, 180, 0), 3)
            cv2.circle(screen, (sx, sy), 18, (0, 255, 0), -1)
            cv2.putText(screen, f"Point {idx+1}/{total} - Look at the dot",
                        (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)
            if point_blink_skips or point_light_skips:
                cv2.putText(screen, f"skipped: {point_blink_skips} blink, {point_light_skips} lighting",
                            (50, 85), cv2.FONT_HERSHEY_SIMPLEX, 0.55, (140, 140, 140), 1)
            cv2.imshow("Calibration", screen)

            if cv2.waitKey(1) & 0xFF == 27:
                break

        session_blink_skips += point_blink_skips
        session_light_skips += point_light_skips

        if not samples:
            print(f"Point {idx+1}: no samples, skipping.")
            points_with_low_yield.append(idx + 1)
            continue

        if len(samples) < MIN_SAMPLES_OK:
            print(f"  ⚠ Point {idx+1}: only {len(samples)} valid samples "
                  f"({point_blink_skips} blink-skipped, {point_light_skips} light-skipped) — may be unreliable")
            points_with_low_yield.append(idx + 1)

        samples_arr = np.array(samples)
        avg         = np.median(samples_arr, axis=0)
        avg_hp_deg  = float(np.median(hp_list))

        # WITHIN-POINT DISPERSION = the noise floor. While you stared at ONE
        # fixed dot, how much did the model's predicted pitch/yaw wobble
        # frame-to-frame? This is the single most important number in the whole
        # calibration and it was never being measured. Compare it against the
        # BETWEEN-point spread (the actual signal) at the end of the run: if
        # noise >= signal on an axis, that axis is unrecoverable by any
        # calibration math, and no amount of RBF tuning will fix it.
        # MAD (median absolute deviation) rather than std, so one blink-tail
        # frame can't inflate it.
        mad_pitch = float(np.median(np.abs(samples_arr[:, 0] - avg[0])))
        mad_yaw   = float(np.median(np.abs(samples_arr[:, 1] - avg[1])))

        gaze_vectors.append(list(avg))
        screen_points.append([sx, sy])
        head_pitches_deg.append(avg_hp_deg)
        point_dispersion.append([mad_pitch, mad_yaw])
        # Eye aperture for this point. Looking DOWN lowers the eyelid, so this is
        # a physically independent vertical cue — and the CNN's pitch output has
        # proven nearly blind vertically (r~0.41). calibrate() will measure both
        # and pick whichever actually tracks screen-Y.
        point_ear.append(float(np.median(ear_list)))

        # Keep a BUNDLE of genuinely different frames for this point — not one
        # "representative" frame. Both the fine-tune and the post-fine-tune RBF
        # refit read from this bundle, so both get real frame diversity and real
        # noise averaging instead of anchoring on a single arbitrary frame.
        n_avail = len(l_list)
        n_take  = min(FRAMES_KEPT_PER_POINT, n_avail)
        indices = np.linspace(0, n_avail - 1, n_take, dtype=int)
        point_frames.append({
            "lefts":  [l_list[i]  for i in indices],
            "rights": [r_list[i]  for i in indices],
            "poses":  [p_list[i]  for i in indices],
            "hps":    [hp_list[i] for i in indices],
        })

        print(f"Point {idx+1:2d}/{total} | pitch={avg[0]:+.4f} yaw={avg[1]:+.4f} "
              f"| jitter(pitch)={mad_pitch:.4f} jitter(yaw)={mad_yaw:.4f} "
              f"| {len(indices)} frames kept "
              f"| skipped: {point_blink_skips} blink, {point_light_skips} lighting")

    cap.release()
    cv2.destroyAllWindows()

    if len(gaze_vectors) < 4:
        print(f"ERROR: only {len(gaze_vectors)} points collected, need at least 4.")
        exit(1)

    # =============================================================================
    # OUTLIER FILTER
    # =============================================================================
    (gaze_vectors, screen_points, head_pitches_deg,
     point_dispersion, point_frames, point_ear) = filter_unreliable_points(
        gaze_vectors, screen_points, head_pitches_deg,
        point_dispersion, point_frames, point_ear, factor=3.0
    )

    n_pts = len(screen_points)
    if n_pts < 4:
        print("ERROR: too many points removed, need at least 4 reliable points.")
        exit(1)

    # =============================================================================
    # FINE-TUNE (optional — see FINETUNE_ENABLED comment for why it defaults OFF)
    # Now fed the REAL diverse frames from each point's bundle. Previously this
    # received the same single frame duplicated 5x per point, so it trained on
    # 16 unique images while reporting 80.
    # =============================================================================
    finetuned = False
    if FINETUNE_ENABLED:
        ft_lefts, ft_rights, ft_poses, ft_sp = [], [], [], []
        for i in range(n_pts):
            pf = point_frames[i]
            for k in range(len(pf["lefts"])):
                ft_lefts.append(pf["lefts"][k])
                ft_rights.append(pf["rights"][k])
                ft_poses.append(pf["poses"][k])
                ft_sp.append(screen_points[i])

        n_unique = len(ft_lefts)

        # Derive the fine-tune target geometry from REAL numbers instead of the
        # fabricated 0.6/0.4 constants.
        _diag_mm  = SCREEN_DIAGONAL_INCHES * 25.4
        _aspect   = SCREEN_W / float(SCREEN_H)
        _scr_h_mm = _diag_mm / np.sqrt(_aspect ** 2 + 1.0)
        _scr_w_mm = _aspect * _scr_h_mm
        if session_distance_mm:
            _dist_mm = float(np.median(session_distance_mm))
        else:
            _dist_mm = 500.0
        # Guard against a wild solvePnP outlier producing nonsense geometry.
        _dist_mm = float(np.clip(_dist_mm, 300.0, 900.0))

        if USE_MEASURED_GEOMETRY:
            K_H = (_scr_w_mm / 2.0) / _dist_mm
            K_V = (_scr_h_mm / 2.0) / _dist_mm
        else:
            K_H, K_V = 0.6, 0.4   # original constants — currently working

        print(f"\n--- Fine-tuning CNN on your eyes ({n_unique} REAL frames, "
              f"{n_pts} points) ---")
        print(f"[Geometry] screen {SCREEN_DIAGONAL_INCHES}\" -> "
              f"{_scr_w_mm:.0f}x{_scr_h_mm:.0f}mm | measured viewing distance "
              f"{_dist_mm:.0f}mm")
        print(f"[Geometry] fine-tune targets: k_h={K_H:.3f}  k_v={K_V:.3f}  "
              f"(ratio {K_H/max(K_V,1e-6):.2f})  "
              f"[{'measured' if USE_MEASURED_GEOMETRY else 'original constants'}]")
        finetuned = finetune_on_calibration_v4(
            onnx_path     = ONNX_PATH,
            left_patches  = ft_lefts,
            right_patches = ft_rights,
            head_poses    = ft_poses,
            screen_points = ft_sp,
            screen_w      = SCREEN_W,
            screen_h      = SCREEN_H,
            ckpt_path     = CKPT_PATH,
            out_onnx_path = USER_ONNX_OUT_PATH,
            steps         = 200,
            lr            = 1e-4,
            k_h           = K_H,
            k_v           = K_V,
        )
    else:
        print("\n--- Fine-tuning SKIPPED (FINETUNE_ENABLED = False) ---")
        print("    Using the base model as-is. The RBF maps its output to screen")
        print("    coordinates, so the model's absolute scale does not need to be")
        print("    'correct' — only monotonic. See the FINETUNE_ENABLED comment.")

    # =============================================================================
    # FIT RBF CALIBRATION
    #
    # CRITICAL FIX: re-predict across EVERY kept frame for each point and take the
    # MEDIAN, instead of predicting from one arbitrary "representative" frame.
    #
    # The old code computed a careful median over 50-150 frames during collection,
    # then threw it away and refit the RBF from a single frame per point. That
    # discarded all noise averaging and anchored the entire calibration on 16
    # single, possibly-noisy frames — which is exactly the kind of thing that
    # makes edges and corners unstable.
    # =============================================================================
    print("\n--- Fitting RBF calibration ---")
    pipeline = InsightUXPipeline(USER_ONNX_OUT_PATH if finetuned else _current_onnx_path())

    adapted_gaze_vectors = []
    adapted_dispersion   = []
    for i in range(n_pts):
        pf = point_frames[i]
        per_frame = []
        for k in range(len(pf["lefts"])):
            l    = pf["lefts"][k]
            r    = pf["rights"][k]
            pose = pf["poses"][k]
            hp   = pf["hps"][k]
            hy   = float(pose[1]) * POSE_NORM_SCALE   # recover raw head yaw (deg)

            _, _, raw_pitch, raw_yaw = pipeline.predict_gaze_vector(l, pose, r)
            pitch = compensate_pitch(raw_pitch, hp)
            yaw   = compensate_yaw(raw_yaw, hy)
            per_frame.append([pitch, yaw])

        per_frame = np.array(per_frame)
        med = np.median(per_frame, axis=0)
        adapted_gaze_vectors.append(list(med))
        adapted_dispersion.append([
            float(np.median(np.abs(per_frame[:, 0] - med[0]))),
            float(np.median(np.abs(per_frame[:, 1] - med[1]))),
        ])

    assert len(adapted_gaze_vectors) == len(screen_points), \
        f"Length mismatch: {len(adapted_gaze_vectors)} gaze vs {len(screen_points)} screen"

    pipeline.calibration.calibrate(
        np.array(adapted_gaze_vectors),
        np.array(screen_points),
        ear=np.array(point_ear),
        screen_size=(SCREEN_W, SCREEN_H)   # <-- fixes the 1493x933 clamp bug
    )
    pipeline.calibration.save(os.path.join(DATA_DIR, "calibration.pkl"))

    import pickle
    _all_hp   = [hp for pf in point_frames for hp in pf["hps"]]
    _all_pose = [p  for pf in point_frames for p  in pf["poses"]]
    baseline_pitch = float(np.median(_all_hp)) if _all_hp else 0.0
    baseline_yaw = float(np.median([float(p[1]) * POSE_NORM_SCALE for p in _all_pose])) if _all_pose else 0.0
    baseline_roll = float(np.median([float(p[2]) * POSE_NORM_SCALE for p in _all_pose])) if _all_pose else 0.0
    with open(os.path.join(DATA_DIR, "baseline_pose.pkl"), "wb") as f:
        pickle.dump({"pitch": baseline_pitch, "yaw": baseline_yaw, "roll": baseline_roll}, f)
    print(f"Saved baseline_pose.pkl: pitch={baseline_pitch:+.2f}, yaw={baseline_yaw:+.2f}, roll={baseline_roll:+.2f}")

    print(f"\nCalibration complete! {n_pts}/{total} points used.")
    if finetuned:
        print("CNN fine-tuned on your eyes + RBF calibration fitted.")
    else:
        print("RBF calibration fitted (CNN fine-tuning skipped).")

    # =============================================================================
    # SESSION DIAGNOSTIC SUMMARY — the "tell the person if something's off" ask.
    # Printed once, in plain language, instead of them having to interpret
    # per-point numbers themselves.
    # =============================================================================
    # =============================================================================
    # SIGNAL-TO-NOISE — THE VERDICT
    #
    # This is the number that decides whether an axis is fixable at all.
    #
    #   SIGNAL = how much the model's output changes between DIFFERENT screen
    #            positions (spread of the per-point medians). This is the real
    #            information the RBF has to work with.
    #   NOISE  = how much the model's output wobbles frame-to-frame while you
    #            stare at ONE fixed dot (median within-point jitter).
    #
    # If NOISE >= SIGNAL on an axis, the model cannot tell "you looked lower"
    # apart from "the same look, one frame later". No RBF, no gain correction,
    # no clamp tuning can recover information that was never there. That axis
    # needs a better MODEL, not better calibration math.
    # =============================================================================
    gv_arr = np.array(adapted_gaze_vectors)
    disp_arr = np.array(adapted_dispersion)
    sp_arr = np.array(screen_points)

    signal_pitch = float(np.std(gv_arr[:, 0]))
    signal_yaw   = float(np.std(gv_arr[:, 1]))
    noise_pitch  = float(np.median(disp_arr[:, 0]))
    noise_yaw    = float(np.median(disp_arr[:, 1]))

    snr_pitch = signal_pitch / noise_pitch if noise_pitch > 1e-9 else float("inf")
    snr_yaw   = signal_yaw   / noise_yaw   if noise_yaw   > 1e-9 else float("inf")

    # CORRELATION — the metric that actually decides usability.
    #
    # An earlier version of this block reported only SNR (spread vs jitter). That
    # was MISLEADING and nearly sent us chasing the wrong fix: a signal can have
    # plenty of spread and still be useless if that spread is not ORDERED by the
    # thing you're predicting. Pitch once reported SNR 1.76 ("marginal") while its
    # per-row averages went down, up, up — i.e. no monotonic relationship with
    # screen row at all. Spread was real; ordering was not.
    #
    # Pearson r between model output and true screen coordinate is the honest
    # test. r near +/-1 means the model tracks that axis. r near 0 means it does
    # not, no matter how much spread there is.
    def _pearson(a, b):
        a = np.asarray(a, dtype=float); b = np.asarray(b, dtype=float)
        if a.std() < 1e-12 or b.std() < 1e-12:
            return 0.0
        return float(np.corrcoef(a, b)[0, 1])

    r_yaw   = _pearson(gv_arr[:, 1], sp_arr[:, 0])   # yaw   vs screen X
    r_pitch = _pearson(gv_arr[:, 0], sp_arr[:, 1])   # pitch vs screen Y
    r_ear   = _pearson(np.array(point_ear), sp_arr[:, 1])   # eye aperture vs screen Y

    def _corr_verdict(r):
        ar = abs(r)
        if ar >= 0.90:
            return "EXCELLENT — model tracks this axis cleanly"
        if ar >= 0.70:
            return "GOOD — usable, some slop"
        if ar >= 0.50:
            return "WEAK — expect significant error on this axis"
        return "BROKEN — model output barely relates to this axis at all"

    print("\n========== DOES THE MODEL TRACK THE SCREEN? (the real test) ==========")
    print(f"HORIZONTAL  yaw vs screen-X : r = {r_yaw:+.3f}   (SNR {snr_yaw:.2f})")
    print(f"                             -> {_corr_verdict(r_yaw)}")
    print(f"VERTICAL    pitch vs screen-Y: r = {r_pitch:+.3f}   (SNR {snr_pitch:.2f})")
    print(f"                             -> {_corr_verdict(r_pitch)}")
    print(f"VERTICAL    EYE APERTURE vs screen-Y: r = {r_ear:+.3f}")
    print(f"                             -> {_corr_verdict(r_ear)}")
    print()
    _best_vert = max(abs(r_pitch), abs(r_ear))
    print(f"Best available vertical cue: "
          f"{'EYE APERTURE' if abs(r_ear) > abs(r_pitch) else 'CNN PITCH'} "
          f"(r={_best_vert:+.3f})")
    print()
    if _best_vert < 0.5:
        print("VERDICT: the model's PITCH output does not meaningfully track where you")
        print("look vertically. This is NOT a calibration problem — the RBF cannot map")
        print("an input that carries no ordered information about screen height. No")
        print("amount of clamp/smoothing/gain tuning will fix it. The fix is the MODEL:")
        print("its pitch head needs retraining, or vertical gaze needs a different")
        print("feature (e.g. eyelid aperture / iris-centre offset within the socket),")
        print("which the current eye-patch CNN is evidently not learning.")
    elif _best_vert < abs(r_yaw) - 0.15:
        print("VERDICT: vertical tracks the screen, but noticeably worse than")
        print("horizontal. Calibration is doing its job; expect up/down to stay the")
        print("looser axis until the model improves.")
    else:
        print("VERDICT: both axes track the screen. Any remaining error is in the")
        print("calibration mapping or the noise floor, not in the model's ability")
        print("to see where you're looking.")
    print("=====================================================================")

    print("\n================ CALIBRATION QUALITY SUMMARY ================")
    if session_brightness:
        avg_bright = float(np.mean(session_brightness))
        pct_dark  = 100 * sum(1 for b in session_brightness if b < LIGHT_MIN_BRIGHTNESS) / len(session_brightness)
        pct_bright = 100 * sum(1 for b in session_brightness if b > LIGHT_MAX_BRIGHTNESS) / len(session_brightness)
        print(f"Average brightness: {avg_bright:.0f} (comfortable range: {LIGHT_MIN_BRIGHTNESS}-{LIGHT_MAX_BRIGHTNESS})")
        if pct_dark > 10:
            print(f"⚠ Lighting was too DARK for {pct_dark:.0f}% of frames. "
                  f"Add a light source facing your face, or face a window, before recalibrating.")
        if pct_bright > 10:
            print(f"⚠ Lighting was too BRIGHT for {pct_bright:.0f}% of frames "
                  f"(backlight or a light directly behind you?). Try facing away from strong light sources.")
        if pct_dark <= 10 and pct_bright <= 10:
            print("Lighting was consistently good throughout.")

    if session_blink_skips > 0:
        print(f"Blinking accounted for {session_blink_skips} skipped frames across the session "
              f"— normal, this is expected and was handled automatically.")

    if points_with_low_yield:
        print(f"⚠ These points had low sample counts and may be less accurate: "
              f"{', '.join(str(p) for p in points_with_low_yield)}. "
              f"If tracking feels off in that part of the screen, consider recalibrating.")
    else:
        print("All points collected a healthy number of samples.")
    print("===============================================================")

    # --- Diagnostic: verify vertical pitch separation ---
    print("\n--- Pitch by screen row (should increase top -> bottom) ---")
    gv_arr = np.array(adapted_gaze_vectors)
    sp_arr = np.array(screen_points)
    thresholds = [(0,          SCREEN_H*0.25, "Top    (y<25%)  "),
                  (SCREEN_H*0.25, SCREEN_H*0.5,  "Mid-hi (25-50%) "),
                  (SCREEN_H*0.5,  SCREEN_H*0.75, "Mid-lo (50-75%) "),
                  (SCREEN_H*0.75, SCREEN_H+1,    "Bottom (y>75%)  ")]
    for lo, hi, label in thresholds:
        mask = (sp_arr[:, 1] >= lo) & (sp_arr[:, 1] < hi)
        if mask.any():
            print(f"  {label}: avg pitch = {gv_arr[mask, 0].mean():.4f}")

    print("\n--- Yaw by screen column (should increase left -> right) ---")
    col_thresholds = [(0,            SCREEN_W*0.25, "Left   (x<25%)  "),
                       (SCREEN_W*0.25, SCREEN_W*0.5,  "Mid-lf (25-50%) "),
                       (SCREEN_W*0.5,  SCREEN_W*0.75, "Mid-rt (50-75%) "),
                       (SCREEN_W*0.75, SCREEN_W+1,    "Right  (x>75%)  ")]
    for lo, hi, label in col_thresholds:
        mask = (sp_arr[:, 0] >= lo) & (sp_arr[:, 0] < hi)
        if mask.any():
            print(f"  {label}: avg yaw = {gv_arr[mask, 1].mean():.4f}")


if __name__ == "__main__":
    main()