115 lines
4.2 KiB
Python
115 lines
4.2 KiB
Python
"""Time-based alignment utilities for audio, video frames, and frame labels.
|
|
|
|
The model runs on a fixed 25 Hz timeline. Source videos remain at their native
|
|
frame rate; frames and labels are sampled onto the model timeline at load time.
|
|
"""
|
|
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Sequence
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class AlignmentConfig:
|
|
model_fps: float = 25.0
|
|
audio_sample_rate: int = 16_000
|
|
|
|
def __post_init__(self) -> None:
|
|
if self.model_fps <= 0:
|
|
raise ValueError("model_fps must be positive")
|
|
if self.audio_sample_rate <= 0:
|
|
raise ValueError("audio_sample_rate must be positive")
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class AVSample:
|
|
sample_id: str
|
|
audio_path: Path
|
|
visual_path: Path
|
|
source_fps: float
|
|
labels: Sequence[int]
|
|
start_sec: float = 0.0
|
|
|
|
def __post_init__(self) -> None:
|
|
if self.source_fps <= 0:
|
|
raise ValueError(f"{self.sample_id}: source_fps must be positive")
|
|
if self.start_sec < 0:
|
|
raise ValueError(f"{self.sample_id}: start_sec cannot be negative")
|
|
|
|
|
|
def model_frame_count(source_frames: int, source_fps: float, model_fps: float) -> int:
|
|
"""Return the number of complete model intervals in a source-frame span."""
|
|
if source_frames < 0 or source_fps <= 0 or model_fps <= 0:
|
|
raise ValueError("frame counts and frame rates must be valid")
|
|
return int(np.floor((source_frames / source_fps) * model_fps + 1e-9))
|
|
|
|
|
|
def source_indices(
|
|
source_frames: int, source_fps: float, model_fps: float
|
|
) -> np.ndarray:
|
|
"""Map model-frame centres to the source intervals that contain them."""
|
|
count = model_frame_count(source_frames, source_fps, model_fps)
|
|
model_centres = (np.arange(count, dtype=np.float64) + 0.5) / model_fps
|
|
indices = np.floor(model_centres * source_fps).astype(np.int64)
|
|
return np.clip(indices, 0, max(source_frames - 1, 0))
|
|
|
|
|
|
def model_indices(
|
|
source_frames: int, source_fps: float, model_fps: float
|
|
) -> np.ndarray:
|
|
"""Map source-frame centres to the model intervals that contain them.
|
|
|
|
The inverse direction of `source_indices`, used to lift model-rate
|
|
predictions back onto the native annotation timeline for evaluation.
|
|
Always returns exactly `source_frames` indices, so a caller can score one
|
|
prediction per native frame regardless of the source frame rate.
|
|
"""
|
|
if source_frames < 0:
|
|
raise ValueError("source_frames cannot be negative")
|
|
count = model_frame_count(source_frames, source_fps, model_fps)
|
|
source_centres = (np.arange(source_frames, dtype=np.float64) + 0.5) / source_fps
|
|
indices = np.floor(source_centres * model_fps).astype(np.int64)
|
|
return np.clip(indices, 0, max(count - 1, 0))
|
|
|
|
|
|
def align_visual_and_labels(
|
|
visual: np.ndarray,
|
|
labels: Sequence[int],
|
|
source_fps: float,
|
|
model_fps: float,
|
|
strict: bool = True,
|
|
) -> tuple[np.ndarray, np.ndarray]:
|
|
"""Sample faces and their labels with one shared timestamp mapping.
|
|
|
|
A length disagreement means the face track and the annotation disagree
|
|
about how long the clip is, which silently shifts every later label.
|
|
`strict` reports that instead of trimming to the shorter of the two.
|
|
"""
|
|
labels_array = np.asarray(labels, dtype=np.float32)
|
|
if strict and len(visual) != len(labels_array):
|
|
raise ValueError(
|
|
f"visual has {len(visual)} frames but labels have {len(labels_array)}"
|
|
)
|
|
usable_frames = min(len(visual), len(labels_array))
|
|
if usable_frames == 0:
|
|
raise ValueError("visual and labels must contain at least one frame")
|
|
|
|
indices = source_indices(usable_frames, source_fps, model_fps)
|
|
return visual[:usable_frames][indices], labels_array[:usable_frames][indices]
|
|
|
|
|
|
def fit_audio_to_frames(
|
|
waveform: torch.Tensor,
|
|
frame_count: int,
|
|
sample_rate: int,
|
|
model_fps: float,
|
|
) -> torch.Tensor:
|
|
"""Pad/crop audio to the exact duration represented by model frames."""
|
|
target_samples = round(frame_count * sample_rate / model_fps)
|
|
waveform = waveform[..., :target_samples]
|
|
return F.pad(waveform, (0, max(0, target_samples - waveform.shape[-1])))
|