129 lines
5.2 KiB
Python
129 lines
5.2 KiB
Python
"""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()
|