300 lines
9.4 KiB
Python
300 lines
9.4 KiB
Python
"""Shared head, labels, and metrics for the EasyCom EgoSpeak experiment."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any, Literal, Sequence
|
|
|
|
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
|
|
|
|
|
|
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,
|
|
) -> None:
|
|
super().__init__()
|
|
self.hidden_size = hidden_size
|
|
self.num_horizons = num_horizons
|
|
self.num_classes = num_classes
|
|
self.dropout = dropout
|
|
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:
|
|
logits = self.net(states)
|
|
return logits.view(
|
|
*states.shape[:-1], self.num_horizons, self.num_classes
|
|
)
|
|
|
|
|
|
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)}"
|
|
|
|
|
|
@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])
|
|
|
|
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",
|
|
) -> SessionData:
|
|
"""Load and validate one saved Omni feature payload and target array."""
|
|
|
|
name = normalize_session_name(session)
|
|
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 != 2:
|
|
raise ValueError(f"states must be [T,D] in {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)
|
|
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)}")
|