847 lines
34 KiB
Python
847 lines
34 KiB
Python
"""Extract frozen 25 Hz causal per-person AV states for Omni floor prediction.
|
|
|
|
The global Qwen2.5-Omni states are not recomputed. A released MuVAP ASD model
|
|
provides one shared traditional CPC audio stream and one causal visual/activity
|
|
state per provided face track on a 25 Hz timeline. Only states whose right edge
|
|
lands on the verified 5 Hz task grid are saved. ``highrate-person-av-attention``
|
|
therefore compares ``global Omni + traditional audio`` against the same frozen
|
|
base plus person-specific visual/speaking evidence.
|
|
|
|
No speech/activity label is read by this extractor. Provided boxes are an
|
|
experimental isolation condition, not an online detector claim.
|
|
"""
|
|
|
|
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_turn_local_features import (
|
|
EGO4D_MAX_TERMINAL_FRAME_OVERRUN,
|
|
box_coordinate_offset,
|
|
build_media_and_boxes,
|
|
load_ego4d_boxes,
|
|
select_samples,
|
|
validate_source_timeline,
|
|
)
|
|
from local_turn_features import (
|
|
HIGHRATE_PERSON_AV_FORMAT,
|
|
crop_box,
|
|
reference_alignment,
|
|
)
|
|
|
|
|
|
MODEL_FPS = 25.0
|
|
OUTPUT_HZ = 5.0
|
|
FACE_SIZE = 112
|
|
SAMPLE_RATE = 16_000
|
|
SAMPLES_PER_MODEL_FRAME = SAMPLE_RATE // int(MODEL_FPS)
|
|
PERSON_FEATURE_KIND = "muvap-asd-causal-visual-plus-activity-logit"
|
|
AUDIO_FEATURE_KIND = "muvap-cpc-causal-audio-state"
|
|
PERSON_FEATURE_EXTRA_DIMS = 1
|
|
|
|
|
|
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 _release_paths(path: Path) -> tuple[Path, Path]:
|
|
path = path.expanduser().resolve()
|
|
weights = path / "weights.pt" if path.is_dir() else path
|
|
config = weights.parent / "config.yaml"
|
|
if not weights.is_file():
|
|
raise FileNotFoundError(weights)
|
|
if not config.is_file():
|
|
raise FileNotFoundError(
|
|
f"MuVAP release config must sit beside the weights: {config}"
|
|
)
|
|
return weights, config
|
|
|
|
|
|
class FrozenMuVAPASDEncoder:
|
|
"""Load the published causal ASD model without copying its source code."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
muvap_root: Path,
|
|
revision: str,
|
|
release: Path,
|
|
device: str,
|
|
dtype: str,
|
|
) -> None:
|
|
root = muvap_root.expanduser().resolve()
|
|
if not re.fullmatch(r"[0-9a-fA-F]{7,40}", revision):
|
|
raise ValueError(
|
|
"--muvap-revision must be a 7--40 character hexadecimal Git "
|
|
"commit prefix"
|
|
)
|
|
required = root / "models" / "asd.py"
|
|
if not required.is_file():
|
|
raise FileNotFoundError(
|
|
f"MuVAP source is missing: {required}. See third_party/README.md"
|
|
)
|
|
revision_file = root / ".upstream-revision"
|
|
if revision_file.is_file():
|
|
actual_revision = 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_revision = (
|
|
completed.stdout.strip().lower()
|
|
if completed.returncode == 0
|
|
else ""
|
|
)
|
|
if (
|
|
not re.fullmatch(r"[0-9a-f]{40}", actual_revision)
|
|
or not actual_revision.startswith(revision.lower())
|
|
):
|
|
raise RuntimeError(
|
|
f"MuVAP checkout is {actual_revision or 'unreadable'}, expected "
|
|
f"revision prefix {revision!r}. See third_party/README.md and "
|
|
f"{revision_file}"
|
|
)
|
|
self.source_revision = actual_revision
|
|
self.weights_path, config_path = _release_paths(release)
|
|
try:
|
|
import yaml
|
|
except ImportError as error:
|
|
raise RuntimeError("PyYAML is required to read the MuVAP release") from error
|
|
|
|
# MuVAP uses top-level imports such as ``models.asd``. Keep the
|
|
# dependency external/ignored and verify that Python resolves it from
|
|
# the explicitly requested checkout.
|
|
sys.path.insert(0, str(root))
|
|
try:
|
|
from models.asd import AudioVisualASD
|
|
from models.release import load_weights
|
|
except Exception as error:
|
|
raise RuntimeError(
|
|
f"Cannot import MuVAP from {root}; ensure the dependencies in "
|
|
"requirements.txt are installed and the pinned checkout is complete"
|
|
) from error
|
|
imported = Path(sys.modules[AudioVisualASD.__module__].__file__).resolve()
|
|
if root not in imported.parents:
|
|
raise RuntimeError(
|
|
f"Python resolved {AudioVisualASD.__module__} from {imported}, not {root}"
|
|
)
|
|
|
|
payload = yaml.safe_load(config_path.read_text()) or {}
|
|
config = payload.get("asd")
|
|
if not isinstance(config, dict):
|
|
raise ValueError(f"MuVAP release has no `asd` config: {config_path}")
|
|
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(
|
|
"This extractor currently freezes the published 25 Hz CPC ASD "
|
|
f"release; got audio_encoder={audio_name!r}"
|
|
)
|
|
for block in ("cross", "temporal"):
|
|
if config.get(block, {}).get("causal", True) is not True:
|
|
raise ValueError(
|
|
f"MuVAP ASD {block} stack must be causal for floor prediction"
|
|
)
|
|
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 = AudioVisualASD(config)
|
|
except FileNotFoundError as error:
|
|
raise RuntimeError(
|
|
"MuVAP's frozen CPC frontend is missing. Run "
|
|
f"`python {root / 'tools' / 'fetch_frontends.py'} cpc` first."
|
|
) from error
|
|
self.model = load_weights(model, self.weights_path).to(self.device).eval()
|
|
self.model.requires_grad_(False)
|
|
self.audio_dim = int(config["cross"]["d_model"])
|
|
self.visual_dim = int(config["cross"]["d_model"])
|
|
self.fused_dim = int(config["temporal"]["d_model"])
|
|
self.feature_dim = self.visual_dim + PERSON_FEATURE_EXTRA_DIMS
|
|
self.release_config = config
|
|
self.release_sha256 = _sha256(self.weights_path)
|
|
|
|
def _autocast(self):
|
|
return torch.autocast(
|
|
device_type=self.device.type,
|
|
dtype=self.autocast_dtype,
|
|
enabled=self.autocast_dtype is not None,
|
|
)
|
|
|
|
@torch.inference_mode()
|
|
def encode_audio(self, waveform: np.ndarray, num_frames: int) -> torch.Tensor:
|
|
"""Encode mixed audio once as causal 25 Hz CPC states."""
|
|
|
|
expected_samples = num_frames * SAMPLES_PER_MODEL_FRAME
|
|
if waveform.dtype != np.float32 or waveform.shape != (expected_samples,):
|
|
raise ValueError(
|
|
f"Audio must be float32 [{expected_samples}], got {waveform.shape}"
|
|
)
|
|
audio = torch.from_numpy(np.ascontiguousarray(waveform)).view(1, 1, -1)
|
|
audio = audio.to(self.device, dtype=torch.float32)
|
|
with self._autocast():
|
|
embedding = self.model.norm_audio(
|
|
self.model.encoder_proj(self.model.audio_encoder(audio))
|
|
)
|
|
if embedding.shape != (1, num_frames, self.audio_dim):
|
|
raise RuntimeError(
|
|
f"Unexpected MuVAP audio shape: {tuple(embedding.shape)}"
|
|
)
|
|
result = embedding[0].float().cpu()
|
|
if not torch.isfinite(result).all():
|
|
raise RuntimeError("MuVAP produced invalid shared audio states")
|
|
return result
|
|
|
|
@torch.inference_mode()
|
|
def encode_person(
|
|
self,
|
|
faces: np.ndarray,
|
|
audio_states: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
"""Encode one face track; keep visual state plus one AV activity logit."""
|
|
|
|
if (
|
|
faces.dtype != np.uint8
|
|
or faces.ndim != 3
|
|
or faces.shape[1:] != (FACE_SIZE, FACE_SIZE)
|
|
):
|
|
raise ValueError(f"Faces must be uint8 [T,112,112], got {faces.shape}")
|
|
if (
|
|
not isinstance(audio_states, torch.Tensor)
|
|
or audio_states.shape != (len(faces), self.audio_dim)
|
|
or not audio_states.is_floating_point()
|
|
or not torch.isfinite(audio_states).all()
|
|
):
|
|
raise ValueError(
|
|
"Person ASD requires finite aligned shared audio states "
|
|
f"[T,{self.audio_dim}]"
|
|
)
|
|
visual = torch.from_numpy(np.ascontiguousarray(faces)).unsqueeze(0)
|
|
visual = visual.to(self.device, dtype=torch.float32)
|
|
audio_embedding = audio_states[None].to(self.device, dtype=torch.float32)
|
|
with self._autocast():
|
|
visual_embedding = self.model.norm_visual(
|
|
self.model.visual_encoder(visual)
|
|
)
|
|
audio_cross = self.model.cross_v2a(
|
|
audio_embedding,
|
|
src=visual_embedding,
|
|
)
|
|
visual_cross = self.model.cross_a2v(
|
|
visual_embedding,
|
|
src=audio_embedding,
|
|
)
|
|
audio_logits = self.model.audio_head(audio_cross)
|
|
visual_logits = self.model.visual_head(visual_cross)
|
|
activity_logits = (audio_logits + visual_logits) / 2
|
|
result = torch.cat((visual_embedding, activity_logits), dim=-1)[0]
|
|
result = result.float().cpu()
|
|
if result.shape != (len(faces), self.feature_dim) or not torch.isfinite(result).all():
|
|
raise RuntimeError("MuVAP produced invalid person AV states")
|
|
return result
|
|
|
|
@torch.inference_mode()
|
|
def encode_person_fused(
|
|
self,
|
|
faces: np.ndarray,
|
|
audio_states: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
"""Return the released ASD backbone's fused ``Z_ASD`` state.
|
|
|
|
The ordinary local-feature experiment intentionally stores a visual
|
|
state and a scalar speaking score. MuVAP Stage 3 instead consumes the
|
|
fused output of the causal ASD temporal stack. Keeping this as a
|
|
separate method prevents the two feature contracts from being
|
|
silently interchanged.
|
|
"""
|
|
|
|
if (
|
|
faces.dtype != np.uint8
|
|
or faces.ndim != 3
|
|
or faces.shape[1:] != (FACE_SIZE, FACE_SIZE)
|
|
):
|
|
raise ValueError(f"Faces must be uint8 [T,112,112], got {faces.shape}")
|
|
if (
|
|
not isinstance(audio_states, torch.Tensor)
|
|
or audio_states.shape != (len(faces), self.audio_dim)
|
|
or not audio_states.is_floating_point()
|
|
or not torch.isfinite(audio_states).all()
|
|
):
|
|
raise ValueError(
|
|
"Person ASD requires finite aligned shared audio states "
|
|
f"[T,{self.audio_dim}]"
|
|
)
|
|
visual = torch.from_numpy(np.ascontiguousarray(faces)).unsqueeze(0)
|
|
visual = visual.to(self.device, dtype=torch.float32)
|
|
audio_embedding = audio_states[None].to(self.device, dtype=torch.float32)
|
|
with self._autocast():
|
|
visual_embedding = self.model.norm_visual(
|
|
self.model.visual_encoder(visual)
|
|
)
|
|
audio_cross = self.model.cross_v2a(
|
|
audio_embedding,
|
|
src=visual_embedding,
|
|
)
|
|
visual_cross = self.model.cross_a2v(
|
|
visual_embedding,
|
|
src=audio_embedding,
|
|
)
|
|
fused = self.model.transformer(
|
|
torch.cat((audio_cross, visual_cross), dim=-1)
|
|
)
|
|
result = fused[0].float().cpu()
|
|
if result.shape != (len(faces), self.fused_dim):
|
|
raise RuntimeError(
|
|
f"Unexpected MuVAP fused ASD shape: {tuple(result.shape)}"
|
|
)
|
|
if not torch.isfinite(result).all():
|
|
raise RuntimeError("MuVAP produced invalid fused ASD states")
|
|
return result
|
|
|
|
|
|
def _square_gray_face(crop: Any) -> np.ndarray:
|
|
"""Match MuVAP's 112x112 grayscale face geometry without distortion."""
|
|
|
|
from PIL import Image
|
|
|
|
gray = crop.convert("L")
|
|
size = max(gray.size)
|
|
canvas = Image.new("L", (size, size), color=128)
|
|
canvas.paste(gray, ((size - gray.width) // 2, (size - gray.height) // 2))
|
|
resampling = getattr(Image, "Resampling", Image).BILINEAR
|
|
return np.asarray(canvas.resize((FACE_SIZE, FACE_SIZE), resampling), dtype=np.uint8)
|
|
|
|
|
|
def _chunked_audio_encode(
|
|
encoder: FrozenMuVAPASDEncoder,
|
|
waveform: np.ndarray,
|
|
*,
|
|
num_frames: int,
|
|
output_frames: int,
|
|
context_frames: int,
|
|
) -> torch.Tensor:
|
|
"""Encode shared mixed audio causally and cover every 25 Hz frame once."""
|
|
|
|
if output_frames <= 0 or context_frames < 0:
|
|
raise ValueError("ASD chunk/context frame counts are invalid")
|
|
if len(waveform) != num_frames * SAMPLES_PER_MODEL_FRAME:
|
|
raise ValueError("Audio/high-rate frame 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)
|
|
sample_start = input_start * SAMPLES_PER_MODEL_FRAME
|
|
sample_end = output_end * SAMPLES_PER_MODEL_FRAME
|
|
encoded = encoder.encode_audio(
|
|
waveform[sample_start:sample_end],
|
|
output_end - input_start,
|
|
)
|
|
pieces.append(encoded[output_start - input_start :])
|
|
result = torch.cat(pieces) if pieces else torch.empty(0, encoder.audio_dim)
|
|
if result.shape != (num_frames, encoder.audio_dim):
|
|
raise RuntimeError("Chunked audio output did not cover every input frame once")
|
|
return result
|
|
|
|
|
|
def _chunked_person_encode(
|
|
encoder: FrozenMuVAPASDEncoder,
|
|
faces: np.ndarray,
|
|
audio_states: torch.Tensor,
|
|
*,
|
|
output_frames: int,
|
|
context_frames: int,
|
|
) -> torch.Tensor:
|
|
"""Encode one causal face track against the already shared audio stream."""
|
|
|
|
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(
|
|
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.feature_dim)
|
|
if result.shape != (len(faces), encoder.feature_dim):
|
|
raise RuntimeError("Chunked person output did not cover every input frame once")
|
|
return result
|
|
|
|
|
|
def _model_tick_bounds(clip: Any, end: float) -> tuple[int, int]:
|
|
start_value = float(clip.dense_start) * MODEL_FPS
|
|
stop_value = min(float(clip.dense_end), end) * MODEL_FPS
|
|
start = round(start_value)
|
|
stop = round(stop_value)
|
|
if abs(start_value - start) > 1e-6 or abs(stop_value - stop) > 1e-6:
|
|
raise ValueError(
|
|
"Source clip boundaries must align to the 25 Hz ASD lattice; "
|
|
f"got [{clip.dense_start}, {min(float(clip.dense_end), end)}]"
|
|
)
|
|
return start, stop
|
|
|
|
|
|
def _collect_clip_tracks(
|
|
media: Any,
|
|
annotations: list[dict[int, list[dict[str, Any]]]],
|
|
clip_index: int,
|
|
start_tick: int,
|
|
stop_tick: int,
|
|
args: argparse.Namespace,
|
|
) -> list[dict[str, Any]]:
|
|
"""Decode one source clip once and return consecutive anonymous tracks."""
|
|
|
|
active: dict[str, dict[str, Any]] = {}
|
|
segments: list[dict[str, Any]] = []
|
|
clip = media.clips[clip_index]
|
|
for tick in tqdm(
|
|
range(start_tick, stop_tick),
|
|
desc=f"ASD faces {clip.path.name}",
|
|
unit="frame",
|
|
disable=args.no_progress,
|
|
leave=False,
|
|
):
|
|
centre = (tick + 0.5) / MODEL_FPS
|
|
mapped_clip, frame_index, source_pts = media._source_frame_for_time(centre)
|
|
if mapped_clip != clip_index:
|
|
raise RuntimeError("25 Hz interval crossed a source-clip boundary")
|
|
image = media.frame_at_or_before(centre)
|
|
offset_x, offset_y = box_coordinate_offset(clip, args.box_coordinate_space)
|
|
seen: set[str] = set()
|
|
for person in annotations[clip_index].get(frame_index, []):
|
|
person_id = str(person.get("id", "")).strip()
|
|
if not person_id or person_id.lower() in {"-1", "none", "null", "unknown"}:
|
|
# Unknown detections cannot be joined across time. Give the
|
|
# one observation its own stable-within-extraction key.
|
|
person_id = f"unknown:{tick}:{len(seen)}"
|
|
if person_id in seen:
|
|
raise ValueError(
|
|
f"Duplicate person ID {person_id!r} at {centre:.3f}s"
|
|
)
|
|
seen.add(person_id)
|
|
box = crop_box(
|
|
person["xyxy"],
|
|
*image.size,
|
|
offset_x=offset_x,
|
|
offset_y=offset_y,
|
|
expansion=args.crop_expansion,
|
|
)
|
|
if box is None:
|
|
continue
|
|
record = {
|
|
"id": person_id,
|
|
"box": list(box),
|
|
"source_clip_index": clip_index,
|
|
"source_frame_index": frame_index,
|
|
"source_time": float(clip.dense_start + source_pts),
|
|
"asd_model_tick": tick,
|
|
"asd_interval_end": (tick + 1) / MODEL_FPS,
|
|
}
|
|
if record["source_time"] >= record["asd_interval_end"] - 1e-8:
|
|
raise RuntimeError(
|
|
"High-rate face sample is not strictly earlier than its "
|
|
f"ASD right edge: source={record['source_time']}, "
|
|
f"right_edge={record['asd_interval_end']}"
|
|
)
|
|
current = active.get(person_id)
|
|
if current is None or current["ticks"][-1] + 1 != tick:
|
|
if current is not None:
|
|
segments.append(current)
|
|
current = {"id": person_id, "ticks": [], "faces": [], "records": []}
|
|
active[person_id] = current
|
|
current["ticks"].append(tick)
|
|
current["faces"].append(_square_gray_face(image.crop(box)))
|
|
current["records"].append(record)
|
|
segments.extend(active.values())
|
|
return segments
|
|
|
|
|
|
def _audio_for_tick_range(media: Any, start_tick: int, stop_tick: int) -> np.ndarray:
|
|
if start_tick < 0 or stop_tick <= start_tick:
|
|
raise ValueError("ASD audio tick range must be nonempty and increasing")
|
|
waveform = media.audio_slice(
|
|
start_tick / MODEL_FPS,
|
|
stop_tick / MODEL_FPS,
|
|
).astype(np.float32, copy=False)
|
|
wanted = (stop_tick - start_tick) * SAMPLES_PER_MODEL_FRAME
|
|
result = np.zeros(wanted, dtype=np.float32)
|
|
used = min(wanted, len(waveform))
|
|
if used:
|
|
# AAC may decode a little outside nominal float PCM bounds. Clipping
|
|
# each observed sample is causal; content-dependent clip/prefix RMS
|
|
# normalization is intentionally not used because it would make the
|
|
# same past waveform depend on a later prediction cutoff.
|
|
result[:used] = np.clip(waveform[:used], -1.0, 1.0)
|
|
return result
|
|
|
|
|
|
def extract_highrate_person_states(
|
|
media: Any,
|
|
annotations: list[dict[int, list[dict[str, Any]]]],
|
|
encoder: FrozenMuVAPASDEncoder,
|
|
times: torch.Tensor,
|
|
args: argparse.Namespace,
|
|
) -> dict[str, Any]:
|
|
"""Produce one variable-cardinality person set at every 5 Hz right edge."""
|
|
|
|
rows: list[list[tuple[torch.Tensor, dict[str, Any]]]] = [
|
|
[] for _ in range(len(times))
|
|
]
|
|
shared_audio = torch.zeros(
|
|
len(times),
|
|
encoder.audio_dim,
|
|
dtype=torch.float32,
|
|
)
|
|
shared_audio_valid = torch.zeros(len(times), dtype=torch.bool)
|
|
required_end = len(times) / OUTPUT_HZ
|
|
output_frames = round(args.asd_output_seconds * MODEL_FPS)
|
|
context_frames = round(args.asd_context_seconds * MODEL_FPS)
|
|
total_ticks = len(times) * int(MODEL_FPS / OUTPUT_HZ)
|
|
session_waveform = _audio_for_tick_range(media, 0, total_ticks)
|
|
session_audio_states = _chunked_audio_encode(
|
|
encoder,
|
|
session_waveform,
|
|
num_frames=total_ticks,
|
|
output_frames=output_frames,
|
|
context_frames=context_frames,
|
|
)
|
|
shared_audio.copy_(
|
|
session_audio_states[
|
|
int(MODEL_FPS / OUTPUT_HZ) - 1 :: int(MODEL_FPS / OUTPUT_HZ)
|
|
]
|
|
)
|
|
shared_audio_valid.fill_(True)
|
|
for clip_index, clip in enumerate(media.clips):
|
|
if 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,
|
|
)
|
|
progress = tqdm(
|
|
segments,
|
|
desc=f"ASD tracks {clip.path.name}",
|
|
unit="track",
|
|
disable=args.no_progress,
|
|
leave=False,
|
|
)
|
|
for segment in progress:
|
|
ticks = segment["ticks"]
|
|
faces = np.stack(segment["faces"])
|
|
positions = torch.tensor(ticks, dtype=torch.long)
|
|
if bool(((positions < 0) | (positions >= len(session_audio_states))).any()):
|
|
raise RuntimeError("Person track extends beyond its shared audio stream")
|
|
states = _chunked_person_encode(
|
|
encoder,
|
|
faces,
|
|
session_audio_states[positions],
|
|
output_frames=output_frames,
|
|
context_frames=context_frames,
|
|
)
|
|
for offset, tick in enumerate(ticks):
|
|
if (tick + 1) % int(MODEL_FPS / OUTPUT_HZ):
|
|
continue
|
|
row = (tick + 1) // int(MODEL_FPS / OUTPUT_HZ) - 1
|
|
if not 0 <= row < len(rows):
|
|
continue
|
|
record = segment["records"][offset]
|
|
cutoff = float(times[row])
|
|
if (
|
|
abs(float(record["asd_interval_end"]) - cutoff) > 1e-7
|
|
or float(record["source_time"]) >= cutoff
|
|
):
|
|
raise RuntimeError(
|
|
"High-rate person observation is not strictly causal at "
|
|
f"row {row}: source={record['source_time']}, "
|
|
f"interval_end={record['asd_interval_end']}, cutoff={cutoff}"
|
|
)
|
|
rows[row].append((states[offset].to(torch.bfloat16), record))
|
|
|
|
if not bool(shared_audio_valid.all()):
|
|
missing = torch.nonzero(~shared_audio_valid).flatten().tolist()
|
|
raise RuntimeError(
|
|
f"Shared traditional audio misses 5 Hz rows: {missing[:20]}"
|
|
)
|
|
|
|
features = []
|
|
records = []
|
|
offsets = [0]
|
|
pooled = torch.zeros(len(rows), encoder.feature_dim, dtype=torch.float32)
|
|
valid = torch.zeros(len(rows), dtype=torch.bool)
|
|
for row_index, observations in enumerate(rows):
|
|
if observations:
|
|
values = torch.stack([item[0] for item in observations])
|
|
features.append(values)
|
|
records.extend(item[1] for item in observations)
|
|
pooled[row_index] = values.float().mean(0)
|
|
valid[row_index] = True
|
|
offsets.append(offsets[-1] + len(observations))
|
|
flat = (
|
|
torch.cat(features)
|
|
if features
|
|
else torch.empty(0, encoder.feature_dim, dtype=torch.bfloat16)
|
|
)
|
|
return {
|
|
"audio": shared_audio.to(torch.bfloat16),
|
|
"audio_valid": shared_audio_valid,
|
|
"visual": pooled.to(torch.bfloat16),
|
|
"visual_valid": valid,
|
|
"persons": {
|
|
"features": flat,
|
|
# These are interval right edges. Equality with the task cutoff is
|
|
# explicitly allowed by the high-rate loader.
|
|
"frame_times": times.to(torch.float64).clone(),
|
|
"frame_offsets": torch.tensor(offsets, dtype=torch.long),
|
|
"records": records,
|
|
},
|
|
}
|
|
|
|
|
|
def requested_config(
|
|
args: argparse.Namespace,
|
|
encoder: FrozenMuVAPASDEncoder,
|
|
) -> dict[str, Any]:
|
|
return {
|
|
"encoder": "MuVAP AudioVisualASD",
|
|
"encoder_source_revision": encoder.source_revision,
|
|
"encoder_release_sha256": encoder.release_sha256,
|
|
"encoder_release_config": encoder.release_config,
|
|
"dtype": args.dtype,
|
|
"model_fps": MODEL_FPS,
|
|
"output_hz": OUTPUT_HZ,
|
|
"asd_output_seconds": args.asd_output_seconds,
|
|
"asd_context_seconds": args.asd_context_seconds,
|
|
"face_size": [FACE_SIZE, FACE_SIZE],
|
|
"person_feature_kind": PERSON_FEATURE_KIND,
|
|
"person_visual_dim": encoder.visual_dim,
|
|
"person_activity_logit_dim": PERSON_FEATURE_EXTRA_DIMS,
|
|
"audio_feature_kind": AUDIO_FEATURE_KIND,
|
|
"visual_dim": encoder.feature_dim,
|
|
"audio_dim": encoder.audio_dim,
|
|
"audio_source": "single-shared-traditional-cpc-stream",
|
|
"person_tokens_duplicate_audio": False,
|
|
"asd_audio_preprocessing": "causal-float-pcm-clip-only-no-rms-normalization",
|
|
"asd_temporal_attention": "causal",
|
|
"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",
|
|
"causal_policy": "right-edge-causal-highrate-asd",
|
|
"uses_speech_labels": False,
|
|
}
|
|
|
|
|
|
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("--feature-root", type=Path, required=True, help="Global Omni .pt root")
|
|
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("--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("--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, help="Smoke only; partial output is rejected by training")
|
|
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 0 <= --shard-index < --num-shards")
|
|
if args.max_steps is not None and args.max_steps <= 0:
|
|
parser.error("--max-steps must be positive")
|
|
if not math.isfinite(args.asd_output_seconds) or args.asd_output_seconds <= 0:
|
|
parser.error("--asd-output-seconds must be finite and positive")
|
|
if not math.isfinite(args.asd_context_seconds) or args.asd_context_seconds < 0:
|
|
parser.error("--asd-context-seconds must be finite and nonnegative")
|
|
if round(args.asd_output_seconds * MODEL_FPS) <= 0:
|
|
parser.error("--asd-output-seconds is shorter than one model frame")
|
|
if args.box_coordinate_space is None:
|
|
args.box_coordinate_space = "encoded" if args.dataset == "easycom" else "camera"
|
|
# ``build_media_and_boxes`` consumes this compatibility field.
|
|
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 clips")
|
|
return
|
|
print(f"Shard {args.shard_index}/{args.num_shards}: {len(names)} assigned samples")
|
|
references = {
|
|
name: resolve_feature_path(args.feature_root, name, dataset=args.dataset)
|
|
for name in names
|
|
}
|
|
for name in names:
|
|
required = [references[name]]
|
|
if args.dataset == "ego4d":
|
|
required.append(args.data_root / "v1" / "clips" / f"{name}.mp4")
|
|
for path in required:
|
|
if not path.is_file():
|
|
raise FileNotFoundError(path)
|
|
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"
|
|
)
|
|
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} states ({steps / OUTPUT_HZ:g}s)")
|
|
return
|
|
|
|
paths: list[Path] = []
|
|
if args.dataset == "ego4d":
|
|
paths = args.ego4d_annotations or [
|
|
args.data_root / "v1" / "annotations" / f"av_{split}.json"
|
|
for split in ("train", "val")
|
|
]
|
|
ego_boxes = load_ego4d_boxes(paths, names) if args.dataset == "ego4d" else {}
|
|
encoder = 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_contract = None
|
|
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, indices = times[: args.max_steps], 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_highrate_person_states(
|
|
media,
|
|
annotations,
|
|
encoder,
|
|
times,
|
|
args,
|
|
)
|
|
config = requested_config(args, encoder)
|
|
contract = {"format": HIGHRATE_PERSON_AV_FORMAT, "modalities": "av", "config": config}
|
|
if expected_contract is not None and contract != expected_contract:
|
|
raise ValueError("High-rate encoder/audio contract differs across samples")
|
|
expected_contract = contract
|
|
output.update(
|
|
{
|
|
"format": HIGHRATE_PERSON_AV_FORMAT,
|
|
"dataset": args.dataset,
|
|
"sample_id": name,
|
|
"times": times,
|
|
"label_indices": indices,
|
|
"partial": len(times) != full_steps,
|
|
"reference_features": str(references[name].resolve()),
|
|
"direct_media": media.metadata(),
|
|
"config": 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)
|
|
count = len(output["persons"]["features"])
|
|
coverage = float(output["visual_valid"].float().mean())
|
|
print(
|
|
f"Saved {destination}; states={len(times)}, person observations={count}, "
|
|
f"person-row coverage={coverage:.1%}, partial={output['partial']}",
|
|
flush=True,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|