580 lines
23 KiB
Python
580 lines
23 KiB
Python
"""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()
|