| """Batched equivalent of generate_kd_clean_cache.py for CPU throughput. |
| |
| Landmark detection and every preprocessing operation remain row-serial and use the |
| same functions as the audited generator. Only strict-teacher forward calls are |
| grouped into batches; output rows retain manifest order. Existing files are never |
| overwritten. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| from datetime import datetime, timezone |
| import json |
| from pathlib import Path |
| 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 numpy as np |
| import torch |
|
|
| import scripts.generate_kd_clean_cache as base |
| from src.models.teacher_strict import audit_to_dict, load_teacher_model_strict |
| from src.utils.preprocess import GazePreprocessor |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--participant", required=True, choices=[f"p{i:02d}" for i in range(15)]) |
| parser.add_argument("--tag", default="official448_full") |
| parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") |
| parser.add_argument("--teacher-batch-size", type=int, default=16) |
| parser.add_argument("--max-source-rows", type=int) |
| parser.add_argument("--max-accepted", type=int) |
| args = parser.parse_args() |
| if args.teacher_batch_size < 1: |
| raise ValueError("teacher batch size must be positive") |
|
|
| manifest_path = base.MANIFEST_ROOT / f"{args.participant}.source.jsonl" |
| summary_path = base.MANIFEST_ROOT / f"{args.participant}.source.summary.json" |
| validation_path = base.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 base.file_sha256(manifest_path) != summary["manifest_sha256"]: |
| raise RuntimeError("source manifest is not validated or its hash changed") |
|
|
| base.OUTPUT_ROOT.mkdir(parents=True, exist_ok=True) |
| output_path = base.OUTPUT_ROOT / f"{args.participant}.{args.tag}.h5" |
| decision_path = base.OUTPUT_ROOT / f"{args.participant}.{args.tag}.processing.jsonl" |
| summary_output = base.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(base.CHECKPOINT, device=args.device) |
| preprocessor = GazePreprocessor(model_path=str(base.LANDMARK_MODEL)) |
| accepted, decisions, pending = [], [], [] |
|
|
| def flush_teacher(): |
| if not pending: |
| return |
| tensor = torch.from_numpy(np.stack([item[1] for item in pending])).to(args.device) |
| with torch.inference_mode(): |
| raw_pitch, raw_yaw = teacher(tensor) |
| raw_pitch = raw_pitch.detach().cpu().numpy().astype(np.float32) |
| raw_yaw = raw_yaw.detach().cpu().numpy().astype(np.float32) |
| for batch_index, (output_index, _) in enumerate(pending): |
| record = accepted[output_index] |
| pitch_logits, yaw_logits = raw_pitch[batch_index], raw_yaw[batch_index] |
| aligned_pitch, aligned_yaw, teacher_vector = base.roll_align_teacher_distribution( |
| pitch_logits, yaw_logits, float(record["left_roll_deg"]) |
| ) |
| teacher_pitch, teacher_yaw = base.angles_from_vectors(teacher_vector[None, :]) |
| record.update({ |
| "teacher_pitch_logits_raw": pitch_logits, |
| "teacher_yaw_logits_raw": yaw_logits, |
| "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[0]), |
| "teacher_yaw_deg": np.float32(teacher_yaw[0]), |
| "teacher_error_deg": np.float32(base.angular_error_deg(teacher_vector, record.pop("_left_aligned_vector"))), |
| }) |
| pending.clear() |
|
|
| examined = 0 |
| for encoded in manifest_path.read_text(encoding="utf-8").splitlines(): |
| 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 = base.ORIGINAL_ROOT / Path(row["relative_frame_path"]) |
| if base.file_sha256(image_path) != row["raw_image_sha256"]: |
| base.append_or_reject(decisions, row, False, "raw_image_hash_mismatch", None); continue |
| frame = cv2.imread(str(image_path)) |
| if frame is None: |
| base.append_or_reject(decisions, row, False, "opencv_decode_failed", None); continue |
| landmarks = preprocessor.get_landmarks(frame) |
| if landmarks is None: |
| base.append_or_reject(decisions, row, False, "face_landmarks_not_found", None); continue |
| crop = base.face_crop(frame, landmarks) |
| if crop is None: |
| base.append_or_reject(decisions, row, False, "teacher_face_crop_empty", None); continue |
| left_eye, left_angle, left_matrix = base.normalize_eye_with_matrix(frame, landmarks, preprocessor.LEFT_CORNERS, preprocessor) |
| right_eye, right_angle, right_matrix = base.normalize_eye_with_matrix(frame, landmarks, preprocessor.RIGHT_CORNERS, preprocessor) |
| landmark_array = np.asarray([[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) |
| 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 = base.rotation_matrix_z(left_angle) @ left_vector |
| right_aligned = base.rotation_matrix_z(right_angle) @ right_vector |
| record = dict(row) |
| record.update({ |
| "left_patches": preprocessor.extract_patches(left_eye, patch_size=16), |
| "right_patches": preprocessor.extract_patches(right_eye, patch_size=16), |
| "landmarks": landmark_array - (left_center + right_center) / 2.0, |
| "left_gaze": np.asarray(preprocessor.gaze_3d_to_mag(left_aligned), dtype=np.float32), |
| "right_gaze": np.asarray(preprocessor.gaze_3d_to_mag(right_aligned), dtype=np.float32), |
| "left_affine_matrix": left_matrix, "right_affine_matrix": right_matrix, |
| "left_roll_deg": np.float32(left_angle), "right_roll_deg": np.float32(right_angle), |
| "_left_aligned_vector": left_aligned, |
| }) |
| output_index = len(accepted); accepted.append(record) |
| pending.append((output_index, crop)) |
| base.append_or_reject(decisions, row, True, "", output_index) |
| if len(pending) >= args.teacher_batch_size: |
| flush_teacher() |
| flush_teacher() |
| 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": base.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": base.file_sha256(Path(__file__)), |
| "serial_reference_script_sha256": base.file_sha256(Path(base.__file__)), |
| "teacher_strict_script_sha256": base.file_sha256(ROOT / "src" / "models" / "teacher_strict.py"), |
| "preprocess_script_sha256": base.file_sha256(ROOT / "src" / "utils" / "preprocess.py"), |
| "landmark_model_sha256": base.file_sha256(base.LANDMARK_MODEL), |
| "patch_size": 16, "teacher_input_size": 448, "teacher_batch_size": args.teacher_batch_size, |
| "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", |
| } |
| base.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") |
| errors = np.asarray([row["teacher_error_deg"] for row in accepted]) |
| 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": base.file_sha256(output_path), |
| "processing_manifest_path": str(decision_path.resolve()), "processing_manifest_sha256": base.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(errors.mean()), "teacher_error_median_deg": float(np.median(errors)), |
| "teacher_batch_size": args.teacher_batch_size, |
| } |
| 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() |
|
|
