icassp2027-experiment/easycom_head.py

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


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])

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