"""Run the Omni anticipation head with EgoSpeak's EasyCom/Ego4D metric.""" from __future__ import annotations import argparse import json from pathlib import Path from typing import Any import numpy as np import torch from tqdm.auto import tqdm from easycom_head import ( CLASS_NAMES, DATASET_NAMES, LOCAL_FUSION_SCOPES, NUM_CLASSES, NUM_HORIZONS, OFFICIAL_TEST_SESSIONS, STEP_SECONDS, TurnPredictionHead, average_precision_metrics, checkpoint_head_config, checkpoint_loss_name, load_ego4d_manifest_split, load_session_data, marginal_probabilities_from_logits, normalize_sample_name, output_classes_for_loss, output_space_for_loss, transition_targets_from_activity, ) from local_turn_features import FramePersonFeatures def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--feature-root", type=Path, required=True) parser.add_argument("--target-root", type=Path, required=True) parser.add_argument("--checkpoint", type=Path, required=True) parser.add_argument("--output-dir", type=Path, required=True) parser.add_argument( "--local-feature-root", type=Path, help="Required for a local-fusion checkpoint; modalities come from the checkpoint.", ) parser.add_argument( "--dataset", choices=DATASET_NAMES, default="easycom", help="Dataset naming/split contract (default: easycom).", ) parser.add_argument( "--manifest", type=Path, help="Required for Ego4D; TSV containing split and uid columns.", ) parser.add_argument( "--sessions", nargs="+", default=None, help=( "Evaluation samples. Defaults to EasyCom Session 3/2/1 or all " "Ego4D manifest test UIDs." ), ) parser.add_argument("--batch-size", type=int, default=1024) parser.add_argument("--device", default="cuda:0") parser.add_argument( "--expected-mode", default="causal-prefix", help="Expected feature mode, or 'any' to disable the check.", ) parser.add_argument("--no-progress", action="store_true") return parser.parse_args() def _resolve_device(name: str) -> torch.device: device = torch.device(name) if device.type == "cuda" and not torch.cuda.is_available(): raise RuntimeError(f"CUDA requested but unavailable: {name}") return device def _load_head( checkpoint: Path, device: torch.device, ) -> tuple[TurnPredictionHead, dict[str, Any]]: payload = torch.load(checkpoint, map_location="cpu", weights_only=False) if not isinstance(payload, dict): raise ValueError(f"Invalid checkpoint: {checkpoint}") state_dict = payload.get("head") if not isinstance(state_dict, dict): raise ValueError(f"Checkpoint has no head state_dict: {checkpoint}") config = checkpoint_head_config(payload) if int(config["num_horizons"]) != NUM_HORIZONS: raise ValueError( f"Checkpoint has {config['num_horizons']} horizons, expected {NUM_HORIZONS}" ) loss_name = checkpoint_loss_name(payload) expected_classes = output_classes_for_loss(loss_name) if int(config["num_classes"]) != expected_classes: raise ValueError( f"Checkpoint has {config['num_classes']} outputs, expected " f"{expected_classes} for loss {loss_name!r}" ) head = TurnPredictionHead(**config) head.load_state_dict(state_dict) return head.to(device).eval(), payload @torch.inference_mode() def _predict_states( head: TurnPredictionHead, states: torch.Tensor, device: torch.device, batch_size: int, *, description: str, show_progress: bool, loss_name: str, local_features: torch.Tensor | None = None, local_person_features: FramePersonFeatures | None = None, ) -> tuple[np.ndarray, np.ndarray | None]: chunks: list[torch.Tensor] = [] transition_chunks: list[torch.Tensor] = [] starts = range(0, len(states), batch_size) progress = tqdm( starts, total=(len(states) + batch_size - 1) // batch_size, desc=description, unit="batch", disable=not show_progress, ) for start in progress: batch = states[start : start + batch_size].to( device=device, dtype=torch.float32, non_blocking=True, ) side = ( local_features[start : start + batch_size].to(device=device, dtype=torch.float32) if local_features is not None else None ) person_batch = None person_mask = None if local_person_features is not None: rows = torch.arange( start, min(start + batch_size, len(states)), ) person_batch, person_mask = local_person_features.padded(rows) person_batch = person_batch.to( device=device, dtype=torch.float32, non_blocking=True, ) person_mask = person_mask.to(device=device, non_blocking=True) if head.transition_auxiliary: activity_logits, transition_logits = head.forward_with_transition( batch, side, person_batch, person_mask, ) transition_chunks.append(transition_logits.sigmoid().cpu()) else: activity_logits = head( batch, side, person_batch, person_mask, ) chunks.append( marginal_probabilities_from_logits(activity_logits, loss_name).cpu() ) transition_probabilities = ( torch.cat(transition_chunks, dim=0).numpy() if transition_chunks else None ) return torch.cat(chunks, dim=0).numpy(), transition_probabilities def _align_to_egospeak_timeline( probabilities: np.ndarray, label_indices: torch.LongTensor, num_target_steps: int, ) -> tuple[np.ndarray, np.ndarray]: """Return EgoSpeak ``[T,3,10]`` scores and their causal-valid mask.""" predictions = np.zeros( (num_target_steps, NUM_CLASSES, NUM_HORIZONS), dtype=np.float32, ) valid = np.zeros((num_target_steps, NUM_HORIZONS), dtype=bool) current = label_indices.numpy() for horizon in range(NUM_HORIZONS): target_indices = current + horizon + 1 usable = (target_indices >= 0) & (target_indices < num_target_steps) rows = target_indices[usable] predictions[rows, :, horizon] = probabilities[usable, horizon, :] valid[rows, horizon] = True return predictions, valid def _align_transitions_to_timeline( probabilities: np.ndarray, label_indices: torch.LongTensor, num_target_steps: int, ) -> np.ndarray: """Align [current_state, future_horizon] transition probabilities.""" predictions = np.zeros( (num_target_steps, NUM_HORIZONS), dtype=np.float32, ) current = label_indices.numpy() for horizon in range(NUM_HORIZONS): target_indices = current + horizon + 1 usable = (target_indices >= 0) & (target_indices < num_target_steps) predictions[target_indices[usable], horizon] = probabilities[usable, horizon] return predictions def _evaluate_transitions( session_predictions: dict[str, np.ndarray], session_valid: dict[str, np.ndarray], session_targets: dict[str, np.ndarray], ) -> dict[str, Any]: """Evaluate auxiliary transition probabilities on causal-valid rows.""" from sklearn.metrics import average_precision_score horizons = [] for horizon in range(NUM_HORIZONS): scores = np.concatenate( [ session_predictions[name][session_valid[name][:, horizon], horizon] for name in session_predictions ] ) targets = np.concatenate( [ session_targets[name][session_valid[name][:, horizon]] for name in session_predictions ] ) predicted = scores >= 0.5 positive = targets == 1 true_positive = int(np.logical_and(predicted, positive).sum()) false_positive = int(np.logical_and(predicted, ~positive).sum()) false_negative = int(np.logical_and(~predicted, positive).sum()) precision = true_positive / max(1, true_positive + false_positive) recall = true_positive / max(1, true_positive + false_negative) f1 = 2.0 * precision * recall / max(1e-12, precision + recall) horizons.append( { "horizon": horizon + 1, "seconds": (horizon + 1) * STEP_SECONDS, "num_steps": int(len(scores)), "positive_rate": float(positive.mean()), "average_precision": float( average_precision_score(targets, scores) ), "precision_at_0.5": precision, "recall_at_0.5": recall, "f1_at_0.5": f1, } ) return { "definition": ( "A positive at target row u means the three-label activity state " "changed between rows u-1 and u." ), "horizons": horizons, "mean_average_precision": float( np.mean([item["average_precision"] for item in horizons]) ), } def _evaluate( session_predictions: dict[str, np.ndarray], session_valid: dict[str, np.ndarray], session_targets: dict[str, np.ndarray], ) -> dict[str, Any]: full_results = [] valid_results = [] for horizon in range(NUM_HORIZONS): full_scores = np.concatenate( [scores[:, :, horizon] for scores in session_predictions.values()], axis=0, ) full_targets = np.concatenate(list(session_targets.values()), axis=0) full = average_precision_metrics(full_scores, full_targets) full["horizon"] = horizon + 1 full["seconds"] = (horizon + 1) * STEP_SECONDS full_results.append(full) valid_scores = np.concatenate( [ session_predictions[name][session_valid[name][:, horizon], :, horizon] for name in session_predictions ], axis=0, ) valid_targets = np.concatenate( [ session_targets[name][session_valid[name][:, horizon]] for name in session_predictions ], axis=0, ) causal = average_precision_metrics(valid_scores, valid_targets) causal["horizon"] = horizon + 1 causal["seconds"] = (horizon + 1) * STEP_SECONDS valid_results.append(causal) return { "egospeak_full_timeline": { "description": ( "Exact EgoSpeak array/evaluator layout. Unavailable pre-session " "causal predictions remain zero in the preallocated arrays." ), "horizons": full_results, "mean_map": float(np.mean([item["map"] for item in full_results])), }, "causal_valid_only": { "description": ( "Primary causal metric; excludes only the first h rows of each " "session for horizon h because no pre-session context exists." ), "horizons": valid_results, "mean_map": float(np.mean([item["map"] for item in valid_results])), }, } def _print_metrics(metrics: dict[str, Any]) -> None: for protocol_name in ("egospeak_full_timeline", "causal_valid_only"): protocol = metrics[protocol_name] print(f"\n{protocol_name}:") for result in protocol["horizons"]: print( f"Action anticipation ({result['seconds']:.2f}s) " f"perframe mAP: {result['map'] * 100:.5f}" ) class_text = ", ".join( f"{name}={value * 100:.5f}" for name, value in result["class_ap"].items() ) print(f" AP: {class_text}; n={result['num_steps']}") print( f"Action anticipation (mean) perframe mAP: " f"{protocol['mean_map'] * 100:.5f}" ) def main() -> None: args = parse_args() if args.batch_size <= 0: raise ValueError("--batch-size must be positive") if args.dataset == "ego4d" and args.manifest is None: raise ValueError("--manifest is required when --dataset ego4d") device = _resolve_device(args.device) expected_mode = None if args.expected_mode == "any" else args.expected_mode if args.dataset == "ego4d": assert args.manifest is not None official_test = load_ego4d_manifest_split(args.manifest, "test") requested_values = args.sessions or official_test else: official_test = list(OFFICIAL_TEST_SESSIONS) requested_values = args.sessions or official_test sessions = [ normalize_sample_name(item, args.dataset) for item in requested_values ] if len(sessions) != len(set(sessions)): raise ValueError("Duplicate evaluation samples were supplied") if args.dataset == "ego4d": unknown = set(sessions) - set(official_test) if unknown: raise ValueError( "Ego4D evaluation samples must belong to the manifest test " f"split; invalid: {sorted(unknown)}" ) head, checkpoint_payload = _load_head(args.checkpoint, device) checkpoint_dataset = checkpoint_payload.get("dataset") if checkpoint_dataset is not None and checkpoint_dataset != args.dataset: raise ValueError( f"Checkpoint dataset is {checkpoint_dataset!r}, but --dataset is " f"{args.dataset!r}" ) head_config = checkpoint_head_config(checkpoint_payload) loss_name = checkpoint_loss_name(checkpoint_payload) local_contract = checkpoint_payload.get("local_contract") if bool(head.local_input_size) != bool(local_contract): raise ValueError("Checkpoint local dimensions/contract disagree") if bool(head.local_input_size) != (args.local_feature_root is not None): raise ValueError( "--local-feature-root must be supplied exactly when the checkpoint uses local fusion" ) args.output_dir.mkdir(parents=True, exist_ok=True) prediction_dir = args.output_dir / "pred" valid_dir = args.output_dir / "valid" prediction_dir.mkdir(parents=True, exist_ok=True) valid_dir.mkdir(parents=True, exist_ok=True) session_predictions: dict[str, np.ndarray] = {} session_valid: dict[str, np.ndarray] = {} session_targets: dict[str, np.ndarray] = {} session_transition_predictions: dict[str, np.ndarray] = {} session_transition_targets: dict[str, np.ndarray] = {} transition_prediction_dir = args.output_dir / "transition_pred" if head.transition_auxiliary: transition_prediction_dir.mkdir(parents=True, exist_ok=True) print(f"device: {device}") print( f"loss: {loss_name}; output_space: {output_space_for_loss(loss_name)}" ) print( f"local fusion: {head.local_fusion}; " f"scope: {LOCAL_FUSION_SCOPES[head.local_fusion]}" ) if head.person_interaction_gates(): print( "person interaction gates: " + json.dumps(head.person_interaction_gates(), sort_keys=True) ) print(f"sessions: {', '.join(sessions)}") for name in sessions: data = load_session_data( args.feature_root, args.target_root, name, expected_mode=expected_mode, dataset=args.dataset, local_feature_root=args.local_feature_root, local_modalities=local_contract["modalities"] if local_contract else "av", local_contract=local_contract, local_fusion=head.local_fusion, ) if data.hidden_size != int(head_config["hidden_size"]): raise ValueError( f"{name} feature dim is {data.hidden_size}, checkpoint expects " f"{head_config['hidden_size']}" ) if data.num_input_layers != int(head_config["num_input_layers"]): raise ValueError( f"{name} has {data.num_input_layers} feature layers, checkpoint " f"expects {head_config['num_input_layers']}" ) probabilities, transition_probabilities = _predict_states( head, data.states, device, args.batch_size, description=f"infer {name}", show_progress=not args.no_progress, loss_name=loss_name, local_features=data.local_features, local_person_features=data.local_person_features, ) predictions, valid = _align_to_egospeak_timeline( probabilities, data.label_indices, len(data.targets), ) target_array = data.targets.numpy() session_predictions[name] = predictions session_valid[name] = valid session_targets[name] = target_array np.save(prediction_dir / f"{name}.npy", predictions) np.save(valid_dir / f"{name}.npy", valid) if transition_probabilities is not None: transition_predictions = _align_transitions_to_timeline( transition_probabilities, data.label_indices, len(data.targets), ) transition_targets = np.zeros(len(data.targets), dtype=np.float32) transition_targets[1:] = ( transition_targets_from_activity(data.targets).numpy() ) session_transition_predictions[name] = transition_predictions session_transition_targets[name] = transition_targets np.save( transition_prediction_dir / f"{name}.npy", transition_predictions, ) print( f"{name}: saved pred {predictions.shape}, " f"causal-valid counts={valid.sum(axis=0).tolist()}" ) metrics = _evaluate(session_predictions, session_valid, session_targets) transition_metrics = ( _evaluate_transitions( session_transition_predictions, session_valid, session_transition_targets, ) if session_transition_predictions else None ) metrics.update( { "sessions": sessions, "dataset": args.dataset, "manifest": ( str(args.manifest.expanduser().resolve()) if args.manifest is not None else None ), "classes": list(CLASS_NAMES), "loss": loss_name, "output_space": output_space_for_loss(loss_name), "prediction_layout": "[target_time, class, future_horizon]", "checkpoint": str(args.checkpoint.resolve()), "checkpoint_epoch": checkpoint_payload.get("epoch"), "checkpoint_train_sessions": checkpoint_payload.get("train_sessions"), "checkpoint_dev_sessions": checkpoint_payload.get("dev_sessions"), "layer_fusion": head.layer_fusion, "local_fusion": head.local_fusion, "local_fusion_scope": LOCAL_FUSION_SCOPES[head.local_fusion], "person_interaction_gates": head.person_interaction_gates(), "local_contract": local_contract, "transition_auxiliary": head.transition_auxiliary, "transition_objective": checkpoint_payload.get( "training_config", {}, ).get("transition_objective", "none"), "transition_metrics": transition_metrics, "layer_mixture_weights_oldest_to_newest": ( head.layer_mixture_weights().detach().float().cpu().tolist() if head.layer_mixture_weights() is not None else None ), } ) metrics_path = args.output_dir / "metrics.json" metrics_path.write_text(json.dumps(metrics, indent=2) + "\n") _print_metrics(metrics) if transition_metrics is not None: print("\ntransition_auxiliary:") for result in transition_metrics["horizons"]: print( f"{result['seconds']:.1f}s: " f"AP={result['average_precision'] * 100:.5f}, " f"P@0.5={result['precision_at_0.5'] * 100:.5f}, " f"R@0.5={result['recall_at_0.5'] * 100:.5f}, " f"F1@0.5={result['f1_at_0.5'] * 100:.5f}" ) print( "10-horizon transition mean AP: " f"{transition_metrics['mean_average_precision'] * 100:.5f}" ) if metrics["layer_mixture_weights_oldest_to_newest"] is not None: print( "\nLayer mixture weights (oldest->newest): " + ", ".join( f"{value:.6f}" for value in metrics[ "layer_mixture_weights_oldest_to_newest" ] ) ) print(f"\nmetrics: {metrics_path}") print(f"EgoSpeak-compatible predictions: {prediction_dir}") if __name__ == "__main__": main()