Add reproducible EasyCom causal head pipeline
This commit is contained in:
Xin Wang
289
infer_turn_head.py
Normal file
289
infer_turn_head.py
Normal file
@@ -0,0 +1,289 @@
|
||||
"""Run the Omni anticipation head and evaluate it with EgoSpeak's EasyCom 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,
|
||||
NUM_CLASSES,
|
||||
NUM_HORIZONS,
|
||||
OFFICIAL_TEST_SESSIONS,
|
||||
STEP_SECONDS,
|
||||
TurnPredictionHead,
|
||||
average_precision_metrics,
|
||||
checkpoint_head_config,
|
||||
load_session_data,
|
||||
normalize_session_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(
|
||||
"--sessions",
|
||||
nargs="+",
|
||||
default=list(OFFICIAL_TEST_SESSIONS),
|
||||
help="Evaluation sessions; default is the official Session 3/2/1 test split.",
|
||||
)
|
||||
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")
|
||||
device = _resolve_device(args.device)
|
||||
expected_mode = None if args.expected_mode == "any" else args.expected_mode
|
||||
sessions = [normalize_session_name(item) for item in args.sessions]
|
||||
if len(sessions) != len(set(sessions)):
|
||||
raise ValueError("Duplicate sessions were supplied")
|
||||
head, checkpoint_payload = _load_head(args.checkpoint, device)
|
||||
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,
|
||||
)
|
||||
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']}"
|
||||
)
|
||||
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,
|
||||
"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"),
|
||||
}
|
||||
)
|
||||
metrics_path = args.output_dir / "metrics.json"
|
||||
metrics_path.write_text(json.dumps(metrics, indent=2) + "\n")
|
||||
_print_metrics(metrics)
|
||||
print(f"\nmetrics: {metrics_path}")
|
||||
print(f"EgoSpeak-compatible predictions: {prediction_dir}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user