File size: 3,130 Bytes
3a490d9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 | #!/usr/bin/env python3
"""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())
|