trav
618f7ffc22
CPU-only online speaker labelling for caption windows (CAPTION_DIARIZATION=1): pyannote segmentation-3.0 finds local turns, a TitaNet-small voiceprint via sherpa-onnx matches them to per-room speakers numbered in arrival order. /transcribe?speakers=1 returns committed words grouped into segments with a speaker ID and the name assigned for the show. Models are SHA-256 pinned and cached with the Whisper weights; audio and voiceprints are never stored. Rooms gain a persisted cast list and in-memory per-show names through /api/rooms/<id>/speakers. The capture page offers "Identify speakers automatically", publishes one relay caption per speaker turn with label and spk, and shares a speaker-names panel (speaker-names.js) with the editor. Also adds scripts/diarize_eval.py, deploy.py --speakers, and copies rooms.py into the GPU image.
70 lines
3.1 KiB
Python
70 lines
3.1 KiB
Python
"""Replay a 16 kHz mono WAV through caption windows with speaker labels.
|
|
|
|
Prints each labelled caption and the speaker-label cost per window. Use
|
|
--threads 2 to approximate a small CPU host.
|
|
"""
|
|
import argparse
|
|
import sys
|
|
import time
|
|
import wave
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
|
from server import Engine, RATE # noqa: E402
|
|
from speakers import OnnxBackend, SpeakerTracker # noqa: E402
|
|
|
|
|
|
def read_wav(path):
|
|
with wave.open(str(path)) as wav:
|
|
if wav.getframerate() != RATE or wav.getnchannels() != 1 or wav.getsampwidth() != 2:
|
|
raise SystemExit('Expected 16 kHz mono 16-bit PCM WAV; convert with ffmpeg -ar 16000 -ac 1')
|
|
return np.frombuffer(wav.readframes(wav.getnframes()), dtype='<i2').astype(np.float32) / 32768
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('wav')
|
|
parser.add_argument('--model', default='small')
|
|
parser.add_argument('--models', default='/models', help='Directory holding the speakers/ model cache')
|
|
parser.add_argument('--threads', type=int, default=2)
|
|
parser.add_argument('--window', type=int, choices=(3, 6, 9), default=6)
|
|
parser.add_argument('--threshold', type=float, nargs='+', default=[0.5],
|
|
help='One or more cosine thresholds; each gets its own tracker over the same transcript')
|
|
parser.add_argument('--max-speakers', type=int, default=12)
|
|
parser.add_argument('--offline', action='store_true')
|
|
args = parser.parse_args()
|
|
audio = read_wav(args.wav)
|
|
engine = Engine(args.model, threads=args.threads, offline=args.offline)
|
|
backend = OnnxBackend(args.models, args.threads, args.offline)
|
|
trackers = {t: SpeakerTracker(backend, t, args.max_speakers) for t in args.threshold}
|
|
lines = {t: [] for t in trackers}
|
|
start, context, costs, whisper = 0, 0.0, [], []
|
|
while start < len(audio):
|
|
end = min(len(audio), start + round((context + args.window) * RATE))
|
|
final = end == len(audio)
|
|
chunk = audio[start:end]
|
|
began = time.perf_counter()
|
|
text, _, committed, words = engine.window(chunk, 'en', context, final, words=True)
|
|
whisper.append(time.perf_counter() - began)
|
|
for threshold, tracker in trackers.items():
|
|
began = time.perf_counter()
|
|
segments = tracker.label(chunk, words)
|
|
costs.append(time.perf_counter() - began)
|
|
lines[threshold] += [f'[{(start / RATE + context):7.2f}s] {s["speaker"] or "?":>4}: {s["text"]}'
|
|
for s in segments]
|
|
if final:
|
|
break
|
|
keep = min(0.5, committed - context) if committed > context else 0
|
|
start += round((committed - keep) * RATE)
|
|
context = keep
|
|
for threshold, tracker in trackers.items():
|
|
print(f'== threshold {threshold}: {len(tracker.detected)} speakers', *lines[threshold], sep='\n')
|
|
print(f'\nwindows={len(whisper)} label mean={np.mean(costs):.3f}s max={np.max(costs):.3f}s '
|
|
f'whisper mean={np.mean(whisper):.3f}s', file=sys.stderr)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|