scripts/reconstruct_v16_frame_mapping.py

"""Reconstruct V16 H5 row-to-MPIIGaze-frame mappings from labels and landmarks.

Uses the exact archived preprocessing equations. It does not rerun MediaPipe.
"""
from pathlib import Path
import sys

ROOT = Path(r"E:\Gaze_estimation")
sys.path.insert(0, str(ROOT / ".codex_deps"))
import h5py
import numpy as np
import pandas as pd
from PIL import Image

RAW = ROOT / "data" / "MPIIGaze" / "MPIIGaze" / "MPIIGaze" / "Data" / "Original"
OUT = ROOT / "artifacts" / "kd-teacher-trap-diagnostic"
LEFT_CORNERS = (362, 263)
RIGHT_CORNERS = (33, 133)


def parse(line):
    parts = line.split()
    if len(parts) < 41:
        return None
    return {
        "target": np.array([float(parts[i]) for i in (26, 27, 28)]),
        "left_eye": np.array([float(parts[i]) for i in (32, 33, 34)]),
        "right_eye": np.array([float(parts[i]) for i in (35, 36, 37)]),
    }


def angle(landmarks, corners, width, height):
    p1 = landmarks[corners[0]] * np.array([width, height])
    p2 = landmarks[corners[1]] * np.array([width, height])
    d = p2 - p1
    return np.degrees(np.arctan2(d[1], d[0]))


def gaze(ann, eye, angle_deg):
    g = ann["target"] - ann[eye]
    g /= np.linalg.norm(g)
    a = np.radians(angle_deg); c, s = np.cos(a), np.sin(a)
    x, y, z = np.array([[c, -s, 0], [s, c, 0], [0, 0, 1]]) @ g
    return np.array([np.arcsin(-y), np.arctan2(-x, -z)])


rows = []
for subject in ("p01", "p08", "p11"):
    candidates = []
    for day in sorted((RAW / subject).glob("day*")):
        ann_path = day / "annotation.txt"
        if not ann_path.exists():
            continue
        for i, line in enumerate(ann_path.read_text().splitlines()):
            ann = parse(line)
            image = day / f"{i+1:04d}.jpg"
            if ann is not None and image.exists():
                candidates.append((day.name, i + 1, image, ann))
    with h5py.File(ROOT / "data" / "processed" / f"{subject}_v16.h5", "r") as f:
        landmarks = f["landmarks"][:].astype(np.float64)
        left_gt = f["left_gaze"][:].astype(np.float64)
        right_gt = f["right_gaze"][:].astype(np.float64)
    first_size = Image.open(candidates[0][2]).size
    width, height = first_size
    left_base = np.stack([(c[3]["target"] - c[3]["left_eye"]) / np.linalg.norm(c[3]["target"] - c[3]["left_eye"]) for c in candidates])
    right_base = np.stack([(c[3]["target"] - c[3]["right_eye"]) / np.linalg.norm(c[3]["target"] - c[3]["right_eye"]) for c in candidates])
    match_lists = []
    match_errors = []
    for row in range(len(left_gt)):
        la = angle(landmarks[row], LEFT_CORNERS, width, height)
        ra = angle(landmarks[row], RIGHT_CORNERS, width, height)
        def rotate_many(base, degrees):
            a = np.radians(degrees); c, s = np.cos(a), np.sin(a)
            x = c * base[:, 0] - s * base[:, 1]
            y = s * base[:, 0] + c * base[:, 1]
            z = base[:, 2]
            return np.column_stack([np.arcsin(-y), np.arctan2(-x, -z)])
        gl_all = rotate_many(left_base, la)
        gr_all = rotate_many(right_base, ra)
        # Training and the legacy evaluator use left_gaze. Right-eye labels are
        # audited after mapping but are not required for identification because
        # historical right-eye rotation conventions changed independently.
        errors = np.max(np.abs(gl_all - left_gt[row]), axis=1)
        matches = np.flatnonzero(errors < 2e-6)
        if not len(matches):
            ci = int(np.argmin(errors)); day, frame, _, _ = candidates[ci]
            raise RuntimeError(f"{subject} row {row}: no exact label match; best={errors[ci]} at {day}/{frame:04d}")
        match_lists.append(matches)
        match_errors.append(errors[matches])

    # Find the globally consistent strictly increasing path. A quadratic gap
    # penalty avoids choosing repeated gaze labels far from the expected cadence.
    expected_gap = len(candidates) / len(left_gt)
    costs, backs = [], []
    for row, matches in enumerate(match_lists):
        if row == 0:
            costs.append((matches.astype(float) - (expected_gap - 1)) ** 2)
            backs.append(np.full(len(matches), -1, dtype=int))
            continue
        prev_m, prev_c = match_lists[row - 1], costs[row - 1]
        cur_cost = np.full(len(matches), np.inf)
        cur_back = np.full(len(matches), -1, dtype=int)
        for j, ci in enumerate(matches):
            valid = np.flatnonzero(prev_m < ci)
            if len(valid):
                gaps = ci - prev_m[valid]
                candidate_costs = prev_c[valid] + (gaps - expected_gap) ** 2
                k = int(np.argmin(candidate_costs))
                cur_cost[j] = candidate_costs[k]
                cur_back[j] = valid[k]
        if not np.isfinite(cur_cost).any():
            raise RuntimeError(f"{subject} row {row}: no monotonic mapping path")
        costs.append(cur_cost); backs.append(cur_back)
    final_matches = match_lists[-1]
    tail = (len(candidates) - 1) - final_matches
    end_scores = costs[-1] + (tail - (expected_gap - 1)) ** 2
    chosen_pos = int(np.argmin(end_scores))
    chosen = [0] * len(match_lists)
    for row in range(len(match_lists) - 1, -1, -1):
        chosen[row] = int(match_lists[row][chosen_pos])
        chosen_pos = int(backs[row][chosen_pos]) if row else -1

    previous = -1
    for row, ci in enumerate(chosen):
        day, frame, image, ann = candidates[ci]
        la = angle(landmarks[row], LEFT_CORNERS, width, height)
        ra = angle(landmarks[row], RIGHT_CORNERS, width, height)
        gl = gaze(ann, "left_eye", la); gr = gaze(ann, "right_eye", ra)
        err = float(np.max(np.abs(gl - left_gt[row])))
        right_err = float(np.max(np.abs(gr - right_gt[row])))
        rows.append({
            "subject": subject, "h5_row": row, "day": day, "frame_number": frame,
            "frame_path": str(image), "candidate_index": ci,
            "skipped_candidates_since_previous": ci - previous - 1,
            "max_label_abs_error_rad": err, "exact_label_match": err < 2e-6,
            "right_label_abs_error_rad": right_err,
            "left_roll_deg": la, "right_roll_deg": ra,
        })
        previous = ci
    print(subject, "mapped", sum(r["exact_label_match"] for r in rows if r["subject"] == subject), "/", len(left_gt), "last candidate", previous + 1, "/", len(candidates))

df = pd.DataFrame(rows)
df.to_csv(OUT / "reconstructed_v16_frame_mapping.csv", index=False)
summary = df.groupby("subject").agg(
    rows=("h5_row", "count"), exact_matches=("exact_label_match", "sum"),
    max_label_error_rad=("max_label_abs_error_rad", "max"),
    total_skipped_candidates=("skipped_candidates_since_previous", "sum"),
    first_frame=("frame_path", "first"), last_frame=("frame_path", "last"),
).reset_index()
summary.to_csv(OUT / "reconstructed_v16_frame_mapping_summary.csv", index=False)
print(summary.to_string(index=False))