Add inference-only evaluation for EgoSpeak Omni finetuning

2026-09-16 17:28:31 +08:00
parent 64adbef13d
commit 7344d0c469
3 changed files with 181 additions and 0 deletions
@@ -0,0 +1,128 @@
+"""Evaluate an existing stage-two checkpoint without any optimization."""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import numpy as np
+
+from easycom_head import OFFICIAL_TEST_SESSIONS
+from egospeak_omni_finetune import load_stage_two_checkpoint
+from egospeak_omni_residual import align_source_predictions, sha256_file
+from train_egospeak_omni_finetune import raw_endpoint_metrics
+from train_egospeak_omni_residual import (
+    evaluate_predictions,
+    load_sessions,
+    normalize_sessions,
+    omni_contract_signature,
+    predict_session,
+    print_metrics,
+    resolve_device,
+)
+
+
+def main(argv: list[str] | None = None) -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--checkpoint", type=Path, required=True)
+    parser.add_argument("--output-dir", type=Path, required=True)
+    parser.add_argument("--sessions", nargs="+", default=None)
+    parser.add_argument("--egospeak-feature-root", type=Path)
+    parser.add_argument("--omni-feature-root", type=Path)
+    parser.add_argument("--device", default="cuda")
+    parser.add_argument("--eval-chunk-size", type=int, default=0)
+    parser.add_argument("--no-progress", action="store_true")
+    args = parser.parse_args(argv)
+    if args.eval_chunk_size < 0:
+        raise ValueError("Evaluation chunk size must be nonnegative")
+    model, payload = load_stage_two_checkpoint(args.checkpoint)
+    config = payload["training_config"]
+    args.egospeak_feature_root = args.egospeak_feature_root or Path(
+        config["egospeak_feature_root"]
+    )
+    args.omni_feature_root = args.omni_feature_root or Path(config["omni_feature_root"])
+    args.omni_layer_fusion = config["omni_layer_fusion"]
+    args.expected_omni_mode = config["expected_omni_mode"]
+    names = normalize_sessions(args.sessions, OFFICIAL_TEST_SESSIONS)
+    sessions = load_sessions(args, names)
+    if {int(s.omni.shape[-1]) for s in sessions} != {model.omni_dim}:
+        raise ValueError("Inference Omni dimension differs from checkpoint")
+    known = payload.get("omni_sources", {})
+    for session in sessions:
+        signature = omni_contract_signature(session)
+        for source in known.values():
+            c = source["config"]
+            expected = (
+                source["model_id"],
+                source["mode"],
+                c.get("memory_seconds"),
+                c.get("video_fps"),
+                c.get("readout_last_n_layers"),
+                c.get("readout_layer_fusion"),
+                tuple(source["state_shape_before_layer_fusion"][1:]),
+            )
+            if signature != expected:
+                raise ValueError(
+                    "Inference Omni extraction differs from checkpoint provenance"
+                )
+    output = args.output_dir.expanduser().resolve()
+    if output.exists() and (not output.is_dir() or any(output.iterdir())):
+        raise FileExistsError(f"Refusing to overwrite output: {output}")
+    for name in ("pred", "valid", "raw_eval"):
+        (output / name).mkdir(parents=True)
+    device = resolve_device(args.device)
+    model.requires_grad_(False)
+    model.to(device).eval()
+    predictions, valid, targets = {}, {}, {}
+    for session in sessions:
+        raw, _ = predict_session(
+            model,
+            session,
+            device,
+            chunk_size=args.eval_chunk_size,
+            amp=False,
+            show_progress=not args.no_progress,
+        )
+        aligned, mask = align_source_predictions(raw, session.num_steps)
+        predictions[session.name], valid[session.name] = aligned, mask
+        targets[session.name] = np.asarray(session.targets, dtype=np.float32)
+        np.save(output / "raw_eval" / f"{session.name}.npy", raw)
+        np.save(output / "pred" / f"{session.name}.npy", aligned)
+        np.save(output / "valid" / f"{session.name}.npy", mask)
+    dense = evaluate_predictions(predictions, valid, targets)
+    endpoints = raw_endpoint_metrics(predictions, valid, targets)
+    metrics = {
+        **dense,
+        "raw_endpoint_anticipation": endpoints,
+        "sessions": names,
+        "train_sessions": payload["train_sessions"],
+        "seed": config["seed"],
+        "checkpoint": str(args.checkpoint.resolve()),
+        "checkpoint_sha256": sha256_file(args.checkpoint),
+        "inference_only": True,
+        "budgets": payload["budgets"],
+        "omni_sources": {s.name: s.omni_metadata for s in sessions},
+    }
+    (output / "metrics.json").write_text(json.dumps(metrics, indent=2) + "\n")
+    (output / "README.md").write_text(
+        "# Stage-two inference-only evaluation\n\n"
+        f"Checkpoint: {args.checkpoint}. Sessions: {names}. No optimization.\n\n"
+        "metrics.json stores dense/raw endpoint AP as fractions; console uses percent. "
+        "Repeated official-test evaluation is exploratory, not dev selection.\n"
+    )
+    print_metrics("stage_two inference only", dense)
+    print(
+        "Raw endpoint mean AP (percent):",
+        {
+            family: None
+            if endpoints[f"mean_{family}_ap"] is None
+            else endpoints[f"mean_{family}_ap"] * 100
+            for family in ("onset", "offset")
+        },
+    )
+    print(f"Saved {output / 'metrics.json'}; checkpoint unchanged")
+
+
+if __name__ == "__main__":
+    main()