672 lines
26 KiB
Python
672 lines
26 KiB
Python
|
"""Shared data/model code for the MuVAP-only EgoSpeak Stage-3 baseline."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import math
|
||
|
|
from dataclasses import dataclass
|
||
|
|
from pathlib import Path
|
||
|
|
from typing import Any, Iterable, Sequence
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
import torch
|
||
|
|
from torch import nn
|
||
|
|
from torch.utils.data import Dataset
|
||
|
|
|
||
|
|
from easycom_head import (
|
||
|
|
CLASS_NAMES,
|
||
|
|
NUM_CLASSES,
|
||
|
|
NUM_HORIZONS,
|
||
|
|
STEP_SECONDS,
|
||
|
|
_validated_binary_activity_targets,
|
||
|
|
average_precision_metrics,
|
||
|
|
marginal_probabilities_from_logits,
|
||
|
|
normalize_sample_name,
|
||
|
|
)
|
||
|
|
MODEL_FRAME_HZ = 25.0
|
||
|
|
MODEL_FRAMES_PER_TASK_STEP = 5
|
||
|
|
MUVAP_EGOSPEAK_INPUT_FORMAT = "muvap-egospeak-stage3-paper-reconstruction-v1"
|
||
|
|
MODEL_KIND = "muvap-egospeak-stage3-paper-reconstruction"
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass(frozen=True)
|
||
|
|
class MuVAPSessionFeatures:
|
||
|
|
name: str
|
||
|
|
vap: torch.Tensor
|
||
|
|
people: torch.Tensor
|
||
|
|
person_valid: torch.BoolTensor
|
||
|
|
targets: torch.Tensor
|
||
|
|
task_frame_indices: torch.LongTensor
|
||
|
|
label_indices: torch.LongTensor
|
||
|
|
times: torch.Tensor
|
||
|
|
contract: dict[str, Any]
|
||
|
|
|
||
|
|
@property
|
||
|
|
def vap_dim(self) -> int:
|
||
|
|
return int(self.vap.shape[-1])
|
||
|
|
|
||
|
|
@property
|
||
|
|
def asd_dim(self) -> int:
|
||
|
|
return int(self.people.shape[-1])
|
||
|
|
|
||
|
|
|
||
|
|
def _validate_exact_grid(values: torch.Tensor, expected: torch.Tensor, name: str) -> None:
|
||
|
|
if values.dtype != torch.float64 or values.shape != expected.shape:
|
||
|
|
raise ValueError(f"Invalid {name} shape/dtype")
|
||
|
|
if not torch.allclose(values, expected, atol=1e-8, rtol=0):
|
||
|
|
raise ValueError(f"{name} does not follow the required right-edge grid")
|
||
|
|
|
||
|
|
|
||
|
|
def load_muvap_session(
|
||
|
|
feature_root: str | Path,
|
||
|
|
target_root: str | Path,
|
||
|
|
session: str | int,
|
||
|
|
*,
|
||
|
|
expected_contract: dict[str, Any] | None = None,
|
||
|
|
allow_partial: bool = False,
|
||
|
|
) -> MuVAPSessionFeatures:
|
||
|
|
"""Load one 25 Hz frozen-backbone sidecar and fail closed on alignment."""
|
||
|
|
|
||
|
|
name = normalize_sample_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 payload.get("format") != MUVAP_EGOSPEAK_INPUT_FORMAT:
|
||
|
|
raise ValueError(f"Not a MuVAP EgoSpeak Stage-3 sidecar: {feature_path}")
|
||
|
|
if payload.get("dataset") != "easycom" or payload.get("sample_id") != name:
|
||
|
|
raise ValueError(f"Dataset/sample identity mismatch: {feature_path}")
|
||
|
|
if bool(payload.get("partial", False)) and not allow_partial:
|
||
|
|
raise ValueError(f"Partial sidecar cannot be used for training: {feature_path}")
|
||
|
|
contract = {
|
||
|
|
"format": payload["format"],
|
||
|
|
"dataset": payload["dataset"],
|
||
|
|
"config": payload.get("config"),
|
||
|
|
}
|
||
|
|
if not isinstance(contract["config"], dict):
|
||
|
|
raise ValueError(f"Missing MuVAP extraction config: {feature_path}")
|
||
|
|
if contract["config"].get("uses_omni_features") is not False:
|
||
|
|
raise ValueError(f"MuVAP-only sidecar does not exclude Omni: {feature_path}")
|
||
|
|
if contract["config"].get("uses_speech_labels") is not False:
|
||
|
|
raise ValueError(f"Feature extractor consumed speech labels: {feature_path}")
|
||
|
|
if expected_contract is not None and contract != expected_contract:
|
||
|
|
raise ValueError(f"MuVAP feature contract differs across sessions: {feature_path}")
|
||
|
|
|
||
|
|
vap = payload.get("vap")
|
||
|
|
if (
|
||
|
|
not isinstance(vap, torch.Tensor)
|
||
|
|
or vap.ndim != 2
|
||
|
|
or not vap.is_floating_point()
|
||
|
|
or not torch.isfinite(vap).all()
|
||
|
|
):
|
||
|
|
raise ValueError(f"vap must be finite floating [K,D]: {feature_path}")
|
||
|
|
num_ticks = len(vap)
|
||
|
|
frame_times = payload.get("frame_times")
|
||
|
|
if not isinstance(frame_times, torch.Tensor):
|
||
|
|
raise ValueError(f"Missing frame_times: {feature_path}")
|
||
|
|
_validate_exact_grid(
|
||
|
|
frame_times,
|
||
|
|
torch.arange(1, num_ticks + 1, dtype=torch.float64) / MODEL_FRAME_HZ,
|
||
|
|
"MuVAP frame_times",
|
||
|
|
)
|
||
|
|
|
||
|
|
times = payload.get("times")
|
||
|
|
label_indices = payload.get("label_indices")
|
||
|
|
task_frame_indices = payload.get("task_frame_indices")
|
||
|
|
if not all(isinstance(value, torch.Tensor) for value in (times, label_indices, task_frame_indices)):
|
||
|
|
raise ValueError(f"Missing task alignment tensors: {feature_path}")
|
||
|
|
num_steps = len(times)
|
||
|
|
_validate_exact_grid(
|
||
|
|
times,
|
||
|
|
torch.arange(1, num_steps + 1, dtype=torch.float64) * STEP_SECONDS,
|
||
|
|
"task times",
|
||
|
|
)
|
||
|
|
expected_labels = torch.arange(num_steps, dtype=torch.long)
|
||
|
|
expected_frames = (expected_labels + 1) * MODEL_FRAMES_PER_TASK_STEP - 1
|
||
|
|
if label_indices.dtype != torch.long or not torch.equal(label_indices, expected_labels):
|
||
|
|
raise ValueError(f"label_indices must be dense and zero-based: {feature_path}")
|
||
|
|
if task_frame_indices.dtype != torch.long or not torch.equal(
|
||
|
|
task_frame_indices, expected_frames
|
||
|
|
):
|
||
|
|
raise ValueError(f"Invalid task-to-25Hz mapping: {feature_path}")
|
||
|
|
if num_ticks != num_steps * MODEL_FRAMES_PER_TASK_STEP:
|
||
|
|
raise ValueError(f"25Hz/5Hz length mismatch: {feature_path}")
|
||
|
|
|
||
|
|
person = payload.get("persons")
|
||
|
|
if not isinstance(person, dict):
|
||
|
|
raise ValueError(f"Missing persons block: {feature_path}")
|
||
|
|
features = person.get("features")
|
||
|
|
offsets = person.get("frame_offsets")
|
||
|
|
records = person.get("records")
|
||
|
|
if (
|
||
|
|
not isinstance(features, torch.Tensor)
|
||
|
|
or features.ndim != 2
|
||
|
|
or not features.is_floating_point()
|
||
|
|
or not torch.isfinite(features).all()
|
||
|
|
):
|
||
|
|
raise ValueError(f"Invalid fused ASD features: {feature_path}")
|
||
|
|
if (
|
||
|
|
not isinstance(offsets, torch.Tensor)
|
||
|
|
or offsets.dtype != torch.long
|
||
|
|
or offsets.shape != (num_ticks + 1,)
|
||
|
|
or int(offsets[0]) != 0
|
||
|
|
or int(offsets[-1]) != len(features)
|
||
|
|
or bool((offsets[1:] < offsets[:-1]).any())
|
||
|
|
):
|
||
|
|
raise ValueError(f"Invalid person frame offsets: {feature_path}")
|
||
|
|
if not isinstance(records, list) or len(records) != len(features):
|
||
|
|
raise ValueError(f"Person records/features differ: {feature_path}")
|
||
|
|
config = contract["config"]
|
||
|
|
if config.get("vap_dim") != vap.shape[1] or config.get("asd_dim") != features.shape[1]:
|
||
|
|
raise ValueError(f"Saved feature dimensions disagree with config: {feature_path}")
|
||
|
|
if not config.get("vap_release_sha256") or not config.get("asd_release_sha256"):
|
||
|
|
raise ValueError(f"Missing frozen release identity: {feature_path}")
|
||
|
|
counts = offsets[1:] - offsets[:-1]
|
||
|
|
max_people = int(counts.max()) if len(counts) else 0
|
||
|
|
saved_max = payload.get("max_people")
|
||
|
|
if isinstance(saved_max, bool) or not isinstance(saved_max, int) or saved_max != max_people:
|
||
|
|
raise ValueError(f"Incorrect max_people metadata: {feature_path}")
|
||
|
|
people = torch.zeros(num_ticks, max_people, features.shape[1], dtype=features.dtype)
|
||
|
|
person_valid = torch.zeros(num_ticks, max_people, dtype=torch.bool)
|
||
|
|
for tick in range(num_ticks):
|
||
|
|
start, end = int(offsets[tick]), int(offsets[tick + 1])
|
||
|
|
count = end - start
|
||
|
|
if count:
|
||
|
|
seen: set[str] = set()
|
||
|
|
expected_edge = (tick + 1) / MODEL_FRAME_HZ
|
||
|
|
for record in records[start:end]:
|
||
|
|
if not isinstance(record, dict):
|
||
|
|
raise ValueError(f"Invalid person record at tick {tick}: {feature_path}")
|
||
|
|
person_id = str(record.get("id", ""))
|
||
|
|
if person_id in seen:
|
||
|
|
raise ValueError(f"Duplicate person ID at tick {tick}: {feature_path}")
|
||
|
|
seen.add(person_id)
|
||
|
|
if record.get("asd_model_tick") != tick:
|
||
|
|
raise ValueError(f"Person record assigned to wrong tick: {feature_path}")
|
||
|
|
edge = record.get("asd_interval_end")
|
||
|
|
source_time = record.get("source_time")
|
||
|
|
if (
|
||
|
|
not isinstance(edge, (int, float))
|
||
|
|
or not isinstance(source_time, (int, float))
|
||
|
|
or abs(float(edge) - expected_edge) > 1e-8
|
||
|
|
or float(source_time) >= expected_edge - 1e-8
|
||
|
|
):
|
||
|
|
raise ValueError(f"Non-causal person record at tick {tick}: {feature_path}")
|
||
|
|
people[tick, :count] = features[start:end]
|
||
|
|
person_valid[tick, :count] = True
|
||
|
|
|
||
|
|
target_array = np.load(target_path, allow_pickle=False)
|
||
|
|
targets = torch.from_numpy(np.asarray(target_array, dtype=np.float32))
|
||
|
|
if targets.shape != (num_steps, NUM_CLASSES):
|
||
|
|
raise ValueError(
|
||
|
|
f"Target/feature length mismatch {targets.shape} vs ({num_steps},3): {name}"
|
||
|
|
)
|
||
|
|
_validated_binary_activity_targets(targets)
|
||
|
|
return MuVAPSessionFeatures(
|
||
|
|
name=name,
|
||
|
|
vap=vap,
|
||
|
|
people=people,
|
||
|
|
person_valid=person_valid,
|
||
|
|
targets=targets,
|
||
|
|
task_frame_indices=task_frame_indices,
|
||
|
|
label_indices=label_indices,
|
||
|
|
times=times,
|
||
|
|
contract=contract,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass(frozen=True)
|
||
|
|
class Stage3Window:
|
||
|
|
name: str
|
||
|
|
current_row: int
|
||
|
|
vap: torch.Tensor
|
||
|
|
people: torch.Tensor
|
||
|
|
person_valid: torch.BoolTensor
|
||
|
|
history_valid: torch.BoolTensor
|
||
|
|
targets: torch.Tensor
|
||
|
|
future_valid: torch.BoolTensor
|
||
|
|
|
||
|
|
|
||
|
|
class MuVAPWindowDataset(Dataset[Stage3Window]):
|
||
|
|
"""One fixed past-only 25 Hz window for every 5 Hz prediction cutoff."""
|
||
|
|
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
sessions: Sequence[MuVAPSessionFeatures],
|
||
|
|
*,
|
||
|
|
history_seconds: float,
|
||
|
|
require_full_future: bool,
|
||
|
|
) -> None:
|
||
|
|
if not sessions:
|
||
|
|
raise ValueError("At least one session is required")
|
||
|
|
history_frames = round(history_seconds * MODEL_FRAME_HZ)
|
||
|
|
if history_frames <= 0 or not math.isclose(
|
||
|
|
history_frames / MODEL_FRAME_HZ, history_seconds, abs_tol=1e-8
|
||
|
|
):
|
||
|
|
raise ValueError("history_seconds must be a positive 25 Hz multiple")
|
||
|
|
self.sessions = list(sessions)
|
||
|
|
self.history_frames = history_frames
|
||
|
|
self.require_full_future = require_full_future
|
||
|
|
self.examples: list[tuple[int, int]] = []
|
||
|
|
for session_index, session in enumerate(self.sessions):
|
||
|
|
stop = len(session.targets) - NUM_HORIZONS if require_full_future else len(session.targets)
|
||
|
|
self.examples.extend((session_index, row) for row in range(max(0, stop)))
|
||
|
|
|
||
|
|
def __len__(self) -> int:
|
||
|
|
return len(self.examples)
|
||
|
|
|
||
|
|
def __getitem__(self, index: int) -> Stage3Window:
|
||
|
|
session_index, row = self.examples[index]
|
||
|
|
session = self.sessions[session_index]
|
||
|
|
end = int(session.task_frame_indices[row]) + 1
|
||
|
|
start = max(0, end - self.history_frames)
|
||
|
|
observed = end - start
|
||
|
|
left = self.history_frames - observed
|
||
|
|
vap = torch.zeros(self.history_frames, session.vap_dim, dtype=session.vap.dtype)
|
||
|
|
vap[left:] = session.vap[start:end]
|
||
|
|
people = torch.zeros(
|
||
|
|
self.history_frames,
|
||
|
|
session.people.shape[1],
|
||
|
|
session.asd_dim,
|
||
|
|
dtype=session.people.dtype,
|
||
|
|
)
|
||
|
|
person_valid = torch.zeros(
|
||
|
|
self.history_frames, session.people.shape[1], dtype=torch.bool
|
||
|
|
)
|
||
|
|
if session.people.shape[1]:
|
||
|
|
people[left:] = session.people[start:end]
|
||
|
|
person_valid[left:] = session.person_valid[start:end]
|
||
|
|
history_valid = torch.zeros(self.history_frames, dtype=torch.bool)
|
||
|
|
history_valid[left:] = True
|
||
|
|
targets = torch.zeros(NUM_HORIZONS, NUM_CLASSES, dtype=torch.float32)
|
||
|
|
future_valid = torch.zeros(NUM_HORIZONS, dtype=torch.bool)
|
||
|
|
available = min(NUM_HORIZONS, len(session.targets) - row - 1)
|
||
|
|
if available:
|
||
|
|
targets[:available] = session.targets[row + 1 : row + 1 + available]
|
||
|
|
future_valid[:available] = True
|
||
|
|
if self.require_full_future and not bool(future_valid.all()):
|
||
|
|
raise RuntimeError("Training dataset exposed an incomplete future target")
|
||
|
|
return Stage3Window(
|
||
|
|
name=session.name,
|
||
|
|
current_row=row,
|
||
|
|
vap=vap,
|
||
|
|
people=people,
|
||
|
|
person_valid=person_valid,
|
||
|
|
history_valid=history_valid,
|
||
|
|
targets=targets,
|
||
|
|
future_valid=future_valid,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def collate_stage3_windows(items: Sequence[Stage3Window]) -> dict[str, Any]:
|
||
|
|
if not items:
|
||
|
|
raise ValueError("Cannot collate an empty batch")
|
||
|
|
max_people = max(item.people.shape[1] for item in items)
|
||
|
|
batch_people = torch.zeros(
|
||
|
|
len(items),
|
||
|
|
items[0].people.shape[0],
|
||
|
|
max_people,
|
||
|
|
items[0].people.shape[-1],
|
||
|
|
dtype=items[0].people.dtype,
|
||
|
|
)
|
||
|
|
batch_person_valid = torch.zeros(
|
||
|
|
len(items), items[0].people.shape[0], max_people, dtype=torch.bool
|
||
|
|
)
|
||
|
|
for index, item in enumerate(items):
|
||
|
|
count = item.people.shape[1]
|
||
|
|
if count:
|
||
|
|
batch_people[index, :, :count] = item.people
|
||
|
|
batch_person_valid[index, :, :count] = item.person_valid
|
||
|
|
return {
|
||
|
|
"names": [item.name for item in items],
|
||
|
|
"current_rows": torch.tensor([item.current_row for item in items], dtype=torch.long),
|
||
|
|
"vap": torch.stack([item.vap for item in items]),
|
||
|
|
"people": batch_people,
|
||
|
|
"person_valid": batch_person_valid,
|
||
|
|
"history_valid": torch.stack([item.history_valid for item in items]),
|
||
|
|
"targets": torch.stack([item.targets for item in items]),
|
||
|
|
"future_valid": torch.stack([item.future_valid for item in items]),
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def sinusoidal_positions(length: int, dimension: int) -> torch.Tensor:
|
||
|
|
position = torch.arange(length, dtype=torch.float32)[:, None]
|
||
|
|
frequencies = torch.exp(
|
||
|
|
torch.arange(0, dimension, 2, dtype=torch.float32)
|
||
|
|
* (-math.log(10_000.0) / dimension)
|
||
|
|
)
|
||
|
|
result = torch.zeros(length, dimension, dtype=torch.float32)
|
||
|
|
result[:, 0::2] = torch.sin(position * frequencies)
|
||
|
|
result[:, 1::2] = torch.cos(position * frequencies[: result[:, 1::2].shape[1]])
|
||
|
|
return result
|
||
|
|
|
||
|
|
|
||
|
|
class MuVAPEgoSpeakStage3(nn.Module):
|
||
|
|
"""Paper-inspired MuVAP frame/temporal fusion with an EgoSpeak head.
|
||
|
|
|
||
|
|
Published VAP and ASD states are inputs and remain outside this trainable
|
||
|
|
module. The model trains only the Stage-3 projections, same-frame person
|
||
|
|
cross-attention, causal temporal transformer and the task head.
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
*,
|
||
|
|
vap_dim: int,
|
||
|
|
asd_dim: int,
|
||
|
|
model_dim: int = 256,
|
||
|
|
num_heads: int = 4,
|
||
|
|
dropout: float = 0.1,
|
||
|
|
temporal_layers: int = 1,
|
||
|
|
history_frames: int = 200,
|
||
|
|
num_classes: int = NUM_CLASSES,
|
||
|
|
person_evidence: str = "aligned",
|
||
|
|
) -> None:
|
||
|
|
super().__init__()
|
||
|
|
if min(vap_dim, asd_dim, model_dim, num_heads, temporal_layers, history_frames) <= 0:
|
||
|
|
raise ValueError("Stage-3 dimensions must be positive")
|
||
|
|
if model_dim % num_heads:
|
||
|
|
raise ValueError("model_dim must be divisible by num_heads")
|
||
|
|
if person_evidence not in {"aligned", "none"}:
|
||
|
|
raise ValueError("person_evidence must be 'aligned' or 'none'")
|
||
|
|
self.vap_dim = vap_dim
|
||
|
|
self.asd_dim = asd_dim
|
||
|
|
self.model_dim = model_dim
|
||
|
|
self.num_heads = num_heads
|
||
|
|
self.dropout_probability = dropout
|
||
|
|
self.temporal_layers = temporal_layers
|
||
|
|
self.history_frames = history_frames
|
||
|
|
self.num_classes = num_classes
|
||
|
|
self.person_evidence = person_evidence
|
||
|
|
self.vap_projection = nn.Sequential(
|
||
|
|
nn.Linear(vap_dim, model_dim),
|
||
|
|
nn.LayerNorm(model_dim),
|
||
|
|
)
|
||
|
|
self.asd_projection = nn.Sequential(
|
||
|
|
nn.Linear(asd_dim, model_dim),
|
||
|
|
nn.LayerNorm(model_dim),
|
||
|
|
)
|
||
|
|
self.null_person = nn.Parameter(torch.zeros(1, 1, model_dim))
|
||
|
|
nn.init.normal_(self.null_person, std=0.02)
|
||
|
|
self.frame_attention = nn.MultiheadAttention(
|
||
|
|
model_dim,
|
||
|
|
num_heads,
|
||
|
|
dropout=dropout,
|
||
|
|
batch_first=True,
|
||
|
|
)
|
||
|
|
self.frame_attention_norm = nn.LayerNorm(model_dim)
|
||
|
|
self.frame_ff_norm = nn.LayerNorm(model_dim)
|
||
|
|
self.frame_ff = nn.Sequential(
|
||
|
|
nn.Linear(model_dim, model_dim * 4),
|
||
|
|
nn.GELU(),
|
||
|
|
nn.Dropout(dropout),
|
||
|
|
nn.Linear(model_dim * 4, model_dim),
|
||
|
|
)
|
||
|
|
self.dropout = nn.Dropout(dropout)
|
||
|
|
temporal_layer = nn.TransformerEncoderLayer(
|
||
|
|
d_model=model_dim,
|
||
|
|
nhead=num_heads,
|
||
|
|
dim_feedforward=model_dim * 4,
|
||
|
|
dropout=dropout,
|
||
|
|
activation="gelu",
|
||
|
|
batch_first=True,
|
||
|
|
norm_first=True,
|
||
|
|
)
|
||
|
|
self.temporal = nn.TransformerEncoder(
|
||
|
|
temporal_layer,
|
||
|
|
num_layers=temporal_layers,
|
||
|
|
norm=nn.LayerNorm(model_dim),
|
||
|
|
)
|
||
|
|
self.head = nn.Linear(model_dim, NUM_HORIZONS * num_classes)
|
||
|
|
self.register_buffer(
|
||
|
|
"position_encoding",
|
||
|
|
sinusoidal_positions(history_frames, model_dim),
|
||
|
|
persistent=False,
|
||
|
|
)
|
||
|
|
self.register_buffer(
|
||
|
|
"causal_mask",
|
||
|
|
torch.triu(torch.ones(history_frames, history_frames, dtype=torch.bool), diagonal=1),
|
||
|
|
persistent=False,
|
||
|
|
)
|
||
|
|
|
||
|
|
def forward(
|
||
|
|
self,
|
||
|
|
vap: torch.Tensor,
|
||
|
|
people: torch.Tensor,
|
||
|
|
person_valid: torch.BoolTensor,
|
||
|
|
history_valid: torch.BoolTensor,
|
||
|
|
*,
|
||
|
|
return_embedding: bool = False,
|
||
|
|
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
|
||
|
|
if vap.ndim != 3 or vap.shape[-1] != self.vap_dim:
|
||
|
|
raise ValueError(f"Expected VAP [B,H,{self.vap_dim}], got {tuple(vap.shape)}")
|
||
|
|
batch, history, _ = vap.shape
|
||
|
|
if history != self.history_frames:
|
||
|
|
raise ValueError(f"Expected {self.history_frames} history frames, got {history}")
|
||
|
|
if people.ndim != 4 or people.shape[:2] != (batch, history) or people.shape[-1] != self.asd_dim:
|
||
|
|
raise ValueError(
|
||
|
|
f"Expected people [B,H,N,{self.asd_dim}], got {tuple(people.shape)}"
|
||
|
|
)
|
||
|
|
if person_valid.dtype != torch.bool or person_valid.shape != people.shape[:3]:
|
||
|
|
raise ValueError("person_valid must be bool [B,H,N]")
|
||
|
|
if history_valid.dtype != torch.bool or history_valid.shape != (batch, history):
|
||
|
|
raise ValueError("history_valid must be bool [B,H]")
|
||
|
|
if not bool(history_valid[:, -1].all()):
|
||
|
|
raise ValueError("Every Stage-3 window must end at an observed cutoff")
|
||
|
|
if not torch.isfinite(vap).all() or not torch.isfinite(people).all():
|
||
|
|
raise ValueError("Stage-3 inputs contain NaN/Inf")
|
||
|
|
if self.person_evidence == "none":
|
||
|
|
# Capacity-matched VAP-only control. The same frame and temporal
|
||
|
|
# blocks remain trainable; only observed ASD tokens are hidden.
|
||
|
|
person_valid = torch.zeros_like(person_valid)
|
||
|
|
|
||
|
|
query = self.vap_projection(vap).reshape(batch * history, 1, self.model_dim)
|
||
|
|
count = people.shape[2]
|
||
|
|
person_tokens = self.asd_projection(people).reshape(
|
||
|
|
batch * history, count, self.model_dim
|
||
|
|
)
|
||
|
|
null = self.null_person.expand(batch * history, -1, -1)
|
||
|
|
person_tokens = torch.cat((person_tokens, null), dim=1)
|
||
|
|
key_padding = torch.cat(
|
||
|
|
(
|
||
|
|
~person_valid.reshape(batch * history, count),
|
||
|
|
torch.zeros(batch * history, 1, dtype=torch.bool, device=vap.device),
|
||
|
|
),
|
||
|
|
dim=1,
|
||
|
|
)
|
||
|
|
selected, _ = self.frame_attention(
|
||
|
|
query,
|
||
|
|
person_tokens,
|
||
|
|
person_tokens,
|
||
|
|
key_padding_mask=key_padding,
|
||
|
|
need_weights=False,
|
||
|
|
)
|
||
|
|
frame = self.frame_attention_norm(query + self.dropout(selected))
|
||
|
|
frame = frame + self.dropout(self.frame_ff(self.frame_ff_norm(frame)))
|
||
|
|
frame = frame.reshape(batch, history, self.model_dim)
|
||
|
|
frame = frame + self.position_encoding[None].to(frame.dtype)
|
||
|
|
|
||
|
|
# Windows are left-padded so the final token is always the current
|
||
|
|
# 5 Hz cutoff. A plain key-padding mask leaves the earliest padded
|
||
|
|
# queries with no visible key under causality, which can yield NaNs
|
||
|
|
# (and can contaminate later layers). Give each padded query only its
|
||
|
|
# own diagonal while every observed query still masks every padded
|
||
|
|
# key. The mask is repeated because MultiheadAttention accepts one
|
||
|
|
# [H,H] matrix per batch/head pair.
|
||
|
|
temporal_mask = self.causal_mask[None].expand(batch, -1, -1).clone()
|
||
|
|
temporal_mask |= (~history_valid)[:, None, :]
|
||
|
|
padded_batch, padded_query = torch.where(~history_valid)
|
||
|
|
temporal_mask[padded_batch, padded_query, padded_query] = False
|
||
|
|
temporal_mask = temporal_mask.repeat_interleave(self.num_heads, dim=0)
|
||
|
|
temporal = self.temporal(
|
||
|
|
frame,
|
||
|
|
mask=temporal_mask,
|
||
|
|
)
|
||
|
|
embedding = temporal[:, -1]
|
||
|
|
logits = self.head(embedding).reshape(batch, NUM_HORIZONS, self.num_classes)
|
||
|
|
if return_embedding:
|
||
|
|
return logits, embedding
|
||
|
|
return logits
|
||
|
|
|
||
|
|
|
||
|
|
def model_config(model: MuVAPEgoSpeakStage3) -> dict[str, Any]:
|
||
|
|
return {
|
||
|
|
"vap_dim": model.vap_dim,
|
||
|
|
"asd_dim": model.asd_dim,
|
||
|
|
"model_dim": model.model_dim,
|
||
|
|
"num_heads": model.num_heads,
|
||
|
|
"dropout": model.dropout_probability,
|
||
|
|
"temporal_layers": model.temporal_layers,
|
||
|
|
"history_frames": model.history_frames,
|
||
|
|
"num_classes": model.num_classes,
|
||
|
|
"person_evidence": model.person_evidence,
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def build_model(config: dict[str, Any]) -> MuVAPEgoSpeakStage3:
|
||
|
|
required = {
|
||
|
|
"vap_dim",
|
||
|
|
"asd_dim",
|
||
|
|
"model_dim",
|
||
|
|
"num_heads",
|
||
|
|
"dropout",
|
||
|
|
"temporal_layers",
|
||
|
|
"history_frames",
|
||
|
|
"num_classes",
|
||
|
|
"person_evidence",
|
||
|
|
}
|
||
|
|
if set(config) != required:
|
||
|
|
raise ValueError(f"Invalid Stage-3 model config keys: {sorted(set(config) ^ required)}")
|
||
|
|
return MuVAPEgoSpeakStage3(**config)
|
||
|
|
|
||
|
|
|
||
|
|
@torch.inference_mode()
|
||
|
|
def predict_loader(
|
||
|
|
model: MuVAPEgoSpeakStage3,
|
||
|
|
loader: Iterable[dict[str, Any]],
|
||
|
|
*,
|
||
|
|
device: torch.device,
|
||
|
|
loss_name: str,
|
||
|
|
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[str]]:
|
||
|
|
model.eval()
|
||
|
|
probabilities, targets, rows, names = [], [], [], []
|
||
|
|
for batch in loader:
|
||
|
|
logits = model(
|
||
|
|
batch["vap"].to(device=device, dtype=torch.float32, non_blocking=True),
|
||
|
|
batch["people"].to(device=device, dtype=torch.float32, non_blocking=True),
|
||
|
|
batch["person_valid"].to(device=device, non_blocking=True),
|
||
|
|
batch["history_valid"].to(device=device, non_blocking=True),
|
||
|
|
)
|
||
|
|
probabilities.append(marginal_probabilities_from_logits(logits, loss_name).cpu())
|
||
|
|
targets.append(batch["targets"])
|
||
|
|
rows.append(batch["current_rows"])
|
||
|
|
names.extend(batch["names"])
|
||
|
|
return (
|
||
|
|
torch.cat(probabilities),
|
||
|
|
torch.cat(targets),
|
||
|
|
torch.cat(rows),
|
||
|
|
names,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def evaluate_valid_predictions(
|
||
|
|
probabilities: torch.Tensor,
|
||
|
|
targets: torch.Tensor,
|
||
|
|
) -> dict[str, Any]:
|
||
|
|
if probabilities.shape != targets.shape or probabilities.shape[-2:] != (
|
||
|
|
NUM_HORIZONS,
|
||
|
|
NUM_CLASSES,
|
||
|
|
):
|
||
|
|
raise ValueError("Evaluation expects matching [N,10,3] arrays")
|
||
|
|
results = []
|
||
|
|
for horizon in range(NUM_HORIZONS):
|
||
|
|
result = average_precision_metrics(
|
||
|
|
probabilities[:, horizon].numpy(),
|
||
|
|
targets[:, horizon].numpy(),
|
||
|
|
)
|
||
|
|
result["horizon"] = horizon + 1
|
||
|
|
result["seconds"] = (horizon + 1) * STEP_SECONDS
|
||
|
|
results.append(result)
|
||
|
|
return {
|
||
|
|
"horizons": results,
|
||
|
|
"mean_map": float(np.mean([result["map"] for result in results])),
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def align_session_predictions(
|
||
|
|
probabilities: np.ndarray,
|
||
|
|
num_steps: int,
|
||
|
|
) -> tuple[np.ndarray, np.ndarray]:
|
||
|
|
if probabilities.shape != (num_steps, NUM_HORIZONS, NUM_CLASSES):
|
||
|
|
raise ValueError(
|
||
|
|
f"Expected [{num_steps},{NUM_HORIZONS},{NUM_CLASSES}], got "
|
||
|
|
f"{probabilities.shape}"
|
||
|
|
)
|
||
|
|
predictions = np.zeros((num_steps, NUM_CLASSES, NUM_HORIZONS), dtype=np.float32)
|
||
|
|
valid = np.zeros((num_steps, NUM_HORIZONS), dtype=bool)
|
||
|
|
current = np.arange(num_steps)
|
||
|
|
for horizon in range(NUM_HORIZONS):
|
||
|
|
future = current + horizon + 1
|
||
|
|
usable = future < num_steps
|
||
|
|
predictions[future[usable], :, horizon] = probabilities[usable, horizon]
|
||
|
|
valid[future[usable], horizon] = True
|
||
|
|
return predictions, valid
|
||
|
|
|
||
|
|
|
||
|
|
def evaluate_aligned(
|
||
|
|
session_predictions: dict[str, np.ndarray],
|
||
|
|
session_valid: dict[str, np.ndarray],
|
||
|
|
session_targets: dict[str, np.ndarray],
|
||
|
|
) -> dict[str, Any]:
|
||
|
|
protocols: dict[str, Any] = {}
|
||
|
|
for protocol in ("egospeak_full_timeline", "causal_valid_only"):
|
||
|
|
horizons = []
|
||
|
|
for horizon in range(NUM_HORIZONS):
|
||
|
|
if protocol == "egospeak_full_timeline":
|
||
|
|
scores = np.concatenate(
|
||
|
|
[value[:, :, horizon] for value in session_predictions.values()]
|
||
|
|
)
|
||
|
|
targets = np.concatenate(list(session_targets.values()))
|
||
|
|
else:
|
||
|
|
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
|
||
|
|
]
|
||
|
|
)
|
||
|
|
result = average_precision_metrics(scores, targets)
|
||
|
|
result["horizon"] = horizon + 1
|
||
|
|
result["seconds"] = (horizon + 1) * STEP_SECONDS
|
||
|
|
horizons.append(result)
|
||
|
|
protocols[protocol] = {
|
||
|
|
"horizons": horizons,
|
||
|
|
"mean_map": float(np.mean([result["map"] for result in horizons])),
|
||
|
|
}
|
||
|
|
return protocols
|
||
|
|
|
||
|
|
|
||
|
|
def print_metrics(metrics: dict[str, Any]) -> None:
|
||
|
|
for protocol in ("egospeak_full_timeline", "causal_valid_only"):
|
||
|
|
print(f"\n{protocol}:")
|
||
|
|
for result in metrics[protocol]["horizons"]:
|
||
|
|
values = ", ".join(
|
||
|
|
f"{name}={result['class_ap'][name] * 100:.5f}" for name in CLASS_NAMES
|
||
|
|
)
|
||
|
|
print(
|
||
|
|
f"Action anticipation ({result['seconds']:.2f}s) perframe mAP: "
|
||
|
|
f"{result['map'] * 100:.5f}"
|
||
|
|
)
|
||
|
|
print(f" AP: {values}; n={result['num_steps']}")
|
||
|
|
print(
|
||
|
|
f"Action anticipation (mean) perframe mAP: "
|
||
|
|
f"{metrics[protocol]['mean_map'] * 100:.5f}"
|
||
|
|
)
|