419 lines
14 KiB
Python
419 lines
14 KiB
Python
"""Shared head, labels, and metrics for the EasyCom/Ego4D experiments."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import csv
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any, Literal, Sequence
|
|
from uuid import UUID
|
|
|
|
import numpy as np
|
|
import torch
|
|
from torch import nn
|
|
|
|
|
|
OFFICIAL_TRAIN_SESSIONS = (
|
|
"Session_7",
|
|
"Session_11",
|
|
"Session_10",
|
|
"Session_6",
|
|
"Session_9",
|
|
"Session_4",
|
|
"Session_8",
|
|
"Session_5",
|
|
"Session_12",
|
|
)
|
|
OFFICIAL_TEST_SESSIONS = ("Session_3", "Session_2", "Session_1")
|
|
CLASS_NAMES = (
|
|
"background",
|
|
"camera_wearer_speaking",
|
|
"other_people_speaking",
|
|
)
|
|
STEP_SECONDS = 0.2
|
|
NUM_HORIZONS = 10
|
|
NUM_CLASSES = 3
|
|
DATASET_NAMES = ("easycom", "ego4d")
|
|
DatasetName = Literal["easycom", "ego4d"]
|
|
|
|
|
|
class TurnPredictionHead(nn.Module):
|
|
"""Map one Omni state to EgoSpeak-compatible ``[10, 3]`` logits."""
|
|
|
|
def __init__(
|
|
self,
|
|
hidden_size: int = 2048,
|
|
num_horizons: int = NUM_HORIZONS,
|
|
num_classes: int = NUM_CLASSES,
|
|
dropout: float = 0.1,
|
|
num_input_layers: int = 1,
|
|
learnable_layer_mixture: bool = False,
|
|
) -> None:
|
|
super().__init__()
|
|
if num_input_layers <= 0:
|
|
raise ValueError("num_input_layers must be positive")
|
|
if learnable_layer_mixture and num_input_layers < 2:
|
|
raise ValueError(
|
|
"learnable_layer_mixture requires at least two input layers"
|
|
)
|
|
if not learnable_layer_mixture and num_input_layers != 1:
|
|
raise ValueError(
|
|
"Multiple input layers require learnable_layer_mixture=True"
|
|
)
|
|
self.hidden_size = hidden_size
|
|
self.num_horizons = num_horizons
|
|
self.num_classes = num_classes
|
|
self.dropout = dropout
|
|
self.num_input_layers = num_input_layers
|
|
self.learnable_layer_mixture = learnable_layer_mixture
|
|
if learnable_layer_mixture:
|
|
# Zero initialization exactly reproduces a uniform layer mean at
|
|
# step zero while allowing training to select useful query layers.
|
|
self.layer_logits = nn.Parameter(torch.zeros(num_input_layers))
|
|
else:
|
|
self.register_parameter("layer_logits", None)
|
|
self.net = nn.Sequential(
|
|
nn.LayerNorm(hidden_size),
|
|
nn.Linear(hidden_size, hidden_size // 2),
|
|
nn.GELU(),
|
|
nn.Dropout(dropout),
|
|
nn.Linear(hidden_size // 2, num_horizons * num_classes),
|
|
)
|
|
|
|
def forward(self, states: torch.Tensor) -> torch.Tensor:
|
|
if states.shape[-1] != self.hidden_size:
|
|
raise ValueError(
|
|
f"Expected hidden size {self.hidden_size}, got {states.shape[-1]}"
|
|
)
|
|
if self.learnable_layer_mixture:
|
|
if states.ndim < 3 or states.shape[-2] != self.num_input_layers:
|
|
raise ValueError(
|
|
"Layer-mixture input must end in "
|
|
f"[{self.num_input_layers},{self.hidden_size}], got "
|
|
f"{tuple(states.shape)}"
|
|
)
|
|
assert self.layer_logits is not None
|
|
weights = self.layer_logits.softmax(dim=0).to(states.dtype)
|
|
view_shape = (1,) * (states.ndim - 2) + (
|
|
self.num_input_layers,
|
|
1,
|
|
)
|
|
states = (states * weights.view(view_shape)).sum(dim=-2)
|
|
logits = self.net(states)
|
|
return logits.view(
|
|
*states.shape[:-1], self.num_horizons, self.num_classes
|
|
)
|
|
|
|
def layer_mixture_weights(self) -> torch.Tensor | None:
|
|
"""Return normalized oldest-to-newest layer weights, if enabled."""
|
|
|
|
if self.layer_logits is None:
|
|
return None
|
|
return self.layer_logits.softmax(dim=0)
|
|
|
|
|
|
def _validate_loss_inputs(logits: torch.Tensor, targets: torch.Tensor) -> None:
|
|
if logits.shape != targets.shape:
|
|
raise ValueError(
|
|
f"logits/targets shape mismatch: {tuple(logits.shape)} vs "
|
|
f"{tuple(targets.shape)}"
|
|
)
|
|
if logits.shape[-1] != NUM_CLASSES:
|
|
raise ValueError(
|
|
f"Expected {NUM_CLASSES} classes, got {logits.shape[-1]}"
|
|
)
|
|
|
|
|
|
def egospeak_multihorizon_loss(
|
|
logits: torch.Tensor,
|
|
targets: torch.Tensor,
|
|
*,
|
|
reduction: Literal["mean", "sum", "none"] = "mean",
|
|
) -> torch.Tensor:
|
|
"""Reproduce the LSTR ``MCE`` loss used by the supplied shell/config.
|
|
|
|
EasyCom targets may contain both active speaker classes as ``[0, 1, 1]``.
|
|
LSTR preserves those weights directly and sums their negative log-softmax
|
|
terms. It does not normalize the target row.
|
|
"""
|
|
|
|
_validate_loss_inputs(logits, targets)
|
|
targets = targets.to(device=logits.device, dtype=logits.dtype)
|
|
loss = -(targets * nn.functional.log_softmax(logits, dim=-1)).sum(dim=-1)
|
|
if reduction == "mean":
|
|
return loss.mean()
|
|
if reduction == "sum":
|
|
return loss.sum()
|
|
if reduction == "none":
|
|
return loss
|
|
raise ValueError(f"Unknown reduction: {reduction}")
|
|
|
|
|
|
def miniroad_multihorizon_loss(
|
|
logits: torch.Tensor,
|
|
targets: torch.Tensor,
|
|
*,
|
|
reduction: Literal["mean", "sum", "none"] = "mean",
|
|
) -> torch.Tensor:
|
|
"""Reproduce MiniROAD's L2-normalized multi-positive anticipation loss."""
|
|
|
|
_validate_loss_inputs(logits, targets)
|
|
targets = nn.functional.normalize(
|
|
targets.to(device=logits.device, dtype=logits.dtype),
|
|
p=2,
|
|
dim=-1,
|
|
)
|
|
loss = -(targets * nn.functional.log_softmax(logits, dim=-1)).sum(dim=-1)
|
|
if reduction == "mean":
|
|
return loss.mean()
|
|
if reduction == "sum":
|
|
return loss.sum()
|
|
if reduction == "none":
|
|
return loss
|
|
raise ValueError(f"Unknown reduction: {reduction}")
|
|
|
|
|
|
def multihorizon_cross_entropy(
|
|
logits: torch.Tensor,
|
|
targets: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
"""Backward-compatible alias for the LSTR/EgoSpeak loss."""
|
|
|
|
return egospeak_multihorizon_loss(logits, targets)
|
|
|
|
|
|
def normalize_session_name(value: str | int) -> str:
|
|
text = str(value)
|
|
if text.startswith("Session_"):
|
|
suffix = text.removeprefix("Session_")
|
|
else:
|
|
suffix = text
|
|
if not suffix.isdigit():
|
|
raise ValueError(f"Invalid EasyCom session: {value}")
|
|
return f"Session_{int(suffix)}"
|
|
|
|
|
|
def normalize_sample_name(
|
|
value: str | int,
|
|
dataset: DatasetName = "easycom",
|
|
) -> str:
|
|
"""Return the canonical EasyCom session name or Ego4D clip UUID."""
|
|
|
|
if dataset == "easycom":
|
|
return normalize_session_name(value)
|
|
if dataset != "ego4d":
|
|
raise ValueError(f"Unknown dataset: {dataset}")
|
|
|
|
text = str(value).strip()
|
|
try:
|
|
return str(UUID(text))
|
|
except (ValueError, AttributeError) as error:
|
|
raise ValueError(f"Invalid Ego4D clip UUID: {value}") from error
|
|
|
|
|
|
def load_ego4d_manifest_split(
|
|
manifest: str | Path,
|
|
split: Literal["train", "test"],
|
|
) -> list[str]:
|
|
"""Load one ordered UID split from ``ego4d_egospeak.tsv``."""
|
|
|
|
path = Path(manifest)
|
|
if not path.is_file():
|
|
raise FileNotFoundError(path)
|
|
|
|
selected: list[str] = []
|
|
seen: set[str] = set()
|
|
with path.open(newline="") as file:
|
|
reader = csv.DictReader(file, delimiter="\t")
|
|
columns = set(reader.fieldnames or [])
|
|
required = {"split", "uid"}
|
|
if not required.issubset(columns):
|
|
raise ValueError(
|
|
f"Manifest must contain {sorted(required)}, got "
|
|
f"{reader.fieldnames}: {path}"
|
|
)
|
|
for row_number, row in enumerate(reader, start=2):
|
|
row_split = row["split"].strip()
|
|
if row_split not in {"train", "test"}:
|
|
raise ValueError(
|
|
f"Invalid split {row_split!r} at {path}:{row_number}"
|
|
)
|
|
uid = normalize_sample_name(row["uid"], "ego4d")
|
|
if uid in seen:
|
|
raise ValueError(
|
|
f"Duplicate Ego4D UID {uid} at {path}:{row_number}"
|
|
)
|
|
seen.add(uid)
|
|
if row_split == split:
|
|
selected.append(uid)
|
|
|
|
if not selected:
|
|
raise ValueError(f"No {split!r} samples found in {path}")
|
|
return selected
|
|
|
|
|
|
@dataclass
|
|
class SessionData:
|
|
name: str
|
|
states: torch.Tensor
|
|
targets: torch.Tensor
|
|
label_indices: torch.LongTensor
|
|
feature_metadata: dict[str, Any]
|
|
|
|
@property
|
|
def hidden_size(self) -> int:
|
|
return int(self.states.shape[-1])
|
|
|
|
@property
|
|
def num_input_layers(self) -> int:
|
|
return 1 if self.states.ndim == 2 else int(self.states.shape[-2])
|
|
|
|
def future_examples(
|
|
self,
|
|
num_horizons: int = NUM_HORIZONS,
|
|
) -> tuple[torch.Tensor, torch.Tensor, torch.LongTensor]:
|
|
"""Return states, future labels, and target indices for valid rows."""
|
|
|
|
offsets = torch.arange(1, num_horizons + 1, dtype=torch.long)
|
|
future_indices = self.label_indices[:, None] + offsets[None, :]
|
|
valid = (future_indices >= 0).all(dim=1) & (
|
|
future_indices < len(self.targets)
|
|
).all(dim=1)
|
|
rows = future_indices[valid]
|
|
return self.states[valid], self.targets[rows], rows
|
|
|
|
|
|
def load_session_data(
|
|
feature_root: str | Path,
|
|
target_root: str | Path,
|
|
session: str | int,
|
|
*,
|
|
expected_mode: str | None = "causal-prefix",
|
|
dataset: DatasetName = "easycom",
|
|
) -> SessionData:
|
|
"""Load and validate one saved Omni feature payload and target array."""
|
|
|
|
name = normalize_sample_name(session, dataset)
|
|
feature_path = Path(feature_root) / f"{name}.pt"
|
|
target_path = Path(target_root) / f"{name}.npy"
|
|
if not feature_path.is_file():
|
|
raise FileNotFoundError(feature_path)
|
|
if not target_path.is_file():
|
|
raise FileNotFoundError(target_path)
|
|
|
|
payload = torch.load(feature_path, map_location="cpu", weights_only=False)
|
|
if not isinstance(payload, dict) or "states" not in payload:
|
|
raise ValueError(f"Invalid feature payload: {feature_path}")
|
|
states = payload["states"]
|
|
if not isinstance(states, torch.Tensor) or states.ndim not in {2, 3}:
|
|
raise ValueError(f"states must be [T,D] or [T,L,D] in {feature_path}")
|
|
if states.ndim == 3 and states.shape[1] < 2:
|
|
raise ValueError(
|
|
f"Stacked states must contain at least two layers: {feature_path}"
|
|
)
|
|
if not torch.is_floating_point(states):
|
|
raise ValueError(f"states must be floating point in {feature_path}")
|
|
if not torch.isfinite(states).all():
|
|
raise ValueError(f"states contain NaN/Inf in {feature_path}")
|
|
|
|
mode = payload.get("mode")
|
|
if expected_mode is not None and mode != expected_mode:
|
|
raise ValueError(
|
|
f"{feature_path} mode is {mode!r}, expected {expected_mode!r}"
|
|
)
|
|
label_indices = payload.get("label_indices")
|
|
if label_indices is None:
|
|
raise ValueError(f"Missing label_indices in {feature_path}")
|
|
label_indices = torch.as_tensor(label_indices, dtype=torch.long)
|
|
if label_indices.ndim != 1 or len(label_indices) != len(states):
|
|
raise ValueError(f"Invalid label_indices in {feature_path}")
|
|
if len(label_indices) and not torch.equal(
|
|
label_indices,
|
|
torch.arange(len(label_indices), dtype=torch.long),
|
|
):
|
|
raise ValueError(
|
|
f"{feature_path} does not use the verified zero-based dense alignment"
|
|
)
|
|
|
|
target_array = np.load(target_path, allow_pickle=False)
|
|
if target_array.ndim != 2 or target_array.shape[1] != NUM_CLASSES:
|
|
raise ValueError(
|
|
f"Expected target [T,{NUM_CLASSES}], got {target_array.shape}: "
|
|
f"{target_path}"
|
|
)
|
|
targets = torch.from_numpy(np.asarray(target_array, dtype=np.float32))
|
|
if len(label_indices) and int(label_indices[-1]) >= len(targets):
|
|
raise ValueError(
|
|
f"Last label index {int(label_indices[-1])} exceeds target length "
|
|
f"{len(targets)} in {name}"
|
|
)
|
|
if len(states) != len(targets):
|
|
raise ValueError(
|
|
f"{name}: feature/target length mismatch: {len(states)} vs "
|
|
f"{len(targets)}"
|
|
)
|
|
|
|
metadata = {
|
|
key: payload.get(key)
|
|
for key in (
|
|
"model_id",
|
|
"mode",
|
|
"config",
|
|
"start",
|
|
"end",
|
|
"num_steps",
|
|
"video_sampling_grid",
|
|
)
|
|
}
|
|
return SessionData(name, states, targets, label_indices, metadata)
|
|
|
|
|
|
def average_precision_metrics(
|
|
scores: np.ndarray,
|
|
targets: np.ndarray,
|
|
) -> dict[str, Any]:
|
|
"""Compute EgoSpeak frame AP for all three classes, including background."""
|
|
|
|
from sklearn.metrics import average_precision_score
|
|
|
|
scores = np.asarray(scores, dtype=np.float64)
|
|
targets = np.asarray(targets, dtype=np.float64)
|
|
if scores.shape != targets.shape or scores.ndim != 2:
|
|
raise ValueError(
|
|
f"scores/targets must have equal [N,C] shape, got "
|
|
f"{scores.shape} and {targets.shape}"
|
|
)
|
|
if scores.shape[1] != NUM_CLASSES:
|
|
raise ValueError(f"Expected {NUM_CLASSES} classes, got {scores.shape[1]}")
|
|
|
|
aps = average_precision_score(targets, scores, average=None)
|
|
return {
|
|
"num_steps": int(len(scores)),
|
|
"class_ap": {
|
|
name: float(value) for name, value in zip(CLASS_NAMES, aps)
|
|
},
|
|
"map": float(np.mean(aps)),
|
|
}
|
|
|
|
|
|
def checkpoint_head_config(payload: dict[str, Any]) -> dict[str, Any]:
|
|
"""Return a validated head config, including legacy-checkpoint defaults."""
|
|
|
|
config = dict(payload.get("head_config", {}))
|
|
config.setdefault("hidden_size", 2048)
|
|
config.setdefault("num_horizons", NUM_HORIZONS)
|
|
config.setdefault("num_classes", NUM_CLASSES)
|
|
config.setdefault("dropout", 0.1)
|
|
config.setdefault("num_input_layers", 1)
|
|
config.setdefault("learnable_layer_mixture", False)
|
|
return config
|
|
|
|
|
|
def validate_disjoint_sessions(
|
|
first: Sequence[str],
|
|
second: Sequence[str],
|
|
) -> None:
|
|
overlap = set(first) & set(second)
|
|
if overlap:
|
|
raise ValueError(f"Session sets overlap: {sorted(overlap)}")
|