Download scripts/generate_kd_clean_cache.py from thanhhuyvan/Gaze-LIPE: direct link, hf CLI and curl.
- Browser
- Download file 17 kB
-
https://huggingface.co/thanhhuyvan/Gaze-LIPE/resolve/main/scripts/generate_kd_clean_cache.py
- Command line
-
hf download hf://thanhhuyvan/Gaze-LIPE/scripts/generate_kd_clean_cache.py
-
curl -L -o generate_kd_clean_cache.py https://huggingface.co/thanhhuyvan/Gaze-LIPE/resolve/main/scripts/generate_kd_clean_cache.py
17 kB
| """Generate a new manifest-backed, strictly loaded, roll-aligned KD cache. | |
| Legacy H5 files are never opened by this script. Existing outputs are never overwritten. | |
| Use bounded ``--max-*`` arguments for a smoke test before a full participant run. | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| from datetime import datetime, timezone | |
| import hashlib | |
| import json | |
| from pathlib import Path | |
| import subprocess | |
| import sys | |
| ROOT = Path(r"E:\Gaze_estimation") | |
| sys.path.insert(0, str(ROOT / ".codex_deps")) | |
| sys.path.insert(0, str(ROOT)) | |
| import cv2 | |
| import h5py | |
| import numpy as np | |
| import torch | |
| from src.models.teacher_strict import audit_to_dict, file_sha256, load_teacher_model_strict | |
| from src.utils.preprocess import GazePreprocessor | |
| DATASET_ROOT = ROOT / "data" / "MPIIGaze" / "MPIIGaze" / "MPIIGaze" | |
| ORIGINAL_ROOT = DATASET_ROOT / "Data" / "Original" | |
| MANIFEST_ROOT = ROOT / "data" / "processed_kd_clean_v1" / "manifests" | |
| OUTPUT_ROOT = ROOT / "data" / "processed_kd_clean_v1" / "cache" | |
| LANDMARK_MODEL = ROOT / "src" / "utils" / "face_landmarker.task" | |
| CHECKPOINT = ROOT / "checkpoints" / "resnet50.pt" | |
| BIN_CENTERS_DEG = np.arange(90, dtype=np.float64) * 4.0 - 180.0 | |
| def text_sha256(value: str) -> str: | |
| return hashlib.sha256(value.encode("utf-8")).hexdigest().upper() | |
| def softmax(values: np.ndarray) -> np.ndarray: | |
| shifted = values.astype(np.float64) - np.max(values) | |
| exp = np.exp(shifted) | |
| return exp / exp.sum() | |
| def gaze_vector(pitch_rad: float, yaw_rad: float) -> np.ndarray: | |
| return np.array( | |
| [ | |
| -np.cos(pitch_rad) * np.sin(yaw_rad), | |
| -np.sin(pitch_rad), | |
| -np.cos(pitch_rad) * np.cos(yaw_rad), | |
| ], | |
| dtype=np.float64, | |
| ) | |
| def angular_error_deg(first: np.ndarray, second: np.ndarray) -> float: | |
| first = first / np.linalg.norm(first) | |
| second = second / np.linalg.norm(second) | |
| return float(np.rad2deg(np.arccos(np.clip(np.dot(first, second), -1.0, 1.0)))) | |
| def rotation_matrix_z(angle_deg: float) -> np.ndarray: | |
| angle = np.deg2rad(angle_deg) | |
| cosine, sine = np.cos(angle), np.sin(angle) | |
| return np.array(((cosine, -sine, 0.0), (sine, cosine, 0.0), (0.0, 0.0, 1.0))) | |
| def angles_from_vectors(vectors: np.ndarray) -> tuple[np.ndarray, np.ndarray]: | |
| pitch = np.arcsin(np.clip(-vectors[:, 1], -1.0, 1.0)) | |
| yaw = np.arctan2(-vectors[:, 0], -vectors[:, 2]) | |
| return np.rad2deg(pitch), np.rad2deg(yaw) | |
| def bin_indices(angle_deg: np.ndarray) -> np.ndarray: | |
| return np.clip(np.rint((angle_deg + 90.0) / 2.0), 0, 89).astype(np.int64) | |
| def make_teacher_grid() -> np.ndarray: | |
| pitch, yaw = np.meshgrid(BIN_CENTERS_DEG, BIN_CENTERS_DEG, indexing="ij") | |
| pitch_rad, yaw_rad = np.deg2rad(pitch.ravel()), np.deg2rad(yaw.ravel()) | |
| return np.column_stack( | |
| ( | |
| -np.cos(pitch_rad) * np.sin(yaw_rad), | |
| -np.sin(pitch_rad), | |
| -np.cos(pitch_rad) * np.cos(yaw_rad), | |
| ) | |
| ) | |
| TEACHER_GRID = make_teacher_grid() | |
| def roll_align_teacher_distribution( | |
| pitch_logits: np.ndarray, yaw_logits: np.ndarray, roll_deg: float | |
| ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: | |
| """Rotate the independent pitch/yaw joint distribution into the eye-normalized frame.""" | |
| pitch_probability = softmax(pitch_logits) | |
| yaw_probability = softmax(yaw_logits) | |
| joint = np.outer(pitch_probability, yaw_probability).ravel() | |
| rotated = TEACHER_GRID @ rotation_matrix_z(roll_deg).T | |
| rotated_pitch, rotated_yaw = angles_from_vectors(rotated) | |
| pitch_marginal = np.bincount(bin_indices(rotated_pitch), weights=joint, minlength=90) | |
| yaw_marginal = np.bincount(bin_indices(rotated_yaw), weights=joint, minlength=90) | |
| pitch_marginal /= pitch_marginal.sum() | |
| yaw_marginal /= yaw_marginal.sum() | |
| aligned_logits = ( | |
| np.log(np.maximum(pitch_marginal, 1e-30)).astype(np.float32), | |
| np.log(np.maximum(yaw_marginal, 1e-30)).astype(np.float32), | |
| ) | |
| expected_vector = (rotated * joint[:, None]).sum(axis=0) | |
| expected_vector /= np.linalg.norm(expected_vector) | |
| return aligned_logits[0], aligned_logits[1], expected_vector | |
| def face_crop(frame: np.ndarray, landmarks, target_size: tuple[int, int] = (448, 448)) -> np.ndarray | None: | |
| height, width = frame.shape[:2] | |
| coordinates = np.array([[point.x * width, point.y * height] for point in landmarks]) | |
| minimum = coordinates.min(axis=0) | |
| maximum = coordinates.max(axis=0) | |
| center = (minimum + maximum) / 2.0 | |
| size = float(np.max(maximum - minimum) * 1.5) | |
| x1, y1 = np.maximum(0, (center - size / 2.0).astype(int)) | |
| x2 = min(width, int(center[0] + size / 2.0)) | |
| y2 = min(height, int(center[1] + size / 2.0)) | |
| crop = frame[y1:y2, x1:x2] | |
| if crop.size == 0: | |
| return None | |
| crop = cv2.resize(crop, target_size) | |
| crop = cv2.cvtColor(crop, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0 | |
| crop = (crop - np.array([0.485, 0.456, 0.406], dtype=np.float32)) / np.array( | |
| [0.229, 0.224, 0.225], dtype=np.float32 | |
| ) | |
| return np.transpose(crop, (2, 0, 1)) | |
| def normalize_eye_with_matrix( | |
| frame: np.ndarray, landmarks, indices: list[int], preprocessor: GazePreprocessor, | |
| target_size: tuple[int, int] = (64, 32) | |
| ) -> tuple[np.ndarray, float, np.ndarray]: | |
| height, width = frame.shape[:2] | |
| first = np.array([landmarks[indices[0]].x * width, landmarks[indices[0]].y * height]) | |
| second = np.array([landmarks[indices[1]].x * width, landmarks[indices[1]].y * height]) | |
| center = (first + second) / 2.0 | |
| delta = second - first | |
| angle = float(np.degrees(np.arctan2(delta[1], delta[0]))) | |
| scale = (target_size[0] * 0.7) / (np.linalg.norm(delta) + 1e-6) | |
| matrix = cv2.getRotationMatrix2D(tuple(center), angle, scale) | |
| matrix[0, 2] += target_size[0] / 2.0 - center[0] | |
| matrix[1, 2] += target_size[1] / 2.0 - center[1] | |
| normalized = cv2.warpAffine(frame, matrix, target_size, flags=cv2.INTER_CUBIC) | |
| normalized = cv2.cvtColor(normalized, cv2.COLOR_BGR2GRAY) | |
| normalized = cv2.medianBlur(normalized, 3) | |
| normalized = preprocessor.clahe.apply(normalized) | |
| return normalized, angle, matrix.astype(np.float32) | |
| def git_commit() -> str: | |
| try: | |
| return subprocess.check_output( | |
| ["git", "rev-parse", "HEAD"], cwd=ROOT, text=True, stderr=subprocess.DEVNULL | |
| ).strip() | |
| except Exception: | |
| return "UNAVAILABLE" | |
| def append_or_reject(records: list[dict], row: dict, accepted: bool, reason: str, output_index: int | None) -> None: | |
| records.append( | |
| { | |
| "source_index": row["source_index"], | |
| "sample_id": row["sample_id"], | |
| "relative_frame_path": row["relative_frame_path"], | |
| "accepted": accepted, | |
| "rejection_reason": reason, | |
| "output_index": output_index, | |
| } | |
| ) | |
| def create_h5(path: Path, accepted: list[dict], attributes: dict) -> None: | |
| string = h5py.string_dtype(encoding="utf-8") | |
| with h5py.File(path, "x") as handle: | |
| for key, value in attributes.items(): | |
| handle.attrs[key] = value if isinstance(value, (str, int, float, bool)) else json.dumps(value, sort_keys=True) | |
| text_fields = ("sample_id", "relative_frame_path", "participant", "day", "frame_id", "raw_image_sha256") | |
| for field in text_fields: | |
| handle.create_dataset(field, data=np.array([row[field] for row in accepted], dtype=object), dtype=string) | |
| handle.create_dataset("source_index", data=np.array([row["source_index"] for row in accepted], dtype=np.int64)) | |
| handle.create_dataset("annotation_row", data=np.array([row["annotation_row"] for row in accepted], dtype=np.int32)) | |
| numeric_fields = ( | |
| "left_patches", "right_patches", "landmarks", "left_gaze", "right_gaze", | |
| "left_affine_matrix", "right_affine_matrix", "left_roll_deg", "right_roll_deg", | |
| "teacher_pitch_logits_raw", "teacher_yaw_logits_raw", | |
| "teacher_pitch_logits", "teacher_yaw_logits", "teacher_aligned_vector", | |
| "teacher_pitch_deg", "teacher_yaw_deg", "teacher_error_deg", | |
| ) | |
| for field in numeric_fields: | |
| handle.create_dataset(field, data=np.asarray([row[field] for row in accepted]), compression="gzip") | |
| def main() -> None: | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--participant", required=True, choices=[f"p{i:02d}" for i in range(15)]) | |
| parser.add_argument("--max-source-rows", type=int) | |
| parser.add_argument("--max-accepted", type=int) | |
| parser.add_argument("--tag", default="full") | |
| parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") | |
| args = parser.parse_args() | |
| manifest_path = MANIFEST_ROOT / f"{args.participant}.source.jsonl" | |
| summary_path = MANIFEST_ROOT / f"{args.participant}.source.summary.json" | |
| validation_path = MANIFEST_ROOT / f"{args.participant}.source.validation.json" | |
| summary = json.loads(summary_path.read_text(encoding="utf-8")) | |
| validation = json.loads(validation_path.read_text(encoding="utf-8")) | |
| if not validation.get("pass") or file_sha256(manifest_path) != summary["manifest_sha256"]: | |
| raise RuntimeError("source manifest is not validated or its hash changed") | |
| OUTPUT_ROOT.mkdir(parents=True, exist_ok=True) | |
| output_path = OUTPUT_ROOT / f"{args.participant}.{args.tag}.h5" | |
| decision_path = OUTPUT_ROOT / f"{args.participant}.{args.tag}.processing.jsonl" | |
| summary_output = OUTPUT_ROOT / f"{args.participant}.{args.tag}.summary.json" | |
| for path in (output_path, decision_path, summary_output): | |
| if path.exists(): | |
| raise FileExistsError(f"refusing to overwrite clean artifact: {path}") | |
| teacher, teacher_audit = load_teacher_model_strict(CHECKPOINT, device=args.device) | |
| preprocessor = GazePreprocessor(model_path=str(LANDMARK_MODEL)) | |
| manifest_rows = manifest_path.read_text(encoding="utf-8").splitlines() | |
| accepted: list[dict] = [] | |
| decisions: list[dict] = [] | |
| examined = 0 | |
| for encoded in manifest_rows: | |
| if args.max_source_rows is not None and examined >= args.max_source_rows: | |
| break | |
| if args.max_accepted is not None and len(accepted) >= args.max_accepted: | |
| break | |
| row = json.loads(encoded) | |
| examined += 1 | |
| image_path = ORIGINAL_ROOT / Path(row["relative_frame_path"]) | |
| if file_sha256(image_path) != row["raw_image_sha256"]: | |
| append_or_reject(decisions, row, False, "raw_image_hash_mismatch", None) | |
| continue | |
| frame = cv2.imread(str(image_path)) | |
| if frame is None: | |
| append_or_reject(decisions, row, False, "opencv_decode_failed", None) | |
| continue | |
| landmarks = preprocessor.get_landmarks(frame) | |
| if landmarks is None: | |
| append_or_reject(decisions, row, False, "face_landmarks_not_found", None) | |
| continue | |
| crop = face_crop(frame, landmarks) | |
| if crop is None: | |
| append_or_reject(decisions, row, False, "teacher_face_crop_empty", None) | |
| continue | |
| left_eye, left_angle, left_matrix = normalize_eye_with_matrix( | |
| frame, landmarks, preprocessor.LEFT_CORNERS, preprocessor | |
| ) | |
| right_eye, right_angle, right_matrix = normalize_eye_with_matrix( | |
| frame, landmarks, preprocessor.RIGHT_CORNERS, preprocessor | |
| ) | |
| left_patches = preprocessor.extract_patches(left_eye, patch_size=16) | |
| right_patches = preprocessor.extract_patches(right_eye, patch_size=16) | |
| landmark_array = np.array([[point.x, point.y] for point in landmarks], dtype=np.float32) | |
| left_center = landmark_array[preprocessor.LEFT_CORNERS].mean(axis=0) | |
| right_center = landmark_array[preprocessor.RIGHT_CORNERS].mean(axis=0) | |
| centered_landmarks = landmark_array - (left_center + right_center) / 2.0 | |
| target = np.asarray(row["target_ccs"], dtype=np.float64) | |
| left_vector = target - np.asarray(row["left_eye_ccs"], dtype=np.float64) | |
| right_vector = target - np.asarray(row["right_eye_ccs"], dtype=np.float64) | |
| left_vector /= np.linalg.norm(left_vector) | |
| right_vector /= np.linalg.norm(right_vector) | |
| left_aligned_vector = rotation_matrix_z(left_angle) @ left_vector | |
| right_aligned_vector = rotation_matrix_z(right_angle) @ right_vector | |
| left_gaze = np.asarray(preprocessor.gaze_3d_to_mag(left_aligned_vector), dtype=np.float32) | |
| right_gaze = np.asarray(preprocessor.gaze_3d_to_mag(right_aligned_vector), dtype=np.float32) | |
| with torch.inference_mode(): | |
| tensor = torch.from_numpy(crop).unsqueeze(0).to(args.device) | |
| raw_pitch, raw_yaw = teacher(tensor) | |
| raw_pitch_np = raw_pitch[0].detach().cpu().numpy().astype(np.float32) | |
| raw_yaw_np = raw_yaw[0].detach().cpu().numpy().astype(np.float32) | |
| aligned_pitch, aligned_yaw, teacher_vector = roll_align_teacher_distribution( | |
| raw_pitch_np, raw_yaw_np, left_angle | |
| ) | |
| teacher_pitch_deg, teacher_yaw_deg = angles_from_vectors(teacher_vector[None, :]) | |
| accepted_row = dict(row) | |
| accepted_row.update( | |
| { | |
| "left_patches": left_patches, | |
| "right_patches": right_patches, | |
| "landmarks": centered_landmarks, | |
| "left_gaze": left_gaze, | |
| "right_gaze": right_gaze, | |
| "left_affine_matrix": left_matrix, | |
| "right_affine_matrix": right_matrix, | |
| "left_roll_deg": np.float32(left_angle), | |
| "right_roll_deg": np.float32(right_angle), | |
| "teacher_pitch_logits_raw": raw_pitch_np, | |
| "teacher_yaw_logits_raw": raw_yaw_np, | |
| "teacher_pitch_logits": aligned_pitch, | |
| "teacher_yaw_logits": aligned_yaw, | |
| "teacher_aligned_vector": teacher_vector.astype(np.float32), | |
| "teacher_pitch_deg": np.float32(teacher_pitch_deg[0]), | |
| "teacher_yaw_deg": np.float32(teacher_yaw_deg[0]), | |
| "teacher_error_deg": np.float32(angular_error_deg(teacher_vector, left_aligned_vector)), | |
| } | |
| ) | |
| output_index = len(accepted) | |
| accepted.append(accepted_row) | |
| append_or_reject(decisions, row, True, "", output_index) | |
| if not accepted: | |
| raise RuntimeError("no rows were accepted; no cache written") | |
| attributes = { | |
| "schema": "mpiigaze-kd-clean-cache-v2-official448", | |
| "created_utc": datetime.now(timezone.utc).isoformat(), | |
| "git_commit": git_commit(), | |
| "participant": args.participant, | |
| "partial": args.max_source_rows is not None or args.max_accepted is not None, | |
| "source_manifest_sha256": summary["manifest_sha256"], | |
| "teacher_loader_audit": audit_to_dict(teacher_audit), | |
| "teacher_checkpoint_sha256": teacher_audit.checkpoint_sha256, | |
| "generator_script_sha256": file_sha256(Path(__file__)), | |
| "teacher_strict_script_sha256": file_sha256(ROOT / "src" / "models" / "teacher_strict.py"), | |
| "preprocess_script_sha256": file_sha256(ROOT / "src" / "utils" / "preprocess.py"), | |
| "landmark_model_sha256": file_sha256(LANDMARK_MODEL), | |
| "patch_size": 16, | |
| "teacher_input_size": 448, | |
| "teacher_bin_centers_deg": "index * 4 - 180", | |
| "teacher_protocol_basis": "recovered checkpoint key layout plus documented Gaze360 ResNet configuration", | |
| "training_target": "left-eye gaze rotated by left eye affine roll angle", | |
| "teacher_alignment": "independent pitch/yaw joint distribution rotated by same left-eye Z roll; marginals rebinned to 90 bins", | |
| } | |
| create_h5(output_path, accepted, attributes) | |
| with decision_path.open("x", encoding="utf-8", newline="\n") as stream: | |
| for decision in decisions: | |
| stream.write(json.dumps(decision, sort_keys=True, separators=(",", ":")) + "\n") | |
| run_summary = { | |
| "schema": "mpiigaze-kd-clean-cache-summary-v1", | |
| "participant": args.participant, | |
| "source_rows_examined": examined, | |
| "accepted_rows": len(accepted), | |
| "rejected_rows": examined - len(accepted), | |
| "cache_path": str(output_path.resolve()), | |
| "cache_sha256": file_sha256(output_path), | |
| "processing_manifest_path": str(decision_path.resolve()), | |
| "processing_manifest_sha256": file_sha256(decision_path), | |
| "source_manifest_sha256": summary["manifest_sha256"], | |
| "teacher_checkpoint_sha256": teacher_audit.checkpoint_sha256, | |
| "strict_loader_inference_sha256": teacher_audit.inference_sha256, | |
| "teacher_error_mean_deg": float(np.mean([row["teacher_error_deg"] for row in accepted])), | |
| "teacher_error_median_deg": float(np.median([row["teacher_error_deg"] for row in accepted])), | |
| } | |
| summary_output.write_text(json.dumps(run_summary, indent=2) + "\n", encoding="utf-8", newline="\n") | |
| print(json.dumps(run_summary, indent=2)) | |
| if __name__ == "__main__": | |
| main() | |