Add extract_ego4d_batch.py for efficient EgoSpeak clip extraction. This script loads the Qwen model once and processes selected manifest rows sequentially, ensuring each clip is handled independently. Update README with usage instructions and details on the extraction process, including handling of completed outputs and memory management.
This commit is contained in:
Xin Wang
@@ -0,0 +1,454 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Extract EgoSpeak's Ego4D clips while loading Qwen only once.
|
||||
|
||||
Manifest sample indices are one-based and refer to data rows, excluding the
|
||||
header. For example, ``--start-index 31 --end-index 60`` selects exactly the
|
||||
same clips as ``sed -n '32,61p' ego4d_egospeak.tsv``.
|
||||
|
||||
Only immutable model weights and the processor are reused. Every clip invokes
|
||||
the direct extractor independently, which creates a new media reader, rolling
|
||||
KV cache, TMRoPE origin, and readout sequence. Consequently no temporal state
|
||||
is carried between dataset samples.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import csv
|
||||
import gc
|
||||
import math
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
import traceback
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Sequence
|
||||
|
||||
import torch
|
||||
|
||||
import streaming_turn_model as core
|
||||
import streaming_turn_model_directly as direct
|
||||
|
||||
|
||||
TARGET_HZ = 5
|
||||
DEFAULT_OUTPUT_NAME = "open_stream_2fps_20s_decord"
|
||||
EXPECTED_MANIFEST_COLUMNS = (
|
||||
"split",
|
||||
"uid",
|
||||
"steps",
|
||||
"duration",
|
||||
"video_frames_2fps",
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ClipRecord:
|
||||
index: int
|
||||
split: str
|
||||
uid: str
|
||||
steps: int
|
||||
duration: float
|
||||
video_frames_2fps: int
|
||||
|
||||
|
||||
def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument(
|
||||
"--data-root",
|
||||
type=Path,
|
||||
default=Path("/data/wangx/Datasets/ego4d"),
|
||||
help="Ego4D experiment root (default: %(default)s)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--manifest",
|
||||
type=Path,
|
||||
help="Defaults to DATA_ROOT/ego4d_egospeak.tsv",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--video-root",
|
||||
type=Path,
|
||||
help="Defaults to DATA_ROOT/v1/clips",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--target-root",
|
||||
type=Path,
|
||||
help="Defaults to DATA_ROOT/ego4d_target_frame",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-root",
|
||||
type=Path,
|
||||
help=(
|
||||
"Defaults to DATA_ROOT/omni_features/"
|
||||
f"{DEFAULT_OUTPUT_NAME}"
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model-id",
|
||||
default="/data/wangx/Models/Qwen2.5-Omni-3B",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--start-index",
|
||||
type=int,
|
||||
default=1,
|
||||
help="First one-based manifest data row to process (default: 1)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--end-index",
|
||||
type=int,
|
||||
help="Last one-based manifest data row to process, inclusive",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--split",
|
||||
choices=("all", "train", "test"),
|
||||
default="all",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--overwrite",
|
||||
action="store_true",
|
||||
help="Replace completed outputs instead of resuming around them",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--continue-on-error",
|
||||
action="store_true",
|
||||
help="Record a failed clip and continue with the next selected clip",
|
||||
)
|
||||
parser.add_argument("--memory-seconds", type=float, default=20.0)
|
||||
parser.add_argument("--step-seconds", type=float, default=0.2)
|
||||
parser.add_argument("--native-block-seconds", type=float, default=2.0)
|
||||
parser.add_argument("--video-fps", type=float, default=2.0)
|
||||
parser.add_argument("--video-min-pixels", type=int, default=64 * 28 * 28)
|
||||
parser.add_argument("--video-max-pixels", type=int, default=128 * 28 * 28)
|
||||
parser.add_argument("--frame-cache-size", type=int, default=64)
|
||||
parser.add_argument(
|
||||
"--decoder",
|
||||
choices=("decord", "pynvcodec"),
|
||||
default="decord",
|
||||
)
|
||||
parser.add_argument("--decoder-gpu-id", type=int, default=0)
|
||||
parser.add_argument("--decoder-threads", type=int, default=1)
|
||||
parser.add_argument("--max-open-readers", type=int, default=1)
|
||||
parser.add_argument("--verify-every", type=int, default=0)
|
||||
parser.add_argument("--verify-continuous-every", type=int, default=0)
|
||||
parser.add_argument("--no-progress", action="store_true")
|
||||
parser.add_argument(
|
||||
"--dtype",
|
||||
choices=("auto", "bfloat16", "float16", "float32"),
|
||||
default="bfloat16",
|
||||
)
|
||||
parser.add_argument("--flash-attention-2", action="store_true")
|
||||
parser.add_argument("--device", default="cuda:0")
|
||||
parser.add_argument(
|
||||
"--empty-cache-every",
|
||||
type=int,
|
||||
default=0,
|
||||
help="Call torch.cuda.empty_cache every N processed clips; 0 disables",
|
||||
)
|
||||
args = parser.parse_args(argv)
|
||||
|
||||
if args.start_index <= 0:
|
||||
parser.error("--start-index must be positive")
|
||||
if args.end_index is not None and args.end_index < args.start_index:
|
||||
parser.error("--end-index must be >= --start-index")
|
||||
if args.empty_cache_every < 0:
|
||||
parser.error("--empty-cache-every must be >= 0")
|
||||
if args.step_seconds <= 0:
|
||||
parser.error("--step-seconds must be positive")
|
||||
return args
|
||||
|
||||
|
||||
def resolve_paths(args: argparse.Namespace) -> None:
|
||||
args.data_root = args.data_root.expanduser().resolve()
|
||||
args.manifest = (
|
||||
args.manifest or args.data_root / "ego4d_egospeak.tsv"
|
||||
).expanduser().resolve()
|
||||
args.video_root = (
|
||||
args.video_root or args.data_root / "v1" / "clips"
|
||||
).expanduser().resolve()
|
||||
args.target_root = (
|
||||
args.target_root or args.data_root / "ego4d_target_frame"
|
||||
).expanduser().resolve()
|
||||
args.output_root = (
|
||||
args.output_root
|
||||
or args.data_root / "omni_features" / DEFAULT_OUTPUT_NAME
|
||||
).expanduser().resolve()
|
||||
|
||||
|
||||
def load_manifest(path: Path) -> list[ClipRecord]:
|
||||
if not path.is_file():
|
||||
raise FileNotFoundError(f"Missing EgoSpeak manifest: {path}")
|
||||
records: list[ClipRecord] = []
|
||||
seen: set[str] = set()
|
||||
with path.open(newline="") as file:
|
||||
reader = csv.DictReader(file, delimiter="\t")
|
||||
if reader.fieldnames != list(EXPECTED_MANIFEST_COLUMNS):
|
||||
raise RuntimeError(
|
||||
f"Unexpected manifest columns in {path}: {reader.fieldnames}"
|
||||
)
|
||||
for index, row in enumerate(reader, start=1):
|
||||
split = row["split"].strip()
|
||||
uid = row["uid"].strip()
|
||||
if split not in {"train", "test"}:
|
||||
raise RuntimeError(f"Invalid split at manifest row {index}: {split}")
|
||||
if not uid or uid in seen:
|
||||
raise RuntimeError(
|
||||
f"Missing or duplicate UID at manifest row {index}: {uid!r}"
|
||||
)
|
||||
seen.add(uid)
|
||||
steps = int(row["steps"])
|
||||
duration = float(row["duration"])
|
||||
video_frames = int(row["video_frames_2fps"])
|
||||
if steps <= 0:
|
||||
raise RuntimeError(f"Non-positive steps for {uid}: {steps}")
|
||||
expected_duration = steps / TARGET_HZ
|
||||
if not math.isclose(
|
||||
duration, expected_duration, rel_tol=0.0, abs_tol=1e-6
|
||||
):
|
||||
raise RuntimeError(
|
||||
f"Manifest duration mismatch for {uid}: {duration} vs "
|
||||
f"{expected_duration} from {steps} target steps"
|
||||
)
|
||||
records.append(
|
||||
ClipRecord(
|
||||
index=index,
|
||||
split=split,
|
||||
uid=uid,
|
||||
steps=steps,
|
||||
duration=duration,
|
||||
video_frames_2fps=video_frames,
|
||||
)
|
||||
)
|
||||
if not records:
|
||||
raise RuntimeError(f"No clips in {path}")
|
||||
return records
|
||||
|
||||
|
||||
def select_records(
|
||||
records: list[ClipRecord], args: argparse.Namespace
|
||||
) -> list[ClipRecord]:
|
||||
end_index = args.end_index if args.end_index is not None else len(records)
|
||||
selected = [
|
||||
record
|
||||
for record in records
|
||||
if args.start_index <= record.index <= end_index
|
||||
and (args.split == "all" or record.split == args.split)
|
||||
]
|
||||
if not selected:
|
||||
raise RuntimeError(
|
||||
f"No manifest rows selected by {args.start_index=} "
|
||||
f"{args.end_index=} {args.split=}"
|
||||
)
|
||||
return selected
|
||||
|
||||
|
||||
def preflight(selected: list[ClipRecord], args: argparse.Namespace) -> None:
|
||||
if not args.video_root.is_dir():
|
||||
raise FileNotFoundError(f"Missing video root: {args.video_root}")
|
||||
if not args.target_root.is_dir():
|
||||
raise FileNotFoundError(f"Missing target root: {args.target_root}")
|
||||
model_path = Path(args.model_id).expanduser()
|
||||
if model_path.is_absolute() and not model_path.is_dir():
|
||||
raise FileNotFoundError(f"Missing local model: {model_path}")
|
||||
|
||||
problems: list[str] = []
|
||||
for record in selected:
|
||||
video = args.video_root / f"{record.uid}.mp4"
|
||||
target = args.target_root / f"{record.uid}.npy"
|
||||
if not video.is_file() or video.stat().st_size == 0:
|
||||
problems.append(f"missing video: {video}")
|
||||
if not target.is_file() or target.stat().st_size == 0:
|
||||
problems.append(f"missing target: {target}")
|
||||
continue
|
||||
try:
|
||||
target_steps = core.target_num_steps(target)
|
||||
except Exception as error: # report all bad inputs before model loading
|
||||
problems.append(f"invalid target {target}: {error}")
|
||||
continue
|
||||
if target_steps != record.steps:
|
||||
problems.append(
|
||||
f"target length mismatch {record.uid}: manifest={record.steps}, "
|
||||
f"npy={target_steps}"
|
||||
)
|
||||
if problems:
|
||||
preview = "\n".join(f"- {problem}" for problem in problems[:20])
|
||||
suffix = "" if len(problems) <= 20 else f"\n... and {len(problems)-20} more"
|
||||
raise RuntimeError(f"Ego4D preflight failed:\n{preview}{suffix}")
|
||||
|
||||
|
||||
def clip_argv(
|
||||
record: ClipRecord,
|
||||
args: argparse.Namespace,
|
||||
output: Path,
|
||||
) -> list[str]:
|
||||
values = [
|
||||
"--model-id",
|
||||
str(args.model_id),
|
||||
"--video",
|
||||
str(args.video_root / f"{record.uid}.mp4"),
|
||||
"--target",
|
||||
str(args.target_root / f"{record.uid}.npy"),
|
||||
"--mode",
|
||||
"open-stream-replay",
|
||||
"--step-seconds",
|
||||
str(args.step_seconds),
|
||||
"--native-block-seconds",
|
||||
str(args.native_block_seconds),
|
||||
"--memory-seconds",
|
||||
str(args.memory_seconds),
|
||||
"--video-fps",
|
||||
str(args.video_fps),
|
||||
"--video-min-pixels",
|
||||
str(args.video_min_pixels),
|
||||
"--video-max-pixels",
|
||||
str(args.video_max_pixels),
|
||||
"--frame-cache-size",
|
||||
str(args.frame_cache_size),
|
||||
"--decoder",
|
||||
args.decoder,
|
||||
"--decoder-gpu-id",
|
||||
str(args.decoder_gpu_id),
|
||||
"--decoder-threads",
|
||||
str(args.decoder_threads),
|
||||
"--max-open-readers",
|
||||
str(args.max_open_readers),
|
||||
"--verify-every",
|
||||
str(args.verify_every),
|
||||
"--verify-continuous-every",
|
||||
str(args.verify_continuous_every),
|
||||
"--dtype",
|
||||
args.dtype,
|
||||
"--device",
|
||||
args.device,
|
||||
"--output",
|
||||
str(output),
|
||||
]
|
||||
if args.flash_attention_2:
|
||||
values.append("--flash-attention-2")
|
||||
if args.no_progress:
|
||||
values.append("--no-progress")
|
||||
return values
|
||||
|
||||
|
||||
def format_eta(seconds: float) -> str:
|
||||
if not math.isfinite(seconds) or seconds < 0:
|
||||
return "unknown"
|
||||
hours, remainder = divmod(int(seconds), 3600)
|
||||
minutes, secs = divmod(remainder, 60)
|
||||
return f"{hours:02d}:{minutes:02d}:{secs:02d}"
|
||||
|
||||
|
||||
def main(argv: Sequence[str] | None = None) -> None:
|
||||
args = parse_args(argv)
|
||||
resolve_paths(args)
|
||||
core._configure_qwen_logging()
|
||||
|
||||
records = load_manifest(args.manifest)
|
||||
selected = select_records(records, args)
|
||||
preflight(selected, args)
|
||||
(args.output_root / "train").mkdir(parents=True, exist_ok=True)
|
||||
(args.output_root / "test").mkdir(parents=True, exist_ok=True)
|
||||
|
||||
pending = [
|
||||
record
|
||||
for record in selected
|
||||
if args.overwrite
|
||||
or not (args.output_root / record.split / f"{record.uid}.pt").is_file()
|
||||
or (args.output_root / record.split / f"{record.uid}.pt").stat().st_size == 0
|
||||
]
|
||||
skipped = len(selected) - len(pending)
|
||||
print(
|
||||
f"Selected manifest samples {selected[0].index}-{selected[-1].index}: "
|
||||
f"{len(selected)} clip(s), pending={len(pending)}, skipped={skipped}",
|
||||
flush=True,
|
||||
)
|
||||
print(f"Output root: {args.output_root}", flush=True)
|
||||
if not pending:
|
||||
print("All selected clips already have non-empty outputs.")
|
||||
return
|
||||
|
||||
print(f"Loading Qwen once from {args.model_id} ...", flush=True)
|
||||
load_started = time.perf_counter()
|
||||
model, processor = core.load_qwen_omni(
|
||||
args.model_id,
|
||||
dtype=args.dtype,
|
||||
device_map={"": args.device},
|
||||
use_flash_attention_2=args.flash_attention_2,
|
||||
)
|
||||
model.requires_grad_(False)
|
||||
print(
|
||||
f"Model loaded once in {time.perf_counter() - load_started:.1f}s; "
|
||||
"per-clip KV/media state will reset.",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
run_started = time.perf_counter()
|
||||
completed = 0
|
||||
failures: list[tuple[ClipRecord, str]] = []
|
||||
for pending_index, record in enumerate(pending, start=1):
|
||||
output = args.output_root / record.split / f"{record.uid}.pt"
|
||||
temporary = output.with_name(f"{output.name}.partial.{os.getpid()}")
|
||||
print(
|
||||
f"\n[{pending_index}/{len(pending)} | manifest {record.index}] "
|
||||
f"{record.split}/{record.uid}: {record.steps} states, "
|
||||
f"{record.duration:.1f}s",
|
||||
flush=True,
|
||||
)
|
||||
clip_started = time.perf_counter()
|
||||
try:
|
||||
direct.main(
|
||||
clip_argv(record, args, temporary),
|
||||
model=model,
|
||||
processor=processor,
|
||||
)
|
||||
os.replace(temporary, output)
|
||||
completed += 1
|
||||
elapsed = time.perf_counter() - clip_started
|
||||
average = (time.perf_counter() - run_started) / completed
|
||||
remaining = len(pending) - pending_index
|
||||
print(
|
||||
f"Completed {output} in {elapsed:.1f}s; "
|
||||
f"ETA {format_eta(average * remaining)}",
|
||||
flush=True,
|
||||
)
|
||||
except KeyboardInterrupt:
|
||||
print(f"Interrupted; partial path: {temporary}", file=sys.stderr)
|
||||
raise
|
||||
except Exception as error:
|
||||
failures.append((record, str(error)))
|
||||
print(
|
||||
f"FAILED manifest sample {record.index} ({record.uid}): {error}\n"
|
||||
f"Partial path, if created: {temporary}",
|
||||
file=sys.stderr,
|
||||
flush=True,
|
||||
)
|
||||
traceback.print_exc()
|
||||
if not args.continue_on_error:
|
||||
raise
|
||||
finally:
|
||||
# The model and processor stay referenced above. Collection here
|
||||
# releases the completed clip's reader, PIL frames, branch caches,
|
||||
# and Python cache wrappers before opening the next sample.
|
||||
gc.collect()
|
||||
processed = completed + len(failures)
|
||||
if (
|
||||
args.empty_cache_every > 0
|
||||
and processed > 0
|
||||
and processed % args.empty_cache_every == 0
|
||||
and torch.cuda.is_available()
|
||||
):
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
total = time.perf_counter() - run_started
|
||||
print(
|
||||
f"\nBatch finished in {format_eta(total)}: completed={completed}, "
|
||||
f"skipped={skipped}, failed={len(failures)}",
|
||||
flush=True,
|
||||
)
|
||||
if failures:
|
||||
for record, message in failures:
|
||||
print(f"- manifest {record.index} {record.uid}: {message}", file=sys.stderr)
|
||||
raise SystemExit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user