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.

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