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

infer_turn_head.py

@@ -1,4 +1,4 @@
-"""Run the Omni anticipation head and evaluate it with EgoSpeak's EasyCom metric."""
+"""Run the Omni anticipation head with EgoSpeak's EasyCom/Ego4D metric."""
 
 from __future__ import annotations
 
@@ -13,6 +13,7 @@ from tqdm.auto import tqdm
 
 from easycom_head import (
     CLASS_NAMES,
+    DATASET_NAMES,
     NUM_CLASSES,
     NUM_HORIZONS,
     OFFICIAL_TEST_SESSIONS,
@@ -20,8 +21,9 @@ from easycom_head import (
     TurnPredictionHead,
     average_precision_metrics,
     checkpoint_head_config,
+    load_ego4d_manifest_split,
     load_session_data,
-    normalize_session_name,
+    normalize_sample_name,
 )
 
 
@@ -31,11 +33,25 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--target-root", type=Path, required=True)
     parser.add_argument("--checkpoint", type=Path, required=True)
     parser.add_argument("--output-dir", type=Path, required=True)
+    parser.add_argument(
+        "--dataset",
+        choices=DATASET_NAMES,
+        default="easycom",
+        help="Dataset naming/split contract (default: easycom).",
+    )
+    parser.add_argument(
+        "--manifest",
+        type=Path,
+        help="Required for Ego4D; TSV containing split and uid columns.",
+    )
     parser.add_argument(
         "--sessions",
         nargs="+",
-        default=list(OFFICIAL_TEST_SESSIONS),
-        help="Evaluation sessions; default is the official Session 3/2/1 test split.",
+        default=None,
+        help=(
+            "Evaluation samples. Defaults to EasyCom Session 3/2/1 or all "
+            "Ego4D manifest test UIDs."
+        ),
     )
     parser.add_argument("--batch-size", type=int, default=1024)
     parser.add_argument("--device", default="cuda:0")
@@ -211,12 +227,36 @@ def main() -> None:
     args = parse_args()
     if args.batch_size <= 0:
         raise ValueError("--batch-size must be positive")
+    if args.dataset == "ego4d" and args.manifest is None:
+        raise ValueError("--manifest is required when --dataset ego4d")
     device = _resolve_device(args.device)
     expected_mode = None if args.expected_mode == "any" else args.expected_mode
-    sessions = [normalize_session_name(item) for item in args.sessions]
+    if args.dataset == "ego4d":
+        assert args.manifest is not None
+        official_test = load_ego4d_manifest_split(args.manifest, "test")
+        requested_values = args.sessions or official_test
+    else:
+        official_test = list(OFFICIAL_TEST_SESSIONS)
+        requested_values = args.sessions or official_test
+    sessions = [
+        normalize_sample_name(item, args.dataset) for item in requested_values
+    ]
     if len(sessions) != len(set(sessions)):
-        raise ValueError("Duplicate sessions were supplied")
+        raise ValueError("Duplicate evaluation samples were supplied")
+    if args.dataset == "ego4d":
+        unknown = set(sessions) - set(official_test)
+        if unknown:
+            raise ValueError(
+                "Ego4D evaluation samples must belong to the manifest test "
+                f"split; invalid: {sorted(unknown)}"
+            )
     head, checkpoint_payload = _load_head(args.checkpoint, device)
+    checkpoint_dataset = checkpoint_payload.get("dataset")
+    if checkpoint_dataset is not None and checkpoint_dataset != args.dataset:
+        raise ValueError(
+            f"Checkpoint dataset is {checkpoint_dataset!r}, but --dataset is "
+            f"{args.dataset!r}"
+        )
     head_config = checkpoint_head_config(checkpoint_payload)
 
     args.output_dir.mkdir(parents=True, exist_ok=True)
@@ -236,6 +276,7 @@ def main() -> None:
             args.target_root,
             name,
             expected_mode=expected_mode,
+            dataset=args.dataset,
         )
         if data.hidden_size != int(head_config["hidden_size"]):
             raise ValueError(
@@ -270,6 +311,12 @@ def main() -> None:
     metrics.update(
         {
             "sessions": sessions,
+            "dataset": args.dataset,
+            "manifest": (
+                str(args.manifest.expanduser().resolve())
+                if args.manifest is not None
+                else None
+            ),
             "classes": list(CLASS_NAMES),
             "prediction_layout": "[target_time, class, future_horizon]",
             "checkpoint": str(args.checkpoint.resolve()),
@@ -286,4 +333,4 @@ def main() -> None:
 
 
 if __name__ == "__main__":
-    main()
+    main()

train_turn_head.py

@@ -1,4 +1,4 @@
-"""Train the frozen-Omni EasyCom anticipation head.
+"""Train the frozen-Omni EasyCom or Ego4D anticipation head.
 
 The defaults follow EgoSpeak's EasyCom LSTR optimization config where the
 architectures are comparable: 50 epochs, Adam, 7e-5 learning rate, 5e-5 weight
@@ -22,6 +22,7 @@ from tqdm.auto import tqdm
 
 from easycom_head import (
     CLASS_NAMES,
+    DATASET_NAMES,
     NUM_CLASSES,
     NUM_HORIZONS,
     OFFICIAL_TRAIN_SESSIONS,
@@ -29,9 +30,10 @@ from easycom_head import (
     TurnPredictionHead,
     average_precision_metrics,
     egospeak_multihorizon_loss,
+    load_ego4d_manifest_split,
     load_session_data,
     miniroad_multihorizon_loss,
-    normalize_session_name,
+    normalize_sample_name,
     validate_disjoint_sessions,
 )
 
@@ -41,11 +43,25 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--feature-root", type=Path, required=True)
     parser.add_argument("--target-root", type=Path, required=True)
     parser.add_argument("--output", type=Path, required=True)
+    parser.add_argument(
+        "--dataset",
+        choices=DATASET_NAMES,
+        default="easycom",
+        help="Dataset naming/split contract (default: easycom).",
+    )
+    parser.add_argument(
+        "--manifest",
+        type=Path,
+        help="Required for Ego4D; TSV containing split and uid columns.",
+    )
     parser.add_argument(
         "--sessions",
         nargs="+",
-        default=list(OFFICIAL_TRAIN_SESSIONS),
-        help="Training sessions; default is the official nine-session split.",
+        default=None,
+        help=(
+            "Training samples. Defaults to the official EasyCom sessions or "
+            "all Ego4D manifest train UIDs."
+        ),
     )
     parser.add_argument(
         "--dev-sessions",
@@ -97,6 +113,8 @@ def _validate_args(args: argparse.Namespace) -> None:
         raise ValueError("--dropout must be in [0, 1)")
     if args.num_workers < 0:
         raise ValueError("--num-workers cannot be negative")
+    if args.dataset == "ego4d" and args.manifest is None:
+        raise ValueError("--manifest is required when --dataset ego4d")
 
 
 def _resolve_device(name: str) -> torch.device:
@@ -119,6 +137,7 @@ def _load_examples(
     target_root: Path,
     sessions: list[str],
     expected_mode: str | None,
+    dataset: str,
 ) -> tuple[torch.Tensor, torch.Tensor, int, list[dict[str, Any]]]:
     all_states: list[torch.Tensor] = []
     all_targets: list[torch.Tensor] = []
@@ -130,6 +149,7 @@ def _load_examples(
             target_root,
             name,
             expected_mode=expected_mode,
+            dataset=dataset,
         )
         states, targets, _ = data.future_examples(NUM_HORIZONS)
         if hidden_size is None:
@@ -252,6 +272,12 @@ def _checkpoint(
             ],
         },
         "classes": list(CLASS_NAMES),
+        "dataset": args.dataset,
+        "manifest": (
+            str(args.manifest.expanduser().resolve())
+            if args.manifest is not None
+            else None
+        ),
         "epoch": epoch,
         "best_epoch": best_epoch,
         "best_dev_map": best_dev_map,
@@ -283,14 +309,40 @@ def main() -> None:
     device = _resolve_device(args.device)
     expected_mode = None if args.expected_mode == "any" else args.expected_mode
 
-    requested_train = [normalize_session_name(item) for item in args.sessions]
-    dev_sessions = [normalize_session_name(item) for item in args.dev_sessions]
+    if args.dataset == "ego4d":
+        assert args.manifest is not None
+        official_train = load_ego4d_manifest_split(args.manifest, "train")
+        official_test = set(load_ego4d_manifest_split(args.manifest, "test"))
+        requested_values = args.sessions or official_train
+    else:
+        official_train = list(OFFICIAL_TRAIN_SESSIONS)
+        official_test = {"Session_1", "Session_2", "Session_3"}
+        requested_values = args.sessions or official_train
+
+    requested_train = [
+        normalize_sample_name(item, args.dataset) for item in requested_values
+    ]
+    dev_sessions = [
+        normalize_sample_name(item, args.dataset) for item in args.dev_sessions
+    ]
+    if len(requested_train) != len(set(requested_train)):
+        raise ValueError("Duplicate training samples were supplied")
+    if len(dev_sessions) != len(set(dev_sessions)):
+        raise ValueError("Duplicate development samples were supplied")
+
+    if args.dataset == "ego4d":
+        allowed_train = set(official_train)
+        unknown = (set(requested_train) | set(dev_sessions)) - allowed_train
+        if unknown:
+            raise ValueError(
+                "Ego4D training/dev samples must belong to the manifest train "
+                f"split; invalid: {sorted(unknown)}"
+            )
     train_sessions = [item for item in requested_train if item not in dev_sessions]
     validate_disjoint_sessions(train_sessions, dev_sessions)
     if not train_sessions:
         raise ValueError("No training sessions remain after removing dev sessions")
-    forbidden_test = {"Session_1", "Session_2", "Session_3"}
-    leaked = forbidden_test & (set(train_sessions) | set(dev_sessions))
+    leaked = official_test & (set(train_sessions) | set(dev_sessions))
     if leaked:
         raise ValueError(
             f"Official test sessions cannot be used for training/selection: "
@@ -305,6 +357,7 @@ def main() -> None:
         args.target_root,
         train_sessions,
         expected_mode,
+        args.dataset,
     )
     print(f"training examples: {len(states)}, feature_dim={hidden_size}")
     dev_states: torch.Tensor | None = None
@@ -315,6 +368,7 @@ def main() -> None:
             args.target_root,
             dev_sessions,
             expected_mode,
+            args.dataset,
         )
         if dev_hidden_size != hidden_size:
             raise ValueError(
@@ -473,4 +527,4 @@ def main() -> None:
 
 
 if __name__ == "__main__":
-    main()
+    main()