icassp2027-experiment/infer_turn_head.py

"""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,
    HISTORY_VISUAL_SET_FUSION_MODES,
    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)),
            )
            if head.local_fusion in HISTORY_VISUAL_SET_FUSION_MODES:
                person_batch, person_mask = local_person_features.padded_history(
                    rows,
                    max_observations=head.person_history_frames,
                )
            else:
                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,
            person_track_max_gap_seconds=head.person_track_max_gap_seconds,
        )
        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()}"
        )
        if head.person_adapter_diagnostics():
            print(
                "person adapter diagnostics: "
                + json.dumps(
                    head.person_adapter_diagnostics(),
                    sort_keys=True,
                )
            )

    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_training_epochs": checkpoint_payload.get(
                "training_config",
                {},
            ).get("epochs"),
            "checkpoint_train_sessions": checkpoint_payload.get("train_sessions"),
            "checkpoint_dev_sessions": checkpoint_payload.get("dev_sessions"),
            "initialized_pooled_checkpoint": checkpoint_payload.get(
                "training_config", {}
            ).get("initialize_pooled_checkpoint"),
            "initialized_pooled_checkpoint_sha256": checkpoint_payload.get(
                "training_config", {}
            ).get("initialize_pooled_checkpoint_sha256"),
            "initialized_pooled_checkpoint_epoch": checkpoint_payload.get(
                "training_config", {}
            ).get("initialize_pooled_checkpoint_epoch"),
            "pooled_base_frozen": checkpoint_payload.get(
                "training_config", {}
            ).get("pooled_base_frozen"),
            "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(),
            "person_adapter_diagnostics": head.person_adapter_diagnostics(),
            "person_history_frames": head.person_history_frames,
            "person_track_max_gap_seconds": head.person_track_max_gap_seconds,
            "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()