Add Ego4D local AV ablation workflow

extract_turn_local_features.py

@@ -19,10 +19,10 @@ import numpy as np
 import torch
 from tqdm.auto import tqdm
 
-from easycom_head import normalize_sample_name
+from easycom_head import load_ego4d_manifest_split, normalize_sample_name
 from local_turn_features import (
     LOCAL_FORMAT, causal_frame_ticks, crop_box, easycom_boxes,
-    pooled_recent_visual, reference_alignment,
+    load_local_features, pooled_recent_visual, reference_alignment,
 )
 
 
@@ -164,6 +164,7 @@ def ego4d_clip_boxes(clip: dict[str, Any]) -> dict[int, list[dict[str, Any]]]:
 
 def load_ego4d_boxes(paths: list[Path], names: list[str]) -> dict[str, Any]:
     result = {}
+    selected = set(names)
     for path in paths:
         print(f"Reading provided tracks: {path}", flush=True)
         with path.open() as handle:
@@ -171,7 +172,7 @@ def load_ego4d_boxes(paths: list[Path], names: list[str]) -> dict[str, Any]:
         for video in payload["videos"]:
             for clip in video["clips"]:
                 uid = clip["clip_uid"]
-                if uid in names:
+                if uid in selected:
                     if uid in result:
                         raise ValueError(f"Duplicate Ego4D annotation UID: {uid}")
                     result[uid] = ego4d_clip_boxes(clip)
@@ -338,6 +339,54 @@ def extract_local(media, annotations, encoder, times: torch.Tensor, args: argpar
     }
 
 
+def select_samples(args: argparse.Namespace) -> list[str]:
+    if args.manifest is not None:
+        names = load_ego4d_manifest_split(args.manifest, args.split)
+        # One-based inclusive range in the selected split, then an optional limit.
+        names = names[args.start_index - 1:args.end_index]
+        if args.limit is not None:
+            names = names[:args.limit]
+    else:
+        names = [normalize_sample_name(value, args.dataset) for value in args.samples]
+    if not names:
+        raise ValueError("No samples selected")
+    if len(set(names)) != len(names):
+        raise ValueError("Duplicate samples")
+    return names
+
+
+def requested_local_config(args: argparse.Namespace) -> dict[str, Any]:
+    """Settings that must match before reusing an existing complete sidecar."""
+    return {
+        "encoder_id": args.encoder_id or args.model_id,
+        "dtype": args.dtype,
+        "attention": "flash_attention_2" if args.flash_attention_2 else "default",
+        "visual_fps": args.video_fps, "visual_window_seconds": args.visual_window_seconds,
+        "audio_window_seconds": args.audio_window_seconds,
+        "visual_pooling": "spatial-mean_then_equal-person-mean_then_equal-valid-frame-mean",
+        "audio_pooling": "valid-output-token-mean", "audio_source": "video-mixed-mono-16khz",
+        "box_kind": args.box_kind if args.dataset == "easycom" else "ego4d-provided-tracks",
+        "box_coordinate_space": args.box_coordinate_space, "crop_expansion": args.crop_expansion,
+        "max_pixels": args.max_pixels, "uses_speech_labels": False,
+        "causal_policy": "half-open-trailing-window", "box_protocol": "provided-boxes-not-online-detector",
+    }
+
+
+def validate_reusable_sidecar(args, name, reference, version, expected_contract=None):
+    path = args.output_root / f"{name}.pt"
+    _, contract = load_local_features(path, reference, sample_id=name, dataset=args.dataset,
+                                      modalities="av", expected_contract=expected_contract)
+    config = contract["config"]
+    expected = requested_local_config(args)
+    if version is not None:
+        expected["transformers"] = version
+    mismatches = [key for key, value in expected.items() if config.get(key) != value]
+    if mismatches or not config.get("model_config_sha256") or not config.get("transformers"):
+        raise ValueError(f"Cannot reuse {path}: incompatible/incomplete encoder settings {mismatches}; "
+                         "use a new output root or --overwrite")
+    return contract
+
+
 def parse_args(argv=None):
     p = argparse.ArgumentParser(description=__doc__)
     p.add_argument("--dataset", choices=("easycom", "ego4d"), default="easycom")
@@ -345,7 +394,13 @@ def parse_args(argv=None):
     p.add_argument("--feature-root", required=True, type=Path)
     p.add_argument("--target-root", required=True, type=Path)
     p.add_argument("--output-root", required=True, type=Path)
-    p.add_argument("--samples", required=True, nargs="+", help="EasyCom session numbers or Ego4D clip UIDs")
+    selection = p.add_mutually_exclusive_group(required=True)
+    selection.add_argument("--samples", nargs="+", help="EasyCom session numbers or Ego4D clip UIDs")
+    selection.add_argument("--manifest", type=Path, help="Ego4D EgoSpeak TSV; never creates a new split")
+    p.add_argument("--split", choices=("all", "train", "test"), default="all")
+    p.add_argument("--start-index", type=int, default=1, help="One-based index AFTER split filtering")
+    p.add_argument("--end-index", type=int, help="Inclusive index AFTER split filtering")
+    p.add_argument("--limit", type=int, help="First N clips in the selected manifest range")
     p.add_argument("--ego4d-annotations", type=Path, nargs="+")
     p.add_argument("--model-id", default="Qwen/Qwen2.5-Omni-3B")
     p.add_argument("--encoder-id", help="Stable checkpoint identity for relocated copies; default: model-id")
@@ -370,9 +425,25 @@ def parse_args(argv=None):
     p.add_argument("--decoder-gpu-id", type=int, default=0, help="Visible-device index, not physical GPU number")
     p.add_argument("--max-steps", type=int, help="Smoke test only; partial sidecars are rejected by training")
     p.add_argument("--preview-only", action="store_true", help="Save box overlays without loading the model")
-    p.add_argument("--overwrite", action="store_true")
+    existing = p.add_mutually_exclusive_group()
+    existing.add_argument("--overwrite", action="store_true")
+    existing.add_argument("--reuse-existing", action="store_true", help="Validate and skip complete matching sidecars")
+    p.add_argument("--dry-run", action="store_true", help="Validate references/targets and list work; no model, media decode or writes")
     p.add_argument("--no-progress", action="store_true")
     args = p.parse_args(argv)
+    if args.manifest is not None and args.dataset != "ego4d":
+        p.error("--manifest is only supported with --dataset ego4d")
+    if args.manifest is None and (args.split != "all" or args.start_index != 1
+                                  or args.end_index is not None or args.limit is not None):
+        p.error("--split/--start-index/--end-index/--limit require --manifest")
+    if args.start_index < 1 or (args.end_index is not None and args.end_index < args.start_index):
+        p.error("Invalid manifest index range")
+    if args.limit is not None and args.limit <= 0:
+        p.error("--limit must be positive")
+    if args.reuse_existing and (args.max_steps is not None or args.preview_only):
+        p.error("--reuse-existing is for full extraction, not smoke tests or previews")
+    if args.dry_run and args.preview_only:
+        p.error("--dry-run and --preview-only are separate checks")
     if args.box_coordinate_space is None:
         args.box_coordinate_space = "encoded" if args.dataset == "easycom" else "camera"
     for name in ("video_fps", "audio_window_seconds", "visual_window_seconds", "crop_expansion"):
@@ -389,30 +460,60 @@ def parse_args(argv=None):
     return args
 
 
-def main() -> None:
-    args = parse_args()
-    if not args.preview_only:
+def main(argv=None) -> None:
+    args = parse_args(argv)
+    version = None
+    if not args.preview_only and not args.dry_run:
         import transformers
 
         # Fail before indexing/decoding source media, not after expensive I/O.
-        validate_local_transformers_version(transformers.__version__)
-    names = [normalize_sample_name(value, args.dataset) for value in args.samples]
-    if len(set(names)) != len(names):
-        raise ValueError("Duplicate samples")
-    for name in names:
-        for path in (args.feature_root / f"{name}.pt", args.target_root / f"{name}.npy"):
-            if not path.is_file():
-                raise FileNotFoundError(path)
-        if (args.output_root / f"{name}.pt").exists() and not args.overwrite and not args.preview_only:
-            raise FileExistsError(f"{name}.pt exists; use a new output root or --overwrite")
-    ego_boxes = {}
+        version = validate_local_transformers_version(transformers.__version__)
+    names = select_samples(args)
+    paths = []
     if args.dataset == "ego4d":
         paths = args.ego4d_annotations or [args.data_root / "v1" / "annotations" / f"av_{split}.json"
                                           for split in ("train", "val")]
-        ego_boxes = load_ego4d_boxes(paths, names)
+        for path in paths:
+            if not path.is_file():
+                raise FileNotFoundError(path)
+    for name in names:
+        required = [args.feature_root / f"{name}.pt", args.target_root / f"{name}.npy"]
+        if args.dataset == "ego4d":
+            required.append(args.data_root / "v1" / "clips" / f"{name}.mp4")
+        for path in required:
+            if not path.is_file():
+                raise FileNotFoundError(path)
+        if ((args.output_root / f"{name}.pt").exists() and not args.overwrite
+                and not args.preview_only and not args.reuse_existing):
+            raise FileExistsError(f"{name}.pt exists; use --reuse-existing, a new output root or --overwrite")
+    pending, contract = [], None
+    # Stream through references on CPU; never keep all global states in memory.
+    for name in tqdm(names, desc="local preflight", unit="clip", disable=args.no_progress or args.dry_run):
+        reference = torch.load(args.feature_root / f"{name}.pt", map_location="cpu", weights_only=False)
+        times, _ = reference_alignment(reference)
+        target = np.load(args.target_root / f"{name}.npy", allow_pickle=False, mmap_mode="r")
+        if target.shape != (len(times), 3):
+            raise ValueError(f"Reference/target mismatch: {name}, {target.shape} vs {len(times)} states")
+        if args.reuse_existing and (args.output_root / f"{name}.pt").is_file():
+            contract = validate_reusable_sidecar(args, name, reference, version, contract)
+            print(f"{name}: reuse ({len(times)} states)", flush=True)
+        else:
+            pending.append(name)
+            if args.dry_run:
+                steps = min(len(times), args.max_steps or len(times))
+                print(f"{name}: extract ({steps} states, {steps / 5:g}s)", flush=True)
+        del reference, target
+    print(f"Selected {len(names)} clips; reuse={len(names) - len(pending)}, pending={len(pending)}", flush=True)
+    if args.dry_run:
+        print("Dry run only: annotation contents, pixel alignment and model/runtime are not checked.")
+        return
+    if not pending:
+        return
+    ego_boxes = load_ego4d_boxes(paths, pending) if args.dataset == "ego4d" else {}
     args.output_root.mkdir(parents=True, exist_ok=True)
     encoder = None
-    for name in names:
+    for position, name in enumerate(pending, 1):
+        print(f"[{position}/{len(pending)}] {name}", flush=True)
         path = args.feature_root / f"{name}.pt"
         reference = torch.load(path, map_location="cpu", weights_only=False)
         times, indices = reference_alignment(reference)
@@ -430,30 +531,26 @@ def main() -> None:
                   "annotations use --box-coordinate-space encoded; otherwise boxes shift right.", flush=True)
         if args.preview_only:
             preview_boxes(media, annotations, args, name, len(times) / 5)
+            ego_boxes.pop(name, None)
+            del media, annotations, reference
             continue
         if encoder is None:
             encoder = FrozenQwenLocalEncoder(args)
+        config = {**requested_local_config(args),
+                  "model_config_sha256": encoder.model_config_sha256,
+                  "transformers": encoder.transformers_version,
+                  "visual_dim": encoder.visual_dim, "audio_dim": encoder.audio_dim}
+        new_contract = {"format": LOCAL_FORMAT, "modalities": "av", "config": config}
+        if contract is not None and contract != new_contract:
+            raise ValueError("Current encoder differs from reused sidecars; use a new output root")
+        contract = new_contract
         print(f"{name}: {len(times)} states; native visual={encoder.visual_dim}D audio={encoder.audio_dim}D", flush=True)
         output = extract_local(media, annotations, encoder, times, args)
         output.update({
             "format": LOCAL_FORMAT, "dataset": args.dataset, "sample_id": name,
             "times": times, "label_indices": indices, "partial": len(times) != full_steps,
             "reference_features": str(path.resolve()), "direct_media": media.metadata(),
-            "config": {
-                "encoder_id": args.encoder_id or args.model_id,
-                "model_config_sha256": encoder.model_config_sha256,
-                "transformers": encoder.transformers_version, "dtype": args.dtype,
-                "attention": "flash_attention_2" if args.flash_attention_2 else "default",
-                "visual_dim": encoder.visual_dim, "audio_dim": encoder.audio_dim,
-                "visual_fps": args.video_fps, "visual_window_seconds": args.visual_window_seconds,
-                "audio_window_seconds": args.audio_window_seconds,
-                "visual_pooling": "spatial-mean_then_equal-person-mean_then_equal-valid-frame-mean",
-                "audio_pooling": "valid-output-token-mean", "audio_source": "video-mixed-mono-16khz",
-                "box_kind": args.box_kind if args.dataset == "easycom" else "ego4d-provided-tracks",
-                "box_coordinate_space": args.box_coordinate_space, "crop_expansion": args.crop_expansion,
-                "max_pixels": args.max_pixels, "uses_speech_labels": False,
-                "causal_policy": "half-open-trailing-window", "box_protocol": "provided-boxes-not-online-detector",
-            },
+            "config": config,
         })
         destination = args.output_root / f"{name}.pt"
         # Atomic replacement: a failed extraction never looks like a ready file.
@@ -466,6 +563,8 @@ def main() -> None:
             temporary.unlink(missing_ok=True)
         print(f"Saved {destination}; V-valid={output['visual_valid'].float().mean():.1%}, "
               f"A-valid={output['audio_valid'].float().mean():.1%}, partial={output['partial']}", flush=True)
+        ego_boxes.pop(name, None)
+        del media, annotations, reference, output
 
 
 if __name__ == "__main__":