| |
| """Detect speech in an audio file and write the regions out as WebVTT. |
| |
| python example.py speech.wav |
| python example.py speech.mp3 --dtype fp16 --device cpu -o regions.vtt |
| |
| The VTT is cues covering the detected speech, with no text -- useful for |
| eyeballing the VAD's decisions in any media player, and as the segmentation |
| step in front of an ASR model. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| from pathlib import Path |
|
|
| import torch |
|
|
| from vad_torch import TorchVadRunner, DEFAULT_WEIGHTS, SAMPLE_RATE |
|
|
|
|
| def fmt_vtt_time(seconds: float) -> str: |
| h = int(seconds // 3600) |
| m = int((seconds % 3600) // 60) |
| s = seconds % 60 |
| return f"{h:02d}:{m:02d}:{s:06.3f}" |
|
|
|
|
| def main() -> int: |
| p = argparse.ArgumentParser(description=__doc__, |
| formatter_class=argparse.RawDescriptionHelpFormatter) |
| p.add_argument("audio", type=Path) |
| p.add_argument("-o", "--out", type=Path, default=None, |
| help="output .vtt (default: alongside the input)") |
| p.add_argument("--weights", type=Path, default=DEFAULT_WEIGHTS) |
| p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") |
| p.add_argument("--dtype", default="fp32", choices=["fp32", "fp16"], |
| help="fp32 (default) reproduces the ONNX reference exactly; " |
| "fp16 is ~2x faster and moves a few boundaries") |
| p.add_argument("--batch-size", type=int, default=8) |
| p.add_argument("--threshold", type=float, default=0.4, |
| help="in LOGIT space, not probability (default 0.4 ~= 0.60)") |
| a = p.parse_args() |
|
|
| if not a.audio.exists(): |
| raise SystemExit(f"no such file: {a.audio}") |
|
|
| try: |
| import librosa |
| except ImportError: |
| raise SystemExit("this example needs librosa: pip install librosa") |
|
|
| print(f"loading {a.audio.name}") |
| audio, _ = librosa.load(str(a.audio), sr=SAMPLE_RATE, mono=True) |
| duration = len(audio) / SAMPLE_RATE |
| print(f" {duration:.1f}s at {SAMPLE_RATE} Hz mono") |
|
|
| dt = {"fp32": torch.float32, "fp16": torch.float16}[a.dtype] |
| if a.dtype == "fp16" and a.device == "cpu": |
| raise SystemExit("fp16 on CPU is not supported; use --dtype fp32") |
|
|
| runner = TorchVadRunner(weights=a.weights, device=a.device, dtype=dt, |
| batch_size=a.batch_size) |
| print(f" model on {a.device}, {a.dtype}") |
|
|
| intervals = runner.speech_intervals(audio, threshold=a.threshold) |
| speech = sum(e - s for s, e in intervals) |
| print(f"\n{len(intervals)} speech regions, {speech:.1f}s of {duration:.1f}s " |
| f"({speech / max(duration, 1e-9) * 100:.1f}% active)") |
|
|
| out = a.out or a.audio.with_suffix(".vtt") |
| lines = ["WEBVTT", ""] |
| for i, (start, end) in enumerate(intervals, 1): |
| lines.append(str(i)) |
| lines.append(f"{fmt_vtt_time(start)} --> {fmt_vtt_time(end)}") |
| lines.append(f"[speech {i}]") |
| lines.append("") |
| out.write_text("\n".join(lines), encoding="utf-8") |
| print(f"-> {out}") |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
