"""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}" )