wx44wx/icassp2027-experiment
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add Ego4D local AV ablation workflow

Browse Source
This commit is contained in:
Xin Wang
2026-08-28 00:25:54 +08:00
parent 8b2fbd2622
commit b5ce64702a
8 changed files with 731 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files

171
extract_turn_local_features.py
View File

@@ -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__":