Add causal person fusion and MuVAP stage-3 baselines
This commit is contained in:
Xin Wang
579
extract_muvap_egospeak_features.py
Normal file
579
extract_muvap_egospeak_features.py
Normal file
@@ -0,0 +1,579 @@
|
||||
"""Extract frozen MuVAP VAP/ASD states for an EgoSpeak Stage-3 baseline.
|
||||
|
||||
This is a clean-room adaptation of the architecture described in MuVAP Sec.
|
||||
4.3/5.3. The upstream repository has not released its Stage-3 trainer, so the
|
||||
feature contract is deliberately named ``paper-reconstruction`` rather than
|
||||
claiming to reproduce unpublished code.
|
||||
|
||||
The extractor never reads speech targets. A verified global-feature sidecar
|
||||
is used only for its dense 5 Hz timestamps and source-media inventory; none of
|
||||
its Omni tensors are copied into the output. MuVAP's released role-VAP and ASD
|
||||
backbones are frozen and evaluated on a 25 Hz, strictly causal timeline.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import math
|
||||
import os
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from tqdm.auto import tqdm
|
||||
|
||||
from easycom_head import resolve_feature_path
|
||||
from extract_highrate_person_av_features import (
|
||||
FACE_SIZE,
|
||||
MODEL_FPS,
|
||||
OUTPUT_HZ,
|
||||
SAMPLE_RATE,
|
||||
SAMPLES_PER_MODEL_FRAME,
|
||||
FrozenMuVAPASDEncoder,
|
||||
_audio_for_tick_range,
|
||||
_chunked_audio_encode,
|
||||
_collect_clip_tracks,
|
||||
_model_tick_bounds,
|
||||
_release_paths,
|
||||
)
|
||||
from extract_turn_local_features import (
|
||||
EGO4D_MAX_TERMINAL_FRAME_OVERRUN,
|
||||
build_media_and_boxes,
|
||||
load_ego4d_boxes,
|
||||
select_samples,
|
||||
validate_source_timeline,
|
||||
)
|
||||
from local_turn_features import reference_alignment
|
||||
from muvap_egospeak_stage3 import MUVAP_EGOSPEAK_INPUT_FORMAT
|
||||
|
||||
def _sha256(path: Path) -> str:
|
||||
digest = hashlib.sha256()
|
||||
with path.open("rb") as stream:
|
||||
while block := stream.read(1024 * 1024):
|
||||
digest.update(block)
|
||||
return digest.hexdigest()
|
||||
|
||||
|
||||
def _verified_root(root: Path, revision: str) -> tuple[Path, str]:
|
||||
root = root.expanduser().resolve()
|
||||
if not re.fullmatch(r"[0-9a-fA-F]{7,40}", revision):
|
||||
raise ValueError("--muvap-revision must be a hexadecimal commit prefix")
|
||||
if not (root / "models" / "vap.py").is_file():
|
||||
raise FileNotFoundError(f"MuVAP source is incomplete: {root}")
|
||||
revision_file = root / ".upstream-revision"
|
||||
if revision_file.is_file():
|
||||
actual = revision_file.read_text().strip().lower()
|
||||
else:
|
||||
completed = subprocess.run(
|
||||
["git", "-C", str(root), "rev-parse", "HEAD"],
|
||||
check=False,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
actual = completed.stdout.strip().lower() if completed.returncode == 0 else ""
|
||||
if not re.fullmatch(r"[0-9a-f]{40}", actual) or not actual.startswith(
|
||||
revision.lower()
|
||||
):
|
||||
raise RuntimeError(
|
||||
f"MuVAP checkout is {actual or 'unreadable'}, expected {revision!r}"
|
||||
)
|
||||
return root, actual
|
||||
|
||||
|
||||
class FrozenMuVAPVAPEncoder:
|
||||
"""Load the published 25 Hz mono role-VAP and expose ``Z_VAP``."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
muvap_root: Path,
|
||||
revision: str,
|
||||
release: Path,
|
||||
device: str,
|
||||
dtype: str,
|
||||
) -> None:
|
||||
root, self.source_revision = _verified_root(muvap_root, revision)
|
||||
self.weights_path, config_path = _release_paths(release)
|
||||
try:
|
||||
import yaml
|
||||
except ImportError as error:
|
||||
raise RuntimeError("PyYAML is required to read MuVAP releases") from error
|
||||
|
||||
sys.path.insert(0, str(root))
|
||||
try:
|
||||
from models.release import load_weights
|
||||
from models.vap import build_vap
|
||||
except Exception as error:
|
||||
raise RuntimeError(f"Cannot import MuVAP from {root}") from error
|
||||
imported = Path(sys.modules[build_vap.__module__].__file__).resolve()
|
||||
if root not in imported.parents:
|
||||
raise RuntimeError(
|
||||
f"Python resolved {build_vap.__module__} from {imported}, not {root}"
|
||||
)
|
||||
payload = yaml.safe_load(config_path.read_text()) or {}
|
||||
config = payload.get("vap")
|
||||
if not isinstance(config, dict):
|
||||
raise ValueError(f"MuVAP release has no `vap` config: {config_path}")
|
||||
projection = config.get("projection_window", {})
|
||||
audio_spec = config.get("audio_encoder", "cpc")
|
||||
audio_name = audio_spec if isinstance(audio_spec, str) else audio_spec.get("name")
|
||||
if audio_name != "cpc":
|
||||
raise ValueError("The first baseline requires the CPC VAP release")
|
||||
if config.get("arch", "mono") != "mono":
|
||||
raise ValueError("EgoSpeak mixed audio requires MuVAP's mono role-VAP")
|
||||
if projection.get("mode") != "role_relative":
|
||||
raise ValueError("Expected a role-relative VAP release")
|
||||
if float(projection.get("frame_hz", -1)) != MODEL_FPS:
|
||||
raise ValueError(
|
||||
f"Expected a {MODEL_FPS:g} Hz VAP release, got "
|
||||
f"{projection.get('frame_hz')!r}"
|
||||
)
|
||||
if config.get("temporal", {}).get("causal", True) is not True:
|
||||
raise ValueError("VAP temporal attention must be causal")
|
||||
|
||||
self.device = torch.device(device)
|
||||
if self.device.type == "cuda" and not torch.cuda.is_available():
|
||||
raise RuntimeError(f"CUDA requested but unavailable: {device}")
|
||||
self.autocast_dtype = {
|
||||
"float32": None,
|
||||
"float16": torch.float16,
|
||||
"bfloat16": torch.bfloat16,
|
||||
}[dtype]
|
||||
try:
|
||||
model = build_vap(config)
|
||||
except FileNotFoundError as error:
|
||||
raise RuntimeError(
|
||||
"MuVAP's CPC frontend is missing; run "
|
||||
f"`python {root / 'tools' / 'fetch_frontends.py'} cpc`"
|
||||
) from error
|
||||
self.model = load_weights(model, self.weights_path).to(self.device).eval()
|
||||
self.model.requires_grad_(False)
|
||||
self.feature_dim = int(config["temporal"]["d_model"])
|
||||
self.release_config = config
|
||||
self.release_sha256 = _sha256(self.weights_path)
|
||||
|
||||
@torch.inference_mode()
|
||||
def encode_audio(self, waveform: np.ndarray, num_frames: int) -> torch.Tensor:
|
||||
expected = num_frames * SAMPLES_PER_MODEL_FRAME
|
||||
if waveform.dtype != np.float32 or waveform.shape != (expected,):
|
||||
raise ValueError(f"Audio must be float32 [{expected}], got {waveform.shape}")
|
||||
audio = torch.from_numpy(np.ascontiguousarray(waveform))[None, None]
|
||||
audio = audio.to(self.device, dtype=torch.float32)
|
||||
with torch.autocast(
|
||||
device_type=self.device.type,
|
||||
dtype=self.autocast_dtype,
|
||||
enabled=self.autocast_dtype is not None,
|
||||
):
|
||||
_, embedding = self.model(audio, return_embeddings=True)
|
||||
result = embedding[0].float().cpu()
|
||||
if result.shape != (num_frames, self.feature_dim):
|
||||
raise RuntimeError(f"Unexpected VAP shape: {tuple(result.shape)}")
|
||||
if not torch.isfinite(result).all():
|
||||
raise RuntimeError("MuVAP VAP produced NaN/Inf")
|
||||
return result
|
||||
|
||||
|
||||
def _chunked_vap_encode(
|
||||
encoder: FrozenMuVAPVAPEncoder,
|
||||
waveform: np.ndarray,
|
||||
*,
|
||||
num_frames: int,
|
||||
output_frames: int,
|
||||
context_frames: int,
|
||||
) -> torch.Tensor:
|
||||
"""Encode every tick with a fixed, past-only VAP context."""
|
||||
|
||||
if output_frames <= 0 or context_frames < 0:
|
||||
raise ValueError("VAP chunk/context frame counts are invalid")
|
||||
if waveform.shape != (num_frames * SAMPLES_PER_MODEL_FRAME,):
|
||||
raise ValueError("Audio/VAP high-rate lengths differ")
|
||||
pieces = []
|
||||
for output_start in range(0, num_frames, output_frames):
|
||||
input_start = max(0, output_start - context_frames)
|
||||
output_end = min(num_frames, output_start + output_frames)
|
||||
samples = waveform[
|
||||
input_start * SAMPLES_PER_MODEL_FRAME :
|
||||
output_end * SAMPLES_PER_MODEL_FRAME
|
||||
]
|
||||
encoded = encoder.encode_audio(samples, output_end - input_start)
|
||||
pieces.append(encoded[output_start - input_start :])
|
||||
result = torch.cat(pieces) if pieces else torch.empty(0, encoder.feature_dim)
|
||||
if result.shape != (num_frames, encoder.feature_dim):
|
||||
raise RuntimeError("Chunked VAP output did not cover every frame once")
|
||||
return result
|
||||
|
||||
|
||||
def _chunked_fused_asd_encode(
|
||||
encoder: FrozenMuVAPASDEncoder,
|
||||
faces: np.ndarray,
|
||||
audio_states: torch.Tensor,
|
||||
*,
|
||||
output_frames: int,
|
||||
context_frames: int,
|
||||
) -> torch.Tensor:
|
||||
"""Encode one consecutive face track into frozen ``Z_ASD`` states."""
|
||||
|
||||
if output_frames <= 0 or context_frames < 0:
|
||||
raise ValueError("ASD chunk/context frame counts are invalid")
|
||||
if audio_states.shape != (len(faces), encoder.audio_dim):
|
||||
raise ValueError("Face/shared-audio high-rate lengths differ")
|
||||
pieces = []
|
||||
for output_start in range(0, len(faces), output_frames):
|
||||
input_start = max(0, output_start - context_frames)
|
||||
output_end = min(len(faces), output_start + output_frames)
|
||||
encoded = encoder.encode_person_fused(
|
||||
faces[input_start:output_end],
|
||||
audio_states[input_start:output_end],
|
||||
)
|
||||
pieces.append(encoded[output_start - input_start :])
|
||||
result = torch.cat(pieces) if pieces else torch.empty(0, encoder.fused_dim)
|
||||
if result.shape != (len(faces), encoder.fused_dim):
|
||||
raise RuntimeError("Chunked fused ASD output did not cover every frame once")
|
||||
return result
|
||||
|
||||
|
||||
def extract_sample(
|
||||
media: Any,
|
||||
annotations: list[dict[int, list[dict[str, Any]]]],
|
||||
vap: FrozenMuVAPVAPEncoder,
|
||||
asd: FrozenMuVAPASDEncoder,
|
||||
task_times: torch.Tensor,
|
||||
args: argparse.Namespace,
|
||||
) -> dict[str, Any]:
|
||||
"""Extract 25 Hz VAP states and ragged fused-ASD person sets."""
|
||||
|
||||
ticks_per_task = round(MODEL_FPS / OUTPUT_HZ)
|
||||
if abs(ticks_per_task * OUTPUT_HZ - MODEL_FPS) > 1e-8:
|
||||
raise RuntimeError("MuVAP/task grids are not integral")
|
||||
num_ticks = len(task_times) * ticks_per_task
|
||||
waveform = _audio_for_tick_range(media, 0, num_ticks)
|
||||
vap_states = _chunked_vap_encode(
|
||||
vap,
|
||||
waveform,
|
||||
num_frames=num_ticks,
|
||||
output_frames=round(args.vap_output_seconds * MODEL_FPS),
|
||||
context_frames=round(args.vap_context_seconds * MODEL_FPS),
|
||||
)
|
||||
asd_audio = _chunked_audio_encode(
|
||||
asd,
|
||||
waveform,
|
||||
num_frames=num_ticks,
|
||||
output_frames=round(args.asd_output_seconds * MODEL_FPS),
|
||||
context_frames=round(args.asd_context_seconds * MODEL_FPS),
|
||||
)
|
||||
|
||||
rows: list[list[tuple[torch.Tensor, dict[str, Any]]]] = [
|
||||
[] for _ in range(num_ticks)
|
||||
]
|
||||
required_end = len(task_times) / OUTPUT_HZ
|
||||
for clip_index, clip in enumerate(media.clips):
|
||||
if float(clip.dense_start) >= required_end - 1e-8:
|
||||
break
|
||||
start_tick, stop_tick = _model_tick_bounds(clip, required_end)
|
||||
segments = _collect_clip_tracks(
|
||||
media,
|
||||
annotations,
|
||||
clip_index,
|
||||
start_tick,
|
||||
stop_tick,
|
||||
args,
|
||||
)
|
||||
for segment in tqdm(
|
||||
segments,
|
||||
desc=f"fused ASD {clip.path.name}",
|
||||
unit="track",
|
||||
disable=args.no_progress,
|
||||
leave=False,
|
||||
):
|
||||
positions = torch.tensor(segment["ticks"], dtype=torch.long)
|
||||
if bool(((positions < 0) | (positions >= num_ticks)).any()):
|
||||
raise RuntimeError("ASD track extends beyond the VAP timeline")
|
||||
states = _chunked_fused_asd_encode(
|
||||
asd,
|
||||
np.stack(segment["faces"]),
|
||||
asd_audio[positions],
|
||||
output_frames=round(args.asd_output_seconds * MODEL_FPS),
|
||||
context_frames=round(args.asd_context_seconds * MODEL_FPS),
|
||||
)
|
||||
for offset, tick in enumerate(segment["ticks"]):
|
||||
record = segment["records"][offset]
|
||||
if abs(float(record["asd_interval_end"]) - (tick + 1) / MODEL_FPS) > 1e-8:
|
||||
raise RuntimeError("ASD record/tick right edges differ")
|
||||
rows[tick].append((states[offset].to(torch.bfloat16), record))
|
||||
|
||||
features: list[torch.Tensor] = []
|
||||
records: list[dict[str, Any]] = []
|
||||
offsets = [0]
|
||||
max_people = 0
|
||||
for observations in rows:
|
||||
if observations:
|
||||
features.append(torch.stack([value for value, _ in observations]))
|
||||
records.extend(record for _, record in observations)
|
||||
max_people = max(max_people, len(observations))
|
||||
offsets.append(offsets[-1] + len(observations))
|
||||
flat = (
|
||||
torch.cat(features)
|
||||
if features
|
||||
else torch.empty(0, asd.fused_dim, dtype=torch.bfloat16)
|
||||
)
|
||||
task_frame_indices = torch.arange(len(task_times), dtype=torch.long)
|
||||
task_frame_indices = (task_frame_indices + 1) * ticks_per_task - 1
|
||||
if len(task_times):
|
||||
highrate_edges = (task_frame_indices + 1).to(torch.float64) / MODEL_FPS
|
||||
if not torch.allclose(highrate_edges, task_times, atol=1e-8, rtol=0):
|
||||
raise RuntimeError("25 Hz right edges do not match the 5 Hz task cutoffs")
|
||||
return {
|
||||
"vap": vap_states.to(torch.bfloat16),
|
||||
"frame_times": torch.arange(1, num_ticks + 1, dtype=torch.float64) / MODEL_FPS,
|
||||
"persons": {
|
||||
"features": flat,
|
||||
"frame_offsets": torch.tensor(offsets, dtype=torch.long),
|
||||
"records": records,
|
||||
},
|
||||
"task_frame_indices": task_frame_indices,
|
||||
"max_people": max_people,
|
||||
}
|
||||
|
||||
|
||||
def requested_config(
|
||||
args: argparse.Namespace,
|
||||
vap: FrozenMuVAPVAPEncoder,
|
||||
asd: FrozenMuVAPASDEncoder,
|
||||
) -> dict[str, Any]:
|
||||
return {
|
||||
"implementation": "clean-room-paper-reconstruction-not-upstream-stage3-code",
|
||||
"uses_omni_features": False,
|
||||
"muvap_source_revision": vap.source_revision,
|
||||
"vap_release_sha256": vap.release_sha256,
|
||||
"asd_release_sha256": asd.release_sha256,
|
||||
"vap_release_config": vap.release_config,
|
||||
"asd_release_config": asd.release_config,
|
||||
"vap_dim": vap.feature_dim,
|
||||
"asd_dim": asd.fused_dim,
|
||||
"model_fps": MODEL_FPS,
|
||||
"task_hz": OUTPUT_HZ,
|
||||
"sample_rate": SAMPLE_RATE,
|
||||
"face_size": [FACE_SIZE, FACE_SIZE],
|
||||
"vap_output_seconds": args.vap_output_seconds,
|
||||
"vap_context_seconds": args.vap_context_seconds,
|
||||
"asd_output_seconds": args.asd_output_seconds,
|
||||
"asd_context_seconds": args.asd_context_seconds,
|
||||
"audio_preprocessing": "causal-float-pcm-clip-only-no-future-rms",
|
||||
"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,
|
||||
"box_protocol": "provided-boxes-not-online-detector",
|
||||
"uses_speech_labels": False,
|
||||
"causal_policy": "25hz-right-edge-causal-fixed-past-context",
|
||||
}
|
||||
|
||||
|
||||
def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--dataset", choices=("easycom", "ego4d"), default="easycom")
|
||||
parser.add_argument("--data-root", type=Path, required=True)
|
||||
parser.add_argument(
|
||||
"--reference-feature-root",
|
||||
type=Path,
|
||||
required=True,
|
||||
help="Verified 5 Hz sidecars used only for timestamps/source inventory",
|
||||
)
|
||||
parser.add_argument("--output-root", type=Path, required=True)
|
||||
selection = parser.add_mutually_exclusive_group(required=True)
|
||||
selection.add_argument("--samples", nargs="+")
|
||||
selection.add_argument("--manifest", type=Path)
|
||||
parser.add_argument("--split", choices=("all", "train", "test"), default="all")
|
||||
parser.add_argument("--start-index", type=int, default=1)
|
||||
parser.add_argument("--end-index", type=int)
|
||||
parser.add_argument("--limit", type=int)
|
||||
parser.add_argument("--num-shards", type=int, default=1)
|
||||
parser.add_argument("--shard-index", type=int, default=0)
|
||||
parser.add_argument("--ego4d-annotations", type=Path, nargs="+")
|
||||
parser.add_argument(
|
||||
"--ego4d-max-terminal-overrun-frames",
|
||||
type=int,
|
||||
default=EGO4D_MAX_TERMINAL_FRAME_OVERRUN,
|
||||
)
|
||||
parser.add_argument("--muvap-root", type=Path, default=Path("third_party/MuVAP"))
|
||||
parser.add_argument("--muvap-revision", default="a1ef9c8")
|
||||
parser.add_argument("--vap-weights", type=Path, required=True)
|
||||
parser.add_argument("--asd-weights", type=Path, required=True)
|
||||
parser.add_argument("--device", default="cuda:0")
|
||||
parser.add_argument(
|
||||
"--dtype",
|
||||
choices=("float32", "float16", "bfloat16"),
|
||||
default="bfloat16",
|
||||
)
|
||||
parser.add_argument("--vap-output-seconds", type=float, default=10.0)
|
||||
parser.add_argument("--vap-context-seconds", type=float, default=20.0)
|
||||
parser.add_argument("--asd-output-seconds", type=float, default=10.0)
|
||||
parser.add_argument("--asd-context-seconds", type=float, default=2.0)
|
||||
parser.add_argument("--box-kind", choices=("face", "head"), default="face")
|
||||
parser.add_argument(
|
||||
"--box-coordinate-space", choices=("camera", "encoded"), default=None
|
||||
)
|
||||
parser.add_argument("--crop-expansion", type=float, default=1.0)
|
||||
parser.add_argument(
|
||||
"--decoder-backend",
|
||||
choices=("decord", "torchcodec", "pynvcodec"),
|
||||
default="decord",
|
||||
)
|
||||
parser.add_argument("--decoder-threads", type=int, default=1)
|
||||
parser.add_argument("--decoder-gpu-id", type=int, default=0)
|
||||
parser.add_argument("--max-steps", type=int)
|
||||
parser.add_argument("--overwrite", action="store_true")
|
||||
parser.add_argument("--dry-run", action="store_true")
|
||||
parser.add_argument("--no-progress", action="store_true")
|
||||
args = parser.parse_args(argv)
|
||||
if args.manifest is not None and args.dataset != "ego4d":
|
||||
parser.error("--manifest is only supported for 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
|
||||
):
|
||||
parser.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
|
||||
):
|
||||
parser.error("Invalid manifest range")
|
||||
if args.limit is not None and args.limit <= 0:
|
||||
parser.error("--limit must be positive")
|
||||
if args.num_shards <= 0 or not 0 <= args.shard_index < args.num_shards:
|
||||
parser.error("Require --num-shards > 0 and a valid --shard-index")
|
||||
if args.max_steps is not None and args.max_steps <= 0:
|
||||
parser.error("--max-steps must be positive")
|
||||
for name in (
|
||||
"vap_output_seconds",
|
||||
"vap_context_seconds",
|
||||
"asd_output_seconds",
|
||||
"asd_context_seconds",
|
||||
):
|
||||
value = getattr(args, name)
|
||||
if not math.isfinite(value) or value < 0 or (
|
||||
name.endswith("output_seconds") and value == 0
|
||||
):
|
||||
parser.error(f"--{name.replace('_', '-')} is invalid")
|
||||
frames = round(value * MODEL_FPS)
|
||||
if not math.isclose(frames / MODEL_FPS, value, abs_tol=1e-8):
|
||||
parser.error(f"--{name.replace('_', '-')} must align to the 25 Hz grid")
|
||||
if name.endswith("output_seconds") and frames == 0:
|
||||
parser.error(f"--{name.replace('_', '-')} is shorter than one frame")
|
||||
if not math.isfinite(args.crop_expansion) or args.crop_expansion <= 0:
|
||||
parser.error("--crop-expansion must be finite and positive")
|
||||
if args.decoder_threads <= 0 or args.decoder_gpu_id < 0:
|
||||
parser.error("decoder-threads must be positive and decoder-gpu-id nonnegative")
|
||||
if args.box_coordinate_space is None:
|
||||
args.box_coordinate_space = "encoded" if args.dataset == "easycom" else "camera"
|
||||
args.video_fps = MODEL_FPS
|
||||
return args
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> None:
|
||||
args = parse_args(argv)
|
||||
names = select_samples(args)
|
||||
if not names:
|
||||
print(f"Shard {args.shard_index}/{args.num_shards}: no assigned samples")
|
||||
return
|
||||
references = {
|
||||
name: resolve_feature_path(args.reference_feature_root, name, dataset=args.dataset)
|
||||
for name in names
|
||||
}
|
||||
for name in names:
|
||||
if not references[name].is_file():
|
||||
raise FileNotFoundError(references[name])
|
||||
destination = args.output_root / f"{name}.pt"
|
||||
if destination.exists() and not args.overwrite:
|
||||
raise FileExistsError(
|
||||
f"{destination} exists; use a new root or explicit --overwrite"
|
||||
)
|
||||
print(f"Shard {args.shard_index}/{args.num_shards}: {len(names)} sample(s)")
|
||||
if args.dry_run:
|
||||
for name in names:
|
||||
reference = torch.load(references[name], map_location="cpu", weights_only=False)
|
||||
times, _ = reference_alignment(reference)
|
||||
steps = min(len(times), args.max_steps or len(times))
|
||||
print(f"{name}: would extract {steps * 5} MuVAP frames for {steps} targets")
|
||||
return
|
||||
|
||||
ego_paths: list[Path] = []
|
||||
if args.dataset == "ego4d":
|
||||
ego_paths = args.ego4d_annotations or [
|
||||
args.data_root / "v1" / "annotations" / f"av_{split}.json"
|
||||
for split in ("train", "val")
|
||||
]
|
||||
ego_boxes = load_ego4d_boxes(ego_paths, names) if args.dataset == "ego4d" else {}
|
||||
vap = FrozenMuVAPVAPEncoder(
|
||||
muvap_root=args.muvap_root,
|
||||
revision=args.muvap_revision,
|
||||
release=args.vap_weights,
|
||||
device=args.device,
|
||||
dtype=args.dtype,
|
||||
)
|
||||
asd = FrozenMuVAPASDEncoder(
|
||||
muvap_root=args.muvap_root,
|
||||
revision=args.muvap_revision,
|
||||
release=args.asd_weights,
|
||||
device=args.device,
|
||||
dtype=args.dtype,
|
||||
)
|
||||
args.output_root.mkdir(parents=True, exist_ok=True)
|
||||
expected_config = requested_config(args, vap, asd)
|
||||
for position, name in enumerate(names, 1):
|
||||
print(f"[{position}/{len(names)}] {name}", flush=True)
|
||||
reference = torch.load(references[name], map_location="cpu", weights_only=False)
|
||||
times, indices = reference_alignment(reference)
|
||||
full_steps = len(times)
|
||||
if args.max_steps is not None:
|
||||
times = times[: args.max_steps]
|
||||
indices = indices[: args.max_steps]
|
||||
media, annotations = build_media_and_boxes(
|
||||
args, name, len(times) / OUTPUT_HZ, ego_boxes
|
||||
)
|
||||
validate_source_timeline(reference, media)
|
||||
output = extract_sample(media, annotations, vap, asd, times, args)
|
||||
output.update(
|
||||
{
|
||||
"format": MUVAP_EGOSPEAK_INPUT_FORMAT,
|
||||
"dataset": args.dataset,
|
||||
"sample_id": name,
|
||||
"times": times.to(torch.float64),
|
||||
"label_indices": indices.to(torch.long),
|
||||
"partial": len(times) != full_steps,
|
||||
"reference_features": str(references[name].resolve()),
|
||||
"direct_media": media.metadata(),
|
||||
"config": expected_config,
|
||||
}
|
||||
)
|
||||
if args.dataset == "ego4d":
|
||||
output["box_frame_audit"] = ego_boxes[name]["frame_audit"]
|
||||
destination = args.output_root / f"{name}.pt"
|
||||
with tempfile.NamedTemporaryFile(
|
||||
dir=args.output_root,
|
||||
prefix=f".{name}.",
|
||||
suffix=".tmp",
|
||||
delete=False,
|
||||
) as handle:
|
||||
temporary = Path(handle.name)
|
||||
try:
|
||||
torch.save(output, temporary)
|
||||
os.replace(temporary, destination)
|
||||
finally:
|
||||
temporary.unlink(missing_ok=True)
|
||||
print(
|
||||
f"Saved {destination}; MuVAP frames={len(output['vap'])}, "
|
||||
f"person observations={len(output['persons']['features'])}, "
|
||||
f"max_people={output['max_people']}, partial={output['partial']}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user