diff --git a/easycom_head.py b/easycom_head.py index ff5c138..59cd8cb 100644 --- a/easycom_head.py +++ b/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)}") \ No newline at end of file diff --git a/infer_turn_head.py b/infer_turn_head.py index bcd1da6..1f400e0 100644 --- a/infer_turn_head.py +++ b/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() \ No newline at end of file diff --git a/train_turn_head.py b/train_turn_head.py index cacb3d8..92c2e36 100644 --- a/train_turn_head.py +++ b/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() \ No newline at end of file