Add scripts for preparing and extracting 5 FPS audio-video clips, including prepare_ego4d_2fps.py, run_extract_5fps_2pergpu.sh, and run_prepare_all_5fps.sh. Introduce streaming_turn_model_directly.py for direct extraction from original media. Update README.md to reflect changes in Qwen's streaming design and causal extraction modes. Add new AGENTS.md file for documentation.
This commit is contained in:
Xin Wang
Executable
+421
@@ -0,0 +1,421 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Prepare EgoSpeak's Ego4D split as exact 2 FPS audio-video clips.
|
||||
|
||||
The 5 Hz target arrays are the timing authority. For a target with T rows,
|
||||
the prepared clip has duration T / 5 seconds and exactly ceil(2T / 5) video
|
||||
frames. Short source audio is padded with silence and all audio is trimmed to
|
||||
the target duration.
|
||||
|
||||
Processing is intentionally sequential. A single libx264 process uses all
|
||||
available CPU cores on the intended extraction machine, so launching several
|
||||
encoders normally creates CPU and disk contention instead of improving total
|
||||
throughput.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import csv
|
||||
import json
|
||||
import math
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
TARGET_FPS = 5
|
||||
VIDEO_FPS = 2
|
||||
EXPECTED_TRAIN_CLIPS = 346
|
||||
EXPECTED_TEST_CLIPS = 87
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ClipSpec:
|
||||
split: str
|
||||
uid: str
|
||||
steps: int
|
||||
duration: float
|
||||
video_frames: int
|
||||
source: Path
|
||||
target: Path
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description=(
|
||||
"Prepare the 433 clips in EgoSpeak's EGO4DSHUFFLE split as "
|
||||
"strict fixed-grid 2 FPS MP4 files."
|
||||
)
|
||||
)
|
||||
parser.add_argument(
|
||||
"--data-root",
|
||||
type=Path,
|
||||
default=Path("/data/wangx/Datasets/ego4d"),
|
||||
help="Ego4D dataset root (default: %(default)s)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--split-json",
|
||||
type=Path,
|
||||
help="EgoSpeak LSTR data_info.json; defaults to DATA_ROOT/egospeak_data_info.json",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--video-root",
|
||||
type=Path,
|
||||
help="Raw clip directory; defaults to DATA_ROOT/v1/clips",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--target-root",
|
||||
type=Path,
|
||||
help="Target directory; defaults to DATA_ROOT/ego4d_target_frame",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-root",
|
||||
type=Path,
|
||||
help="Prepared video directory; defaults to DATA_ROOT/omni_prepared_2fps",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--manifest",
|
||||
type=Path,
|
||||
help="Output TSV; defaults to DATA_ROOT/ego4d_egospeak.tsv",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--split",
|
||||
choices=("all", "train", "test"),
|
||||
default="all",
|
||||
help="Prepare all, train, or test clips (default: %(default)s)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--uid",
|
||||
action="append",
|
||||
help="Prepare only this split UID; repeat to select several clips",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--limit",
|
||||
type=int,
|
||||
help="Prepare only the first N selected clips, for testing",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--crf",
|
||||
type=int,
|
||||
default=18,
|
||||
help="libx264 constant-rate-factor value (default: %(default)s)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--preset",
|
||||
default="medium",
|
||||
help="libx264 preset (default: %(default)s)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--force",
|
||||
action="store_true",
|
||||
help="Rebuild outputs even when their frame counts are already correct",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dry-run",
|
||||
action="store_true",
|
||||
help="Validate inputs and print work without running ffmpeg",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.limit is not None and args.limit <= 0:
|
||||
parser.error("--limit must be positive")
|
||||
if not 0 <= args.crf <= 51:
|
||||
parser.error("--crf must be between 0 and 51")
|
||||
return args
|
||||
|
||||
|
||||
def resolve_paths(args: argparse.Namespace) -> None:
|
||||
args.data_root = args.data_root.expanduser().resolve()
|
||||
args.split_json = (
|
||||
args.split_json or args.data_root / "egospeak_data_info.json"
|
||||
).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_prepared_2fps"
|
||||
).expanduser().resolve()
|
||||
args.manifest = (
|
||||
args.manifest or args.data_root / "ego4d_egospeak.tsv"
|
||||
).expanduser().resolve()
|
||||
|
||||
|
||||
def require_file(path: Path, description: str) -> None:
|
||||
if not path.is_file():
|
||||
raise FileNotFoundError(f"Missing {description}: {path}")
|
||||
|
||||
|
||||
def load_specs(args: argparse.Namespace) -> list[ClipSpec]:
|
||||
require_file(args.split_json, "EgoSpeak split JSON")
|
||||
if not args.video_root.is_dir():
|
||||
raise FileNotFoundError(f"Missing raw-video directory: {args.video_root}")
|
||||
if not args.target_root.is_dir():
|
||||
raise FileNotFoundError(f"Missing target directory: {args.target_root}")
|
||||
|
||||
with args.split_json.open() as file:
|
||||
payload = json.load(file)
|
||||
try:
|
||||
split_info = payload["EGO4DSHUFFLE"]
|
||||
train_uids = split_info["train_session_set"]
|
||||
test_uids = split_info["test_session_set"]
|
||||
except KeyError as error:
|
||||
raise RuntimeError(
|
||||
f"{args.split_json} is not EgoSpeak's LSTR data_info.json: "
|
||||
f"missing {error}"
|
||||
) from error
|
||||
|
||||
if len(train_uids) != EXPECTED_TRAIN_CLIPS:
|
||||
raise RuntimeError(
|
||||
f"Expected {EXPECTED_TRAIN_CLIPS} EGO4DSHUFFLE train clips, "
|
||||
f"found {len(train_uids)}"
|
||||
)
|
||||
if len(test_uids) != EXPECTED_TEST_CLIPS:
|
||||
raise RuntimeError(
|
||||
f"Expected {EXPECTED_TEST_CLIPS} EGO4DSHUFFLE test clips, "
|
||||
f"found {len(test_uids)}"
|
||||
)
|
||||
|
||||
specs: list[ClipSpec] = []
|
||||
seen: set[str] = set()
|
||||
for split, uids in (("train", train_uids), ("test", test_uids)):
|
||||
for uid in uids:
|
||||
if uid in seen:
|
||||
raise RuntimeError(f"Duplicate UID in EGO4DSHUFFLE: {uid}")
|
||||
seen.add(uid)
|
||||
|
||||
source = args.video_root / f"{uid}.mp4"
|
||||
target_path = args.target_root / f"{uid}.npy"
|
||||
require_file(source, "source video")
|
||||
require_file(target_path, "target array")
|
||||
|
||||
target = np.load(target_path, mmap_mode="r", allow_pickle=False)
|
||||
if target.ndim != 2 or target.shape[1] != 3:
|
||||
raise RuntimeError(
|
||||
f"Expected target shape [T, 3], found {target.shape}: "
|
||||
f"{target_path}"
|
||||
)
|
||||
|
||||
steps = int(target.shape[0])
|
||||
duration = steps / TARGET_FPS
|
||||
video_frames = math.ceil(steps * VIDEO_FPS / TARGET_FPS)
|
||||
specs.append(
|
||||
ClipSpec(
|
||||
split=split,
|
||||
uid=uid,
|
||||
steps=steps,
|
||||
duration=duration,
|
||||
video_frames=video_frames,
|
||||
source=source,
|
||||
target=target_path,
|
||||
)
|
||||
)
|
||||
return specs
|
||||
|
||||
|
||||
def write_manifest(specs: list[ClipSpec], path: Path) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
temporary = path.with_suffix(path.suffix + ".partial")
|
||||
with temporary.open("w", newline="") as file:
|
||||
writer = csv.writer(file, delimiter="\t")
|
||||
writer.writerow(
|
||||
["split", "uid", "steps", "duration", "video_frames_2fps"]
|
||||
)
|
||||
for spec in specs:
|
||||
writer.writerow(
|
||||
[
|
||||
spec.split,
|
||||
spec.uid,
|
||||
spec.steps,
|
||||
f"{spec.duration:.6f}",
|
||||
spec.video_frames,
|
||||
]
|
||||
)
|
||||
temporary.replace(path)
|
||||
|
||||
|
||||
def select_specs(args: argparse.Namespace, specs: list[ClipSpec]) -> list[ClipSpec]:
|
||||
selected = specs
|
||||
if args.split != "all":
|
||||
selected = [spec for spec in selected if spec.split == args.split]
|
||||
if args.uid:
|
||||
requested = set(args.uid)
|
||||
known = {spec.uid for spec in specs}
|
||||
missing = sorted(requested - known)
|
||||
if missing:
|
||||
raise RuntimeError(
|
||||
"Requested UID(s) are not in EGO4DSHUFFLE: " + ", ".join(missing)
|
||||
)
|
||||
selected = [spec for spec in selected if spec.uid in requested]
|
||||
if args.limit is not None:
|
||||
selected = selected[: args.limit]
|
||||
if not selected:
|
||||
raise RuntimeError("No clips selected")
|
||||
return selected
|
||||
|
||||
|
||||
def count_video_frames(path: Path) -> int | None:
|
||||
if not path.is_file() or path.stat().st_size == 0:
|
||||
return None
|
||||
command = [
|
||||
"ffprobe",
|
||||
"-v",
|
||||
"error",
|
||||
"-select_streams",
|
||||
"v:0",
|
||||
"-count_frames",
|
||||
"-show_entries",
|
||||
"stream=nb_read_frames",
|
||||
"-of",
|
||||
"default=noprint_wrappers=1:nokey=1",
|
||||
str(path),
|
||||
]
|
||||
completed = subprocess.run(command, text=True, capture_output=True)
|
||||
if completed.returncode != 0:
|
||||
return None
|
||||
try:
|
||||
return int(completed.stdout.strip())
|
||||
except ValueError:
|
||||
return None
|
||||
|
||||
|
||||
def prepare_one(
|
||||
args: argparse.Namespace,
|
||||
spec: ClipSpec,
|
||||
index: int,
|
||||
total: int,
|
||||
) -> str:
|
||||
output = args.output_root / f"{spec.uid}_2fps.mp4"
|
||||
actual = count_video_frames(output)
|
||||
prefix = f"[{index}/{total}] [{spec.split}] {spec.uid}"
|
||||
|
||||
if not args.force and actual == spec.video_frames:
|
||||
print(f"{prefix}: reuse ({actual} frames)", flush=True)
|
||||
return "reused"
|
||||
|
||||
if actual is None:
|
||||
reason = "missing"
|
||||
else:
|
||||
reason = f"actual={actual}, expected={spec.video_frames}"
|
||||
print(
|
||||
f"{prefix}: prepare ({reason}; {spec.steps} states, "
|
||||
f"{spec.duration:.1f}s, {spec.video_frames} frames)",
|
||||
flush=True,
|
||||
)
|
||||
if args.dry_run:
|
||||
return "dry-run"
|
||||
|
||||
temporary = args.output_root / f".{spec.uid}.partial.mp4"
|
||||
duration = f"{spec.duration:.6f}"
|
||||
video_filter = (
|
||||
f"[0:v:0]trim=start=0:end={duration},setpts=PTS-STARTPTS,"
|
||||
f"fps=fps={VIDEO_FPS}:start_time=0:round=down,"
|
||||
"tpad=stop_mode=clone:stop=2,"
|
||||
f"trim=start_frame=0:end_frame={spec.video_frames},"
|
||||
f"setpts=N/({VIDEO_FPS}*TB),format=yuv420p[vout]"
|
||||
)
|
||||
audio_filter = (
|
||||
f"[0:a:0]apad,atrim=start=0:end={duration},"
|
||||
"asetpts=PTS-STARTPTS[aout]"
|
||||
)
|
||||
command = [
|
||||
"ffmpeg",
|
||||
"-hide_banner",
|
||||
"-loglevel",
|
||||
"warning",
|
||||
"-stats",
|
||||
"-y",
|
||||
"-i",
|
||||
str(spec.source),
|
||||
"-filter_complex",
|
||||
f"{video_filter};{audio_filter}",
|
||||
"-map",
|
||||
"[vout]",
|
||||
"-map",
|
||||
"[aout]",
|
||||
"-map_metadata",
|
||||
"-1",
|
||||
"-c:v",
|
||||
"libx264",
|
||||
"-preset",
|
||||
args.preset,
|
||||
"-crf",
|
||||
str(args.crf),
|
||||
"-c:a",
|
||||
"aac",
|
||||
"-ar",
|
||||
"16000",
|
||||
"-ac",
|
||||
"1",
|
||||
"-b:a",
|
||||
"64k",
|
||||
"-movflags",
|
||||
"+faststart",
|
||||
str(temporary),
|
||||
]
|
||||
completed = subprocess.run(command)
|
||||
if completed.returncode != 0:
|
||||
raise RuntimeError(
|
||||
f"ffmpeg failed for {spec.uid} with exit code {completed.returncode}; "
|
||||
f"partial output was left at {temporary}"
|
||||
)
|
||||
|
||||
actual = count_video_frames(temporary)
|
||||
if actual != spec.video_frames:
|
||||
raise RuntimeError(
|
||||
f"Prepared frame count is {actual}, expected {spec.video_frames}: "
|
||||
f"{temporary}"
|
||||
)
|
||||
temporary.replace(output)
|
||||
print(f"{prefix}: OK ({actual} frames) -> {output}", flush=True)
|
||||
return "prepared"
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
resolve_paths(args)
|
||||
for executable in ("ffmpeg", "ffprobe"):
|
||||
if shutil.which(executable) is None:
|
||||
raise RuntimeError(f"Required executable is not on PATH: {executable}")
|
||||
|
||||
specs = load_specs(args)
|
||||
write_manifest(specs, args.manifest)
|
||||
selected = select_specs(args, specs)
|
||||
args.output_root.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
total_steps = sum(spec.steps for spec in specs)
|
||||
print(
|
||||
f"Validated EGO4DSHUFFLE: {len(specs)} clips "
|
||||
f"({EXPECTED_TRAIN_CLIPS} train, {EXPECTED_TEST_CLIPS} test), "
|
||||
f"{total_steps} target states"
|
||||
)
|
||||
print(f"Manifest: {args.manifest}")
|
||||
print(f"Output: {args.output_root}")
|
||||
print(f"Selected: {len(selected)} clips; sequential ffmpeg execution")
|
||||
|
||||
counts = {"prepared": 0, "reused": 0, "dry-run": 0}
|
||||
for index, spec in enumerate(selected, start=1):
|
||||
result = prepare_one(args, spec, index, len(selected))
|
||||
counts[result] += 1
|
||||
|
||||
print(
|
||||
"Complete: "
|
||||
+ ", ".join(f"{key}={value}" for key, value in counts.items() if value)
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
try:
|
||||
main()
|
||||
except KeyboardInterrupt:
|
||||
print("\nInterrupted; rerun the same command to resume.", file=sys.stderr)
|
||||
raise SystemExit(130)
|
||||
except Exception as error:
|
||||
print(f"ERROR: {error}", file=sys.stderr)
|
||||
raise SystemExit(1)
|
||||
Reference in New Issue
Block a user