Enhance EasyCom and Ego4D integration by adding dataset support and manifest handling. Update normalization functions to accommodate both datasets, and modify training and inference scripts to require manifest files for Ego4D. Improve argument parsing and validation for dataset-specific operations, ensuring proper session management and error handling. Update documentation to reflect changes in dataset usage and evaluation processes.

easycom_head.py

@@ -1,10 +1,12 @@
-"""Shared head, labels, and metrics for the EasyCom EgoSpeak experiment."""
+"""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
@@ -31,6 +33,8 @@ CLASS_NAMES = (
 STEP_SECONDS = 0.2
 NUM_HORIZONS = 10
 NUM_CLASSES = 3
+DATASET_NAMES = ("easycom", "ego4d")
+DatasetName = Literal["easycom", "ego4d"]
 
 
 class TurnPredictionHead(nn.Module):
@@ -144,6 +148,65 @@ def normalize_session_name(value: str | int) -> str:
     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
@@ -177,10 +240,11 @@ def load_session_data(
     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_session_name(session)
+    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():
@@ -296,4 +360,4 @@ def validate_disjoint_sessions(
 ) -> None:
     overlap = set(first) & set(second)
     if overlap:
-        raise ValueError(f"Session sets overlap: {sorted(overlap)}")
+        raise ValueError(f"Session sets overlap: {sorted(overlap)}")