422 lines
13 KiB
Python
422 lines
13 KiB
Python
|
#!/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)
|