Add causal person fusion and MuVAP stage-3 baselines
This commit is contained in:
Xin Wang
+204
@@ -0,0 +1,204 @@
|
||||
"""Pack decoded native-rate samples into memory-mapped NumPy shards.
|
||||
|
||||
Nothing here is specific to the model that will read the shards. Faces and
|
||||
labels keep the source frame rate, audio keeps its full duration and its
|
||||
original gain, and the per-recording level statistics needed to reproduce a
|
||||
normalization travel in the index. Resampling onto a model timeline and
|
||||
applying loudness normalization are the reader's job.
|
||||
"""
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Iterable
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from data.alignment import model_frame_count
|
||||
from data.media import (
|
||||
AudioSpec,
|
||||
load_waveform,
|
||||
slice_audio,
|
||||
validate_visual,
|
||||
waveform_stats,
|
||||
)
|
||||
from preprocess.asd.schema import (
|
||||
FACE_SIZE,
|
||||
FORMAT_NAME,
|
||||
FORMAT_VERSION,
|
||||
MODEL_FPS,
|
||||
SAMPLE_RATE,
|
||||
SourceRecord,
|
||||
)
|
||||
|
||||
|
||||
def _load_faces(record: SourceRecord, face_size: int = FACE_SIZE) -> np.ndarray:
|
||||
if record.visual_kind == "npy":
|
||||
return validate_visual(np.load(record.visual_path, allow_pickle=False))
|
||||
|
||||
paths = sorted(record.visual_path.glob("*.jpg"), key=lambda path: float(path.stem))
|
||||
if not paths:
|
||||
raise FileNotFoundError(f"{record.sample_id}: no JPEG faces")
|
||||
frames = []
|
||||
for path in paths:
|
||||
frame = cv2.imread(str(path), cv2.IMREAD_GRAYSCALE)
|
||||
if frame is None:
|
||||
raise ValueError(f"{record.sample_id}: could not decode {path}")
|
||||
frames.append(
|
||||
cv2.resize(frame, (face_size, face_size), interpolation=cv2.INTER_LINEAR)
|
||||
)
|
||||
return np.stack(frames).astype(np.uint8, copy=False)
|
||||
|
||||
|
||||
def _pcm16(waveform: torch.Tensor) -> np.ndarray:
|
||||
values = waveform.squeeze(0).clamp(-1.0, 1.0)
|
||||
return (values * 32767.0).round().to(torch.int16).cpu().numpy()
|
||||
|
||||
|
||||
def native_audio_samples(source_frames: int, source_fps: float) -> int:
|
||||
"""Audio length covering exactly the span of the stored face frames."""
|
||||
return round(source_frames / source_fps * SAMPLE_RATE)
|
||||
|
||||
|
||||
class ShardWriter:
|
||||
def __init__(self, output_dir: Path, max_shard_bytes: int = 1_000_000_000):
|
||||
self.output_dir = output_dir
|
||||
self.max_shard_bytes = max_shard_bytes
|
||||
self.shard_index = 0
|
||||
self.samples_written = 0
|
||||
self._reset()
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
def _reset(self) -> None:
|
||||
self.faces: list[np.ndarray] = []
|
||||
self.audio: list[np.ndarray] = []
|
||||
self.labels: list[np.ndarray] = []
|
||||
self.entries: list[dict] = []
|
||||
self.byte_count = 0
|
||||
self.face_count = 0
|
||||
self.audio_count = 0
|
||||
self.label_count = 0
|
||||
|
||||
def add_arrays(
|
||||
self,
|
||||
*,
|
||||
dataset: str,
|
||||
sample_id: str,
|
||||
faces: np.ndarray,
|
||||
waveform: torch.Tensor,
|
||||
labels,
|
||||
source_fps: float,
|
||||
start_sec: float = 0.0,
|
||||
) -> None:
|
||||
faces = validate_visual(faces)
|
||||
labels = np.asarray(labels, dtype=np.uint8)
|
||||
if len(faces) != len(labels):
|
||||
raise ValueError(
|
||||
f"{sample_id}: {len(faces)} faces but {len(labels)} labels"
|
||||
)
|
||||
source_frames = len(labels)
|
||||
if source_frames == 0:
|
||||
raise ValueError(f"{sample_id}: sample is empty")
|
||||
|
||||
waveform = slice_audio(
|
||||
waveform,
|
||||
SAMPLE_RATE,
|
||||
target_samples=native_audio_samples(source_frames, source_fps),
|
||||
)
|
||||
stats = waveform_stats(waveform)
|
||||
audio = _pcm16(waveform)
|
||||
|
||||
sample_bytes = faces.nbytes + audio.nbytes + labels.nbytes
|
||||
if self.entries and self.byte_count + sample_bytes > self.max_shard_bytes:
|
||||
self.flush()
|
||||
|
||||
self.entries.append(
|
||||
{
|
||||
"dataset": dataset,
|
||||
"sample_id": sample_id,
|
||||
"source_fps": source_fps,
|
||||
"source_frames": source_frames,
|
||||
"start_sec": start_sec,
|
||||
# Reported for convenience; the loader recomputes it.
|
||||
"model_frames": model_frame_count(source_frames, source_fps, MODEL_FPS),
|
||||
"face_offset": self.face_count,
|
||||
"audio_offset": self.audio_count,
|
||||
"audio_samples": len(audio),
|
||||
"label_offset": self.label_count,
|
||||
"audio_mean": stats.mean,
|
||||
"audio_rms": stats.rms,
|
||||
"audio_peak": stats.peak,
|
||||
"speaking_ratio": float(labels.mean()),
|
||||
}
|
||||
)
|
||||
self.faces.append(faces)
|
||||
self.audio.append(audio)
|
||||
self.labels.append(labels)
|
||||
self.byte_count += sample_bytes
|
||||
self.face_count += len(faces)
|
||||
self.audio_count += len(audio)
|
||||
self.label_count += len(labels)
|
||||
|
||||
def add(self, record: SourceRecord) -> None:
|
||||
faces = _load_faces(record)
|
||||
waveform, sample_rate = load_waveform(
|
||||
record.audio_path,
|
||||
AudioSpec(sample_rate=SAMPLE_RATE, mono=True),
|
||||
npy_sample_rate=record.npy_audio_sample_rate,
|
||||
)
|
||||
waveform = slice_audio(waveform, sample_rate, start_time=record.start_sec)
|
||||
self.add_arrays(
|
||||
dataset=record.dataset,
|
||||
sample_id=record.sample_id,
|
||||
faces=faces,
|
||||
waveform=waveform,
|
||||
labels=record.labels,
|
||||
source_fps=record.source_fps,
|
||||
start_sec=record.start_sec,
|
||||
)
|
||||
|
||||
def flush(self) -> None:
|
||||
if not self.entries:
|
||||
return
|
||||
shard_dir = self.output_dir / f"shard-{self.shard_index:06d}"
|
||||
shard_dir.mkdir(parents=False, exist_ok=False)
|
||||
np.save(shard_dir / "faces.npy", np.concatenate(self.faces), allow_pickle=False)
|
||||
np.save(shard_dir / "audio.npy", np.concatenate(self.audio), allow_pickle=False)
|
||||
np.save(
|
||||
shard_dir / "labels.npy", np.concatenate(self.labels), allow_pickle=False
|
||||
)
|
||||
with (shard_dir / "index.json").open("w") as handle:
|
||||
json.dump(self.entries, handle, separators=(",", ":"))
|
||||
self.samples_written += len(self.entries)
|
||||
self.shard_index += 1
|
||||
self._reset()
|
||||
|
||||
def close(self) -> None:
|
||||
self.flush()
|
||||
manifest = {
|
||||
"format": FORMAT_NAME,
|
||||
"version": FORMAT_VERSION,
|
||||
"timeline": "native",
|
||||
"audio_sample_rate": SAMPLE_RATE,
|
||||
"audio_normalization": None,
|
||||
"face_shape": [FACE_SIZE, FACE_SIZE],
|
||||
"face_color": "grayscale",
|
||||
"face_dtype": "uint8",
|
||||
"audio_dtype": "int16",
|
||||
"label_dtype": "uint8",
|
||||
"shards": self.shard_index,
|
||||
"samples": self.samples_written,
|
||||
}
|
||||
with (self.output_dir / "dataset.json").open("w") as handle:
|
||||
json.dump(manifest, handle, indent=2)
|
||||
|
||||
|
||||
def write_records(
|
||||
records: Iterable[SourceRecord], output_dir: Path, max_shard_bytes: int
|
||||
) -> int:
|
||||
writer = ShardWriter(output_dir, max_shard_bytes)
|
||||
for record in records:
|
||||
writer.add(record)
|
||||
writer.close()
|
||||
return writer.samples_written
|
||||
Reference in New Issue
Block a user