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.

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