357 lines
12 KiB
Python
357 lines
12 KiB
Python
"""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,
|
|
NUM_CLASSES,
|
|
NUM_HORIZONS,
|
|
OFFICIAL_TEST_SESSIONS,
|
|
STEP_SECONDS,
|
|
TurnPredictionHead,
|
|
average_precision_metrics,
|
|
checkpoint_head_config,
|
|
load_ego4d_manifest_split,
|
|
load_session_data,
|
|
normalize_sample_name,
|
|
)
|
|
|
|
|
|
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(
|
|
"--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}"
|
|
)
|
|
if int(config["num_classes"]) != NUM_CLASSES:
|
|
raise ValueError(
|
|
f"Checkpoint has {config['num_classes']} classes, expected {NUM_CLASSES}"
|
|
)
|
|
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,
|
|
) -> np.ndarray:
|
|
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,
|
|
)
|
|
chunks.append(head(batch).softmax(dim=-1).cpu())
|
|
return torch.cat(chunks, dim=0).numpy()
|
|
|
|
|
|
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 _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)
|
|
|
|
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] = {}
|
|
|
|
print(f"device: {device}")
|
|
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,
|
|
)
|
|
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 = _predict_states(
|
|
head,
|
|
data.states,
|
|
device,
|
|
args.batch_size,
|
|
description=f"infer {name}",
|
|
show_progress=not args.no_progress,
|
|
)
|
|
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)
|
|
print(
|
|
f"{name}: saved pred {predictions.shape}, "
|
|
f"causal-valid counts={valid.sum(axis=0).tolist()}"
|
|
)
|
|
|
|
metrics = _evaluate(session_predictions, session_valid, session_targets)
|
|
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),
|
|
"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_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 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()
|