94 lines
3.7 KiB
Python
94 lines
3.7 KiB
Python
"""Validate packed samples and report per-dataset signal statistics."""
|
|
|
|
import argparse
|
|
import json
|
|
import math
|
|
from collections import defaultdict
|
|
from pathlib import Path
|
|
from typing import Sequence
|
|
|
|
import numpy as np
|
|
|
|
from data.alignment import model_frame_count
|
|
from data.dataloaders.packed import PackedASDDataset
|
|
from data.media import WaveformStats, normalization_gain
|
|
from preprocess.asd.schema import MODEL_FPS
|
|
from preprocess.asd.writer import native_audio_samples
|
|
|
|
|
|
def percentile(values: list[float], points=(0, 1, 50, 95, 99, 100)) -> dict:
|
|
return {str(point): float(np.percentile(values, point)) for point in points}
|
|
|
|
|
|
def validate(roots: Path | Sequence[Path]) -> None:
|
|
dataset = PackedASDDataset(roots)
|
|
grouped = defaultdict(lambda: defaultdict(list))
|
|
|
|
for index in range(len(dataset)):
|
|
audio, visual, vad, _, metadata = dataset[index]
|
|
sample_id = metadata["sample_id"]
|
|
expected = model_frame_count(
|
|
metadata["source_frames"], metadata["source_fps"], MODEL_FPS
|
|
)
|
|
frames = visual.shape[0]
|
|
if frames == 0:
|
|
raise ValueError(f"{sample_id}: empty sample")
|
|
if frames != expected:
|
|
raise ValueError(
|
|
f"{sample_id}: {metadata['source_frames']} source frames at "
|
|
f"{metadata['source_fps']} fps imply {expected} model frames, "
|
|
f"loader produced {frames}"
|
|
)
|
|
if vad.shape[0] != frames or audio.shape[-1] != frames * 640:
|
|
raise ValueError(f"{sample_id}: stream-length mismatch")
|
|
if visual.shape[1:] != (112, 112) or visual.min() < 0 or visual.max() > 255:
|
|
raise ValueError(f"{sample_id}: invalid visual tensor")
|
|
if not np.isfinite(metadata["audio_peak"]):
|
|
raise ValueError(f"{sample_id}: invalid audio statistics")
|
|
if metadata["audio_samples"] != native_audio_samples(
|
|
metadata["source_frames"], metadata["source_fps"]
|
|
):
|
|
raise ValueError(f"{sample_id}: stored audio does not span the faces")
|
|
|
|
stats = grouped[metadata["dataset"]]
|
|
stats["source_audio_peak"].append(metadata["audio_peak"])
|
|
stats["source_audio_rms"].append(metadata["audio_rms"])
|
|
stats["source_audio_mean"].append(metadata["audio_mean"])
|
|
stats["speaking_ratio"].append(metadata["speaking_ratio"])
|
|
stats["normalized_rms"].append(
|
|
audio.square().mean().sqrt().item()
|
|
)
|
|
stats["normalized_peak"].append(audio.abs().max().item())
|
|
gain = normalization_gain(
|
|
WaveformStats(
|
|
mean=metadata["audio_mean"],
|
|
rms=metadata["audio_rms"],
|
|
peak=metadata["audio_peak"],
|
|
)
|
|
)
|
|
stats["gain_db"].append(20 * math.log10(max(gain, 1e-8)))
|
|
|
|
report = {"samples": len(dataset), "datasets": {}}
|
|
for name, stats in grouped.items():
|
|
report["datasets"][name] = {
|
|
"samples": len(stats["source_audio_peak"]),
|
|
"source_audio_peak_percentiles": percentile(stats["source_audio_peak"]),
|
|
"source_audio_rms_percentiles": percentile(stats["source_audio_rms"]),
|
|
"normalized_rms_percentiles": percentile(stats["normalized_rms"]),
|
|
"normalized_peak_percentiles": percentile(stats["normalized_peak"]),
|
|
"gain_db_percentiles": percentile(stats["gain_db"]),
|
|
"max_abs_source_mean": max(abs(v) for v in stats["source_audio_mean"]),
|
|
"speaking_ratio_percentiles": percentile(stats["speaking_ratio"]),
|
|
}
|
|
print(json.dumps(report, indent=2))
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("roots", type=Path, nargs="+")
|
|
validate(parser.parse_args().roots)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|