desktop-subtitle-replay/livecap.py

livecap.py (46121 bytes)
1 #!/usr/bin/env python3
2 """
3 livecap.py - live speech-to-captions for OBS.
4 
5 Captures audio (WASAPI loopback = "whatever your speakers play", or any mic),
6 segments it with an energy VAD, transcribes with faster-whisper, and publishes
7 captions over WebSocket to an OBS Browser Source overlay.
8 
9     python livecap.py --list-devices
10     python livecap.py --selftest 12
11     python livecap.py --lang fi --model small
12 """
13 
14 import argparse
15 import asyncio
16 import json
17 import os
18 import queue
19 import re
20 import sys
21 import threading
22 import time
23 import wave
24 from collections import deque
25 from functools import partial
26 from http.server import SimpleHTTPRequestHandler, ThreadingHTTPServer
27 from pathlib import Path
28 
29 import warnings
30 
31 import numpy as np
32 import soundcard as sc
33 from scipy.signal import resample_poly
34 
35 # Loopback capture reports gaps whenever the device goes idle. Harmless for
36 # speech recognition, but it otherwise floods the console during quiet moments.
37 try:
38     warnings.filterwarnings("ignore", category=sc.SoundcardRuntimeWarning)
39 except AttributeError:
40     warnings.filterwarnings("ignore", message="data discontinuity in recording")
41 
42 HERE = Path(__file__).resolve().parent
43 TARGET_SR = 16000
44 CAPTURE_SR = 48000     # WASAPI shared-mode mix rate on virtually every Windows box
45 BLOCK_MS = 32
46 
47 # Whisper invents these when fed silence or noise. Only applied to short results.
48 HALLUCINATIONS = [
49     # Finnish
50     r"^tekstitys",
51     r"^tekstityksen tuotti",
52     r"^k[aa]a?nn[oo]s",
53     r"^kiitos( kun katsoit| paljon| katsomisesta)?[.!]?$",
54     r"^suomennos",
55     # Japanese - Whisper emits these constantly over silence
56     r"ご(視聴|清聴)(いただき)?ありがとうございま",
57     r"チャンネル登録",
58     r"^字幕",
59     r"^おわり[。]?$",
60     r"^(お|ご)?しまい[。]?$",
61     r"^ありがとうございま(した|す)[。]?$",
62     r"^(えー|あー|ん|う)+[ー。]?$",
63     # Russian
64     r"^субтитры",
65     r"^спасибо за просмотр",
66     r"^продолжение следует",
67     r"^редактор субтитров",
68     # English / generic
69     r"^subtitles? by",
70     r"^amara\.org",
71     r"^thanks? for watching",
72     r"^please subscribe",
73     r"^\W*$",
74 ]
75 HALLUCINATION_RE = [re.compile(p, re.I) for p in HALLUCINATIONS]
76 
77 
78 # Captions are frequently non-Latin (ru, ja, ...). Windows otherwise encodes
79 # stdout with the console codepage, which mangles them once output is redirected.
80 for _s in (sys.stdout, sys.stderr):
81     try:
82         _s.reconfigure(encoding="utf-8", errors="replace")
83     except (AttributeError, ValueError):
84         pass
85 
86 
87 def log(*a):
88     print("[" + time.strftime("%H:%M:%S") + "]", *a, flush=True)
89 
90 
91 # ---------------------------------------------------------------- devices ---
92 
93 def candidates(loopback):
94     """Devices usable in the current mode, in a stable order."""
95     mics = sc.all_microphones(include_loopback=True)
96     return [m for m in mics if bool(m.isloopback) == bool(loopback)]
97 
98 
99 def list_devices():
100     print("\n=== LOOPBACK sources (default mode: captures what this device plays) ===")
101     for i, m in enumerate(candidates(True)):
102         print("  %2d  %s  (%dch)" % (i, m.name, m.channels))
103     try:
104         print("\n  default (used when you pass no --audio-device): %s"
105               % sc.default_speaker().name)
106     except Exception:
107         pass
108 
109     print("\n=== INPUT devices (--mic mode: microphones, virtual cables) ===")
110     for i, m in enumerate(candidates(False)):
111         print("  %2d  %s  (%dch)" % (i, m.name, m.channels))
112     try:
113         print("\n  default (--mic with no --audio-device): %s" % sc.default_microphone().name)
114     except Exception:
115         pass
116     print("\nSelect with --audio-device followed by the number above, or any\n"
117           "case-insensitive part of the name (e.g. --audio-device CABLE).\n")
118 
119 
120 def resolve_device(spec, loopback):
121     """spec: None, an index into the listing, or a substring of the device name."""
122     pool = candidates(loopback)
123     if not pool:
124         raise RuntimeError("No %s devices found." % ("loopback" if loopback else "input"))
125 
126     if spec is None:
127         try:
128             if loopback:
129                 return sc.get_microphone(sc.default_speaker().id, include_loopback=True)
130             return sc.default_microphone()
131         except Exception:
132             return pool[0]
133 
134     try:
135         idx = int(spec)
136     except ValueError:
137         pass
138     else:
139         if 0 <= idx < len(pool):
140             return pool[idx]
141         raise RuntimeError("Device index %d out of range (0-%d)." % (idx, len(pool) - 1))
142 
143     want = spec.lower()
144     for m in pool:
145         if want in m.name.lower():
146             return m
147     raise RuntimeError("No %s device matching %r. Run --list-devices."
148                        % ("loopback" if loopback else "input", spec))
149 
150 
151 # ---------------------------------------------------------------- capture ---
152 
153 class Capture:
154     """Pumps mono float32 blocks off a soundcard device onto a queue."""
155 
156     def __init__(self, mic, loopback=True):
157         self.mic = mic
158         self.name = mic.name
159         self.loopback = loopback
160         self.sr = CAPTURE_SR
161         self.channels = max(1, int(mic.channels))
162         self.blocksize = int(self.sr * BLOCK_MS / 1000)
163         self.q = queue.Queue(maxsize=256)
164         self.dropped = 0
165         self.error = None
166         self._stop = threading.Event()
167         self._ready = threading.Event()
168         self._thread = None
169 
170     def _pump(self):
171         try:
172             with self.mic.recorder(samplerate=self.sr, channels=self.channels,
173                                    blocksize=self.blocksize) as rec:
174                 self._ready.set()
175                 while not self._stop.is_set():
176                     data = rec.record(numframes=self.blocksize)
177                     if data.ndim > 1 and data.shape[1] > 1:
178                         mono = data.mean(axis=1)
179                     else:
180                         mono = data.reshape(-1)
181                     try:
182                         self.q.put_nowait(np.ascontiguousarray(mono, dtype=np.float32))
183                     except queue.Full:
184                         self.dropped += 1
185         except Exception as e:
186             self.error = e
187             log("capture failed:", repr(e))
188         finally:
189             self._ready.set()
190 
191     def __enter__(self):
192         self._thread = threading.Thread(target=self._pump, daemon=True)
193         self._thread.start()
194         self._ready.wait(timeout=10)
195         if self.error is not None:
196             raise self.error
197         log("capturing: %s  (%dch @ %dHz, %s)"
198             % (self.name, self.channels, self.sr, "loopback" if self.loopback else "input"))
199         return self
200 
201     def __exit__(self, *a):
202         self._stop.set()
203         if self._thread is not None:
204             self._thread.join(timeout=2)
205 
206 
207 def to_whisper(audio, sr):
208     """native-rate mono float32 -> 16 kHz float32, gently gain-staged."""
209     if sr != TARGET_SR:
210         g = int(np.gcd(sr, TARGET_SR))
211         audio = resample_poly(audio, TARGET_SR // g, sr // g)
212     audio = np.asarray(audio, dtype=np.float32)
213     peak = float(np.max(np.abs(audio))) if audio.size else 0.0
214     if 0.0 < peak < 0.35:
215         audio = audio * (0.35 / peak)
216     return np.clip(audio, -1.0, 1.0).astype(np.float32)
217 
218 
219 # -------------------------------------------------------------- publishing --
220 
221 class Bus:
222     def __init__(self, args):
223         self.args = args
224         self.loop = None
225         self.clients = set()
226         self.history = deque(maxlen=12)
227         self.finals = deque(maxlen=max(1, args.txt_lines))
228         self.txt = HERE / args.txt
229         self.logf = HERE / "captions.log"
230 
231     def attach(self, loop):
232         self.loop = loop
233 
234     def mark_session(self):
235         """A marker so the reader can show this session's history after a reload."""
236         try:
237             with self.logf.open("a", encoding="utf-8") as fh:
238                 fh.write("=== session %s ===\n" % time.strftime("%Y-%m-%d %H:%M:%S"))
239         except OSError as e:
240             log("could not write captions.log:", e)
241 
242     def publish(self, msg):
243         if msg["type"] == "final":
244             self.history.append(msg)
245             self.finals.append(msg["text"])
246             self._write_files(msg)
247         if self.loop is not None:
248             payload = json.dumps(msg, ensure_ascii=False)
249             self.loop.call_soon_threadsafe(self._fanout, payload)
250 
251     def _fanout(self, payload):
252         for ws in list(self.clients):
253             try:
254                 asyncio.get_running_loop().create_task(ws.send(payload))
255             except Exception:
256                 self.clients.discard(ws)
257 
258     def _write_files(self, msg):
259         try:
260             self.txt.write_text("\n".join(self.finals) + "\n", encoding="utf-8")
261             line = time.strftime("%Y-%m-%d %H:%M:%S") + "\t" + msg["text"]
262             if msg.get("tr"):
263                 line += "\t|| " + msg["tr"]
264             with self.logf.open("a", encoding="utf-8") as fh:
265                 fh.write(line + "\n")
266         except OSError as e:
267             log("file write failed:", e)
268 
269 
270 async def ws_server(bus, args, ctl, stop):
271     import websockets
272 
273     async def handler(ws):
274         bus.clients.add(ws)
275         log("client connected (%d total)" % len(bus.clients))
276         try:
277             await ws.send(json.dumps(ctl.status(), ensure_ascii=False))
278             for m in list(bus.history)[-4:]:
279                 await ws.send(json.dumps(m, ensure_ascii=False))
280             async for raw in ws:
281                 ctl.handle(raw)
282         except Exception:
283             pass
284         finally:
285             bus.clients.discard(ws)
286             log("client disconnected (%d total)" % len(bus.clients))
287 
288     async with websockets.serve(handler, "127.0.0.1", args.ws_port, ping_interval=20):
289         log("websocket  ws://127.0.0.1:%d" % args.ws_port)
290         await stop.wait()
291 
292 
293 class QuietHandler(SimpleHTTPRequestHandler):
294     def log_message(self, *a):
295         pass
296 
297 
298 def http_server(args):
299     handler = partial(QuietHandler, directory=str(HERE))
300     srv = ThreadingHTTPServer(("127.0.0.1", args.http_port), handler)
301     threading.Thread(target=srv.serve_forever, daemon=True).start()
302     url = ("http://127.0.0.1:%d/overlay.html?ws=%d&lines=%d&size=%d&hide=%s"
303            % (args.http_port, args.ws_port, args.lines, args.size, args.hide))
304     if args.translate:
305         url += "&tr=1"
306     log("OBS Browser Source URL (copy this):")
307     log("  " + url)
308     log("Control panel (open in any browser):")
309     log("  http://127.0.0.1:%d/control.html?ws=%d" % (args.http_port, args.ws_port))
310     log("Reader for Yomitan mining (open in your real browser):")
311     log("  http://127.0.0.1:%d/reader.html?ws=%d" % (args.http_port, args.ws_port))
312     return srv
313 
314 
315 # ----------------------------------------------------------- transcription --
316 
317 class Controller:
318     """Applies live language/model changes coming from the control panel."""
319 
320     def __init__(self, args, bus):
321         self.args = args
322         self.bus = bus
323         self.tr = None                      # set once the Transcriber exists
324         self.cmds = queue.Queue()           # model swaps run on the worker thread
325         self.loading = False
326         self.detected = None                # what 'auto' has settled on
327 
328     def status(self):
329         return {
330             "type": "status",
331             "lang": self.args.lang,
332             "detected": self.detected,
333             "model": self.args.model,
334             "loading": self.loading,
335             "translate": bool(self.args.translate),
336             "partials": bool(self.args.partials),
337             "langs": self.args.langs,
338             "models": self.args.model_choices,
339         }
340 
341     def push_status(self):
342         self.bus.publish(self.status())
343 
344     def handle(self, raw):
345         try:
346             m = json.loads(raw)
347         except (ValueError, TypeError):
348             return
349         cmd, val = m.get("cmd"), m.get("value")
350 
351         if cmd == "set_lang":
352             if not valid_language(val):
353                 log("ignoring unknown language %r" % (val,))
354                 return
355             if val != self.args.lang:
356                 self.args.lang = val
357                 self.detected = None
358                 if self.tr is not None:
359                     self.tr.context = ""        # old-language prompt would poison it
360                     # Evidence gathered for a different language is worthless.
361                     self.tr.detector.reset()
362                     self.tr._last_detect = 0.0
363                 log("language -> %s" % val)
364                 self.bus.publish({"type": "clear"})
365             self.push_status()
366 
367         elif cmd == "set_model":
368             if val and val != self.args.model:
369                 self.cmds.put(("model", val))
370             else:
371                 self.push_status()
372 
373         elif cmd == "set_translate":
374             self.args.translate = bool(val)
375             log("translate -> %s" % self.args.translate)
376             self.push_status()
377 
378         elif cmd == "set_partials":
379             self.args.partials = bool(val)
380             log("partials -> %s" % self.args.partials)
381             self.push_status()
382 
383         elif cmd == "clear":
384             self.bus.publish({"type": "clear"})
385 
386         elif cmd == "test":
387             # Lets you position the OBS overlay without having to talk.
388             text = val if isinstance(val, str) and val.strip() else \
389                 "Testiteksti — проверка — テスト — caption preview"
390             self.bus.publish({"type": "final", "text": text, "tr": "",
391                               "ts": time.time()})
392 
393         elif cmd == "status":
394             self.push_status()
395 
396 
397 # Whisper spreads probability across close relatives. When the relative is not
398 # one of the languages in use, its mass belongs to the neighbour that is -
399 # otherwise a Finnish clip can lose to Finnish-plus-Estonian splitting the vote.
400 CONFUSABLE = {
401     "et": "fi",                                          # Estonian
402     "uk": "ru", "be": "ru", "bg": "ru", "mk": "ru",      # Cyrillic Slavic
403     "sr": "ru", "kk": "ru",
404     "gl": "pt",                                          # Galician leans Portuguese
405     "ca": "es", "oc": "es", "an": "es",                  # Iberian Romance
406     "cy": "en", "gd": "en",                              # frequent English mis-picks
407 }
408 
409 
410 class LanguageDetector:
411     """Language identification restricted to the languages actually in use.
412 
413     Three things make Whisper's own per-window detection unreliable here:
414     it may answer with any of 99 languages, it decides afresh every call so the
415     answer flaps mid-conversation, and short utterances carry little evidence.
416 
417     So: fold the distribution onto the allowed set, accumulate it over time with
418     a decay, and only switch when a challenger leads by a margin for several
419     observations in a row.
420     """
421 
422     def __init__(self, allowed, margin=1.3, hold=2, decay=0.55, min_audio=1.6,
423                  min_mass=0.5, min_confidence=0.55):
424         self.allowed = [c for c in allowed if c != "auto"] or ["en"]
425         self.margin = margin
426         self.hold = hold
427         self.decay = decay
428         self.min_audio = min_audio
429         # Restricting the distribution deletes the competitors, which inflates
430         # confidence: "en 0.38, ko 0.25, nn 0.10" becomes a commanding en 0.85.
431         # How much mass landed inside the allowed set is the honest signal of
432         # whether this audio is any of these languages at all.
433         self.min_mass = min_mass
434         self.min_confidence = min_confidence
435         self.scores = {c: 0.0 for c in self.allowed}
436         self.current = None
437         self._pending = None
438         self._pending_n = 0
439 
440     def reset(self):
441         self.scores = {c: 0.0 for c in self.allowed}
442         self.current = None
443         self._pending, self._pending_n = None, 0
444 
445     def fold(self, probs):
446         """Restrict the distribution to the allowed set, returning (dict, mass)."""
447         out = {c: 0.0 for c in self.allowed}
448         for lang, p in probs:
449             target = lang if lang in out else CONFUSABLE.get(lang)
450             if target in out:
451                 out[target] += p
452         return out, sum(out.values())
453 
454     def confidence(self):
455         total = sum(self.scores.values())
456         if total <= 0 or self.current is None:
457             return 0.0
458         return self.scores[self.current] / total
459 
460     def observe(self, probs, duration=3.0):
461         """Feed one detection result; returns the language to actually use."""
462         folded, mass = self.fold(probs)
463         if mass <= 0:
464             return self.current
465 
466         # Longer audio is better evidence than a two-word utterance, and mass
467         # outside the allowed set means this probably is not one of them.
468         reliability = min(1.0, mass / self.min_mass) if self.min_mass > 0 else 1.0
469         weight = max(0.05, min(1.0, duration / 3.0)) * reliability
470         for code in self.scores:
471             self.scores[code] = (self.scores[code] * self.decay
472                                  + (folded[code] / mass) * weight)
473 
474         best = max(self.scores, key=lambda c: self.scores[c])
475         if self.current is None:
476             # Adopting a language from one ambiguous reading is how a whole
477             # session ends up locked to the wrong one.
478             total = sum(self.scores.values())
479             share = self.scores[best] / total if total > 0 else 0.0
480             if mass >= self.min_mass and share >= self.min_confidence:
481                 self.current = best
482             return self.current
483 
484         leader, held = self.scores[best], self.scores[self.current]
485         if best != self.current and leader > held * self.margin:
486             self._pending_n = self._pending_n + 1 if self._pending == best else 1
487             self._pending = best
488             if self._pending_n >= self.hold:
489                 self.current = best
490                 self._pending, self._pending_n = None, 0
491         else:
492             self._pending, self._pending_n = None, 0
493         return self.current
494 
495 
496 def valid_language(code):
497     if code == "auto":
498         return True
499     if not isinstance(code, str) or not code:
500         return False
501     try:
502         from faster_whisper.tokenizer import _LANGUAGE_CODES
503         return code in _LANGUAGE_CODES
504     except Exception:
505         return bool(re.fullmatch(r"[a-z]{2,3}", code))
506 
507 
508 def collapse_looping_tail(s, min_reps=3):
509     """Trim a unit that repeats to the end of the string.
510 
511     The loop does not necessarily start at the beginning - "この日の日の日の日"
512     is "こ" followed by "の日" four times - so anchor the search at the end.
513     """
514     n = len(s)
515     for p in range(1, n // min_reps + 1):
516         unit = s[n - p:]
517         k = 0
518         while (k + 1) * p <= n and s[n - (k + 1) * p: n - k * p] == unit:
519             k += 1
520         if k >= min_reps:
521             return s[: n - (k - 1) * p]
522     return s
523 
524 
525 def collapse_repeats(text):
526     """Whisper degenerates into repeating a phrase; keep one copy.
527 
528     Two shapes show up: identical chunks joined from consecutive segments
529     ("A A A"), and a unit looped inside one chunk ("この日の日の日").
530     """
531     s = " ".join(text.split())
532     if not s:
533         return s
534 
535     out = []
536     for c in s.split(" "):
537         # Only whole phrases collapse on a single repeat; short words legitimately
538         # repeat ("very very good"), so those need three in a row.
539         run = 1 if len(c) >= 8 else 2
540         if len(out) >= run and all(x == c for x in out[-run:]):
541             continue
542         out.append(c)
543 
544     return " ".join(collapse_looping_tail(c) for c in out)
545 
546 
547 def is_repetitive(text):
548     """True when a caption is mostly one phrase repeated."""
549     s = " ".join(text.split())
550     if len(s) < 8:
551         return False
552     return len(collapse_repeats(s)) * 2 < len(s)
553 
554 
555 def looks_hallucinated(text):
556     t = text.strip().lower()
557     if len(t) > 40:
558         return False
559     return any(r.search(t) for r in HALLUCINATION_RE)
560 
561 
562 class Transcriber:
563     def __init__(self, args, bus):
564         from faster_whisper import WhisperModel
565 
566         self.args = args
567         self.bus = bus
568         log("loading model %r on %s/%s (first run downloads it)..."
569             % (args.model, args.compute_device, args.compute))
570         t0 = time.time()
571         self.model = WhisperModel(
572             args.model, device=args.compute_device, compute_type=args.compute,
573             cpu_threads=args.threads, num_workers=1)
574         log("model ready in %.1fs" % (time.time() - t0))
575         self.context = ""
576         self.busy = threading.Event()
577         self.detector = LanguageDetector(
578             args.langs, margin=args.detect_margin, hold=args.detect_hold,
579             min_audio=args.detect_min_audio)
580         self._last_detect = 0.0
581 
582     def pick_language(self, audio, duration, ctl=None):
583         """Resolve 'auto' to one of the configured languages.
584 
585         Detection runs on a schedule rather than every segment: it costs an
586         encoder pass, and once the answer is settled re-deciding constantly is
587         what makes it flap. Passing an explicit language into transcribe() also
588         skips Whisper's own internal detection, so this is close to free.
589         """
590         a = self.args
591         if a.lang != "auto":
592             return a.lang
593 
594         due = (self.detector.current is None
595                or time.time() - self._last_detect >= a.detect_every)
596         if due and duration >= self.detector.min_audio:
597             self._last_detect = time.time()
598             try:
599                 _, _, probs = self.model.detect_language(audio)
600             except Exception as e:
601                 log("language detection failed:", repr(e))
602                 return self.detector.current
603             before = self.detector.current
604             now = self.detector.observe(probs, duration)
605             if now != before:
606                 log("detected language: %s (confidence %.2f)"
607                     % (now, self.detector.confidence()))
608                 self.context = ""          # prompt from another language misleads
609                 if ctl is not None:
610                     ctl.detected = now
611                     ctl.push_status()
612         return self.detector.current
613 
614     def reload(self, name, ctl):
615         """Swap the model in place. Runs on the worker thread."""
616         from faster_whisper import WhisperModel
617 
618         a = self.args
619         ctl.loading = True
620         ctl.push_status()
621         log("loading model %r ..." % name)
622         t0 = time.time()
623         try:
624             new = WhisperModel(name, device=a.compute_device, compute_type=a.compute,
625                                cpu_threads=a.threads, num_workers=1)
626         except Exception as e:
627             log("model %r failed to load (%s); staying on %r" % (name, e, a.model))
628             ctl.loading = False
629             ctl.push_status()
630             return
631         old, self.model = self.model, new
632         del old
633         a.model = name
634         self.context = ""
635         ctl.loading = False
636         ctl.push_status()
637         log("model -> %s (%.1fs)" % (name, time.time() - t0))
638 
639     def run(self, audio, final, language=None):
640         a = self.args
641         lang = language or (None if a.lang == "auto" else a.lang)
642         segs, info = self.model.transcribe(
643             audio,
644             language=lang,
645             task="transcribe",
646             beam_size=a.beam if final else 1,
647             temperature=0.0,
648             condition_on_previous_text=False,
649             initial_prompt=(self.context or None) if (final and not a.no_context) else None,
650             vad_filter=final,
651             no_speech_threshold=0.6,
652             log_prob_threshold=-1.0,
653             repetition_penalty=a.repetition_penalty,
654             without_timestamps=True,
655         )
656         parts, nsp = [], []
657         for s in segs:
658             parts.append(s.text.strip())
659             nsp.append(getattr(s, "no_speech_prob", 0.0))
660         text = re.sub(r"\s+", " ", " ".join(parts)).strip()
661         return text, (max(nsp) if nsp else 1.0), info
662 
663     def translate(self, audio):
664         segs, _ = self.model.transcribe(
665             audio, language=None if self.args.lang == "auto" else self.args.lang,
666             task="translate", beam_size=1, temperature=0.0,
667             condition_on_previous_text=False, vad_filter=True, without_timestamps=True)
668         return re.sub(r"\s+", " ", " ".join(s.text.strip() for s in segs)).strip()
669 
670     def worker(self, jobs, stop, ctl):
671         while not stop.is_set():
672             # Model swaps take priority and must happen on this thread.
673             try:
674                 what, value = ctl.cmds.get_nowait()
675             except queue.Empty:
676                 pass
677             else:
678                 if what == "model":
679                     self.busy.set()
680                     try:
681                         while True:                 # drop audio queued for the old model
682                             jobs.get_nowait()
683                             jobs.task_done()
684                     except queue.Empty:
685                         pass
686                     try:
687                         self.reload(value, ctl)
688                     finally:
689                         self.busy.clear()
690                 continue
691 
692             try:
693                 kind, audio, sr = jobs.get(timeout=0.25)
694             except queue.Empty:
695                 continue
696             self.busy.set()
697             try:
698                 t0 = time.time()
699                 a16 = to_whisper(audio, sr)
700                 dur = len(a16) / TARGET_SR
701                 lang = (self.pick_language(a16, dur, ctl) if kind == "final"
702                         else self.detector.current)
703                 text, nsp, _ = self.run(a16, final=(kind == "final"), language=lang)
704                 if not text:
705                     if kind == "final":
706                         log("no speech recognised in %.1fs segment "
707                             "(music/noise, or wrong --lang)" % dur)
708                     continue
709                 if kind == "final":
710                     if nsp > 0.75 or looks_hallucinated(text):
711                         log("dropped (no_speech=%.2f): %r" % (nsp, text))
712                         continue
713                     looped = is_repetitive(text)
714                     text = collapse_repeats(text)
715                     if looped:
716                         # Feeding a looped caption back as the prompt is how the
717                         # loop sustains itself across segments.
718                         self.context = ""
719                     tr = self.translate(a16) if self.args.translate else ""
720                     if not self.args.no_context and not looped:
721                         self.context = (self.context + " " + text)[-220:]
722                     took = time.time() - t0
723                     log("FINAL %4.1fs audio in %4.1fs (rtf %.2f)  %s"
724                         % (dur, took, took / max(dur, 0.01), text))
725                     self.bus.publish({"type": "final", "text": text,
726                                       "tr": tr, "ts": time.time()})
727                 else:
728                     self.bus.publish({"type": "partial", "text": text, "ts": time.time()})
729             except Exception as e:
730                 log("transcribe error:", repr(e))
731             finally:
732                 self.busy.clear()
733                 jobs.task_done()
734 
735 
736 # ---------------------------------------------------------- segmentation ----
737 
738 SENT_END = (".", "!", "?", "…", "。", "！", "？")
739 
740 
741 def norm_word(w):
742     return re.sub(r"[^\w]", "", w.strip().lower(), flags=re.UNICODE)
743 
744 
745 class StreamDecoder:
746     """LocalAgreement streaming (Macháček et al.).
747 
748     Whisper cannot decode incrementally, so instead we re-decode a growing
749     buffer and only commit the prefix that two consecutive hypotheses agree
750     on. Agreement is a good proxy for stability: text that survives another
751     decode with more audio behind it rarely changes again. This trades CPU
752     for latency - words appear while someone is still talking, rather than
753     a whole sentence landing after they stop.
754     """
755 
756     def __init__(self, tr, bus, args):
757         self.tr, self.bus, self.args = tr, bus, args
758         self.native = []            # unconsumed audio at capture rate
759         self.prev = []              # previous hypothesis, uncommitted part
760         self.sentence = []          # committed words of the sentence in progress
761         self.speech = 0.0           # seconds of speech currently buffered
762 
763     # ---- audio -------------------------------------------------------
764     def add(self, block, voiced, dur):
765         self.native.append(block)
766         if voiced:
767             self.speech += dur
768 
769     def buffered_seconds(self, sr):
770         return sum(len(b) for b in self.native) / sr
771 
772     def _audio16(self, sr):
773         return to_whisper(np.concatenate(self.native), sr)
774 
775     def _trim(self, cut_s, sr):
776         """Drop audio up to cut_s seconds, keeping block boundaries simple."""
777         drop = int(cut_s * sr)
778         merged = np.concatenate(self.native)
779         merged = merged[min(drop, len(merged)):]
780         self.native = [merged] if len(merged) else []
781         self.speech = max(0.0, self.speech - cut_s)
782 
783     # ---- decoding ----------------------------------------------------
784     def _hypothesis(self, sr):
785         a = self.args
786         audio = self._audio16(sr)
787         # Whisper invents text when handed a very short buffer, and in
788         # streaming that invention gets committed before real audio arrives.
789         if len(audio) < int(a.stream_min_audio * TARGET_SR):
790             return [], audio
791         segs, _ = self.tr.model.transcribe(
792             audio,
793             language=None if a.lang == "auto" else a.lang,
794             beam_size=1,
795             temperature=0.0,
796             condition_on_previous_text=False,
797             initial_prompt=(self.tr.context or None) if not a.no_context else None,
798             vad_filter=False,
799             word_timestamps=True,
800             no_speech_threshold=0.6,
801             log_prob_threshold=-1.0,
802         )
803         words = []
804         for s in segs:
805             words.extend(getattr(s, "words", None) or [])
806         return words, audio
807 
808     def step(self, sr):
809         words, _ = self._hypothesis(sr)
810         if not words:
811             return
812 
813         k = 0
814         while (k < len(words) and k < len(self.prev)
815                and norm_word(words[k].word) == norm_word(self.prev[k].word)
816                and norm_word(words[k].word)):
817             k += 1
818 
819         confirmed, rest = words[:k], words[k:]
820         if confirmed:
821             cut = confirmed[-1].end
822             self.sentence.extend(confirmed)
823             self._trim(cut, sr)
824             # Remaining words are now measured against a shorter buffer.
825             for w in rest:
826                 w.start = max(0.0, w.start - cut)
827                 w.end = max(0.0, w.end - cut)
828         self.prev = rest
829 
830         text = self._text(self.sentence)
831         tail = self._text(rest)
832         if confirmed and text and text.strip().endswith(SENT_END):
833             self.flush()
834         elif text or tail:
835             self.bus.publish({"type": "partial",
836                               "text": (text + " " + tail).strip(),
837                               "ts": time.time()})
838 
839     @staticmethod
840     def _text(words):
841         return re.sub(r"\s+", " ", "".join(w.word for w in words)).strip()
842 
843     def flush(self, drop_audio=False):
844         """Emit the sentence built so far as a final caption."""
845         text = self._text(self.sentence)
846         self.sentence = []
847         if drop_audio:
848             self.native, self.prev, self.speech = [], [], 0.0
849         if not text or looks_hallucinated(text):
850             if text:
851                 log("dropped: %r" % text)
852             return
853         if not self.args.no_context:
854             self.tr.context = (self.tr.context + " " + text)[-220:]
855         log("FINAL  %s" % text)
856         self.bus.publish({"type": "final", "text": text, "tr": "", "ts": time.time()})
857 
858 
859 def stream_segmenter(cap, tr, args, stop, bus):
860     """Low-latency path: continuous re-decode with LocalAgreement commits."""
861     dur = cap.blocksize / cap.sr
862     noise = 1e-4
863     dec = StreamDecoder(tr, bus, args)
864     silence = 0.0
865     last_step = 0.0
866     log("streaming mode: committing on agreement every %.1fs" % args.stream_interval)
867 
868     while not stop.is_set():
869         drained = 0
870         while True:
871             try:
872                 blk = cap.q.get_nowait()
873             except queue.Empty:
874                 break
875             drained += 1
876             rms = float(np.sqrt(np.mean(blk * blk)) + 1e-12)
877             if rms < noise:
878                 noise = 0.90 * noise + 0.10 * rms
879             else:
880                 noise = 0.995 * noise + 0.005 * rms
881             voiced = rms > max(noise * args.vad_ratio, args.vad_floor)
882             silence = 0.0 if voiced else silence + dur
883             if voiced or dec.native:
884                 dec.add(blk, voiced, dur)
885 
886         if not drained:
887             time.sleep(0.02)
888 
889         now = time.time()
890         buffered = dec.buffered_seconds(cap.sr)
891 
892         # A long pause ends the sentence: commit whatever is left.
893         if dec.native and silence >= args.pause:
894             if dec.speech >= args.min_speech:
895                 dec.step(cap.sr)
896                 if dec.prev:
897                     dec.sentence.extend(dec.prev)
898                     dec.prev = []
899             dec.flush(drop_audio=True)
900             silence = 0.0
901             last_step = now
902             continue
903 
904         if buffered >= args.max_seg:
905             dec.step(cap.sr)
906             if dec.prev:
907                 dec.sentence.extend(dec.prev)
908                 dec.prev = []
909             dec.flush(drop_audio=True)
910             last_step = now
911             continue
912 
913         if (dec.speech >= args.min_speech
914                 and now - last_step >= args.stream_interval):
915             last_step = now
916             try:
917                 dec.step(cap.sr)
918             except Exception as e:
919                 log("stream decode error:", repr(e))
920 
921 
922 def segmenter(cap, tr, jobs, args, stop):
923     dur = cap.blocksize / cap.sr
924     noise = 1e-4
925     preroll = deque(maxlen=max(1, int(0.35 / dur)))
926     seg = []
927     speech_blocks = 0
928     silence = 0.0
929     last_partial = 0.0
930     last_meter = 0.0
931 
932     while not stop.is_set():
933         try:
934             blk = cap.q.get(timeout=0.3)
935         except queue.Empty:
936             continue
937 
938         rms = float(np.sqrt(np.mean(blk * blk)) + 1e-12)
939         if rms < noise:
940             noise = 0.90 * noise + 0.10 * rms
941         else:
942             noise = 0.995 * noise + 0.005 * rms
943         gate = max(noise * args.vad_ratio, args.vad_floor)
944         voiced = rms > gate
945 
946         if args.meter and time.time() - last_meter > 0.25:
947             last_meter = time.time()
948             db = 20 * np.log10(max(rms, 1e-9))
949             bars = int(np.clip((db + 60) / 60 * 40, 0, 40))
950             sys.stdout.write("\r%-40s %6.1f dBFS   gate %6.1f   %s"
951                              % ("#" * bars, db, 20 * np.log10(gate),
952                                 "VOICE" if voiced else "     "))
953             sys.stdout.flush()
954 
955         if not seg:
956             if not voiced:
957                 preroll.append(blk)
958                 continue
959             seg = list(preroll)
960             preroll.clear()
961 
962         seg.append(blk)
963         if voiced:
964             speech_blocks += 1
965             silence = 0.0
966         else:
967             silence += dur
968 
969         seg_dur = len(seg) * dur
970         speech_dur = speech_blocks * dur
971 
972         if speech_dur >= args.min_speech and (silence >= args.pause or seg_dur >= args.max_seg):
973             if args.verbose:
974                 log("segment queued: %.1fs (%.1fs of it speech)" % (seg_dur, speech_dur))
975             try:
976                 jobs.put_nowait(("final", np.concatenate(seg), cap.sr))
977             except queue.Full:
978                 log("backlog full - dropping a segment (model too slow for live)")
979             seg, speech_blocks, silence = [], 0, 0.0
980         elif silence >= args.pause:
981             seg, speech_blocks, silence = [], 0, 0.0          # noise blip
982         elif (args.partials and speech_dur >= 0.7
983               and time.time() - last_partial >= args.partial_every
984               and not tr.busy.is_set() and jobs.empty()):
985             last_partial = time.time()
986             try:
987                 jobs.put_nowait(("partial", np.concatenate(seg), cap.sr))
988             except queue.Full:
989                 pass
990 
991 
992 # -------------------------------------------------------------- selftest ----
993 
994 def selftest(args, seconds):
995     dev = resolve_device(args.audio_device, args.loopback)
996     with Capture(dev, args.loopback) as cap:
997         log("recording %ds -- play / speak Finnish audio NOW..." % seconds)
998         blocks, t0, peak = [], time.time(), 0.0
999         while time.time() - t0 < seconds:
1000             try:
1001                 b = cap.q.get(timeout=0.5)
1002             except queue.Empty:
1003                 continue
1004             blocks.append(b)
1005             peak = max(peak, float(np.max(np.abs(b))))
1006         sr = cap.sr
1007 
1008     if not blocks:
1009         log("NO AUDIO CAPTURED -- wrong device? Run --list-devices.")
1010         return 1
1011     audio = np.concatenate(blocks)
1012     log("captured %.1fs, peak %.1f dBFS" % (len(audio) / sr, 20 * np.log10(max(peak, 1e-9))))
1013     if peak < 0.001:
1014         log("WARNING: that is silence. Pick another device with --audio-device.")
1015 
1016     a16 = to_whisper(audio, sr)
1017     with wave.open(str(HERE / "selftest.wav"), "wb") as w:
1018         w.setnchannels(1)
1019         w.setsampwidth(2)
1020         w.setframerate(TARGET_SR)
1021         w.writeframes((a16 * 32767).astype(np.int16).tobytes())
1022     log("wrote selftest.wav -- play it to confirm you grabbed the right audio")
1023 
1024     t = Transcriber(args, Bus(args))
1025     t0 = time.time()
1026     text, nsp, info = t.run(a16, final=True)
1027     took = time.time() - t0
1028     rtf = took / max(len(a16) / TARGET_SR, 0.01)
1029     log("detected language=%s  no_speech=%.2f" % (getattr(info, "language", "?"), nsp))
1030     verdict = "FAST ENOUGH for live" if rtf < 0.6 else "TOO SLOW -- use a smaller --model"
1031     log("transcribed in %.1fs  ->  RTF %.2f  (%s)" % (took, rtf, verdict))
1032     print("\n  " + (text or "(nothing recognised)") + "\n")
1033     return 0
1034 
1035 
1036 # ------------------------------------------------------------------ main ----
1037 
1038 DEFAULT_MODEL = "small"
1039 FI_MODEL_DIR = Path("build") / "models" / "fi-small-ct2"
1040 
1041 
1042 def resolve_model(requested, lang, here):
1043     """Pick the model, returning (name, message-to-log-or-None).
1044 
1045     With nothing requested, prefer the Finnish fine-tune when it has been built
1046     and Finnish is what we are transcribing: same size and speed as small, but
1047     better Finnish. An explicit --model always wins, which is why the flag
1048     defaults to None rather than to "small".
1049     """
1050     if requested is not None:
1051         return requested, None
1052     local = Path(here) / FI_MODEL_DIR
1053     if lang == "fi" and (local / "model.bin").is_file():
1054         return str(local), "using the Finnish fine-tune (override with --model small)"
1055     return DEFAULT_MODEL, None
1056 
1057 
1058 def build_parser():
1059     p = argparse.ArgumentParser(
1060         description="Live captions for OBS",
1061         formatter_class=argparse.ArgumentDefaultsHelpFormatter)
1062 
1063     g = p.add_argument_group("audio")
1064     g.add_argument("--list-devices", action="store_true", help="show devices and exit")
1065     g.add_argument("--audio-device", default=None, help="index or name substring")
1066     g.add_argument("--loopback", dest="loopback", action="store_true", default=True,
1067                    help="capture desktop output (default)")
1068     g.add_argument("--mic", dest="loopback", action="store_false",
1069                    help="capture an input device instead of desktop output")
1070     g.add_argument("--meter", action="store_true", help="print a live level meter")
1071     g.add_argument("--verbose", action="store_true", help="log VAD segment decisions")
1072 
1073     g = p.add_argument_group("model")
1074     # Default is resolved in main() so an explicit --model small is
1075     # distinguishable from not passing --model at all.
1076     g.add_argument("--model", default=None,
1077                    help="tiny|base|small|medium|large-v3|large-v3-turbo, or a local "
1078                         "CT2 dir (default: small, or the Finnish fine-tune if built)")
1079     g.add_argument("--lang", default="fi", help="fi, ru, ja, es, pt, ... or 'auto'")
1080     g.add_argument("--langs", default="fi,ru,ja,es,pt,en,auto",
1081                    help="languages offered as one-click buttons in the control panel")
1082     g.add_argument("--model-choices", dest="model_choices",
1083                    default="tiny,base,small,medium",
1084                    help="models offered in the control panel dropdown")
1085     g.add_argument("--compute-device", default="cpu", choices=["cpu", "cuda"])
1086     g.add_argument("--compute", default="int8", help="int8|int8_float32|float32|float16")
1087     g.add_argument("--threads", type=int, default=max(2, (os.cpu_count() or 8) - 2))
1088     g.add_argument("--beam", type=int, default=5)
1089     g.add_argument("--detect-every", type=float, default=6.0,
1090                    help="seconds between language-detection passes when --lang auto")
1091     g.add_argument("--detect-min-audio", type=float, default=1.6,
1092                    help="skip detection on segments shorter than this")
1093     g.add_argument("--detect-margin", type=float, default=1.3,
1094                    help="how far a new language must lead before switching")
1095     g.add_argument("--detect-hold", type=int, default=2,
1096                    help="consecutive detections needed before switching")
1097     g.add_argument("--repetition-penalty", type=float, default=1.15,
1098                    help="discourage Whisper from looping on a phrase; "
1099                         "Japanese and Chinese need this more than European "
1100                         "languages. 1.0 disables it")
1101     g.add_argument("--translate", action="store_true",
1102                    help="also emit an English translation line")
1103     g.add_argument("--no-context", action="store_true",
1104                    help="do not feed previous text back as a prompt")
1105 
1106     g = p.add_argument_group("segmentation")
1107     g.add_argument("--pause", type=float, default=0.65, help="silence (s) that closes a caption")
1108     g.add_argument("--min-speech", type=float, default=0.45, help="min speech (s) worth sending")
1109     g.add_argument("--max-seg", type=float, default=11.0, help="force a cut after this many s")
1110     g.add_argument("--vad-floor", type=float, default=0.004, help="absolute RMS gate")
1111     g.add_argument("--vad-ratio", type=float, default=3.0, help="gate = noise floor * this")
1112     g.add_argument("--partials", dest="partials", action="store_true", default=True)
1113     g.add_argument("--no-partials", dest="partials", action="store_false",
1114                    help="only show finished sentences (lower CPU)")
1115     g.add_argument("--partial-every", type=float, default=0.9)
1116     g.add_argument("--stream", action="store_true",
1117                    help="LocalAgreement streaming: words appear while someone is "
1118                         "still talking instead of after they stop. Costs a lot "
1119                         "more CPU - pair it with a smaller --model")
1120     g.add_argument("--stream-interval", type=float, default=0.8,
1121                    help="how often to re-decode the buffer in --stream mode")
1122     g.add_argument("--stream-min-audio", type=float, default=1.5,
1123                    help="do not decode until this much audio is buffered; "
1124                         "shorter buffers make Whisper hallucinate")
1125 
1126     g = p.add_argument_group("output")
1127     g.add_argument("--ws-port", type=int, default=8765)
1128     g.add_argument("--http-port", type=int, default=8777)
1129     g.add_argument("--txt", default="captions.txt", help="plain-text file for a GDI+ text source")
1130     g.add_argument("--txt-lines", type=int, default=2)
1131     g.add_argument("--lines", type=int, default=2, help="lines shown in the overlay")
1132     g.add_argument("--size", type=int, default=42, help="overlay font size in px")
1133     g.add_argument("--hide", type=float, default=8, help="auto-hide overlay after N idle seconds")
1134 
1135     p.add_argument("--selftest", nargs="?", type=int, const=12, default=None, metavar="SECONDS",
1136                    help="record N seconds, transcribe once, report speed, then exit")
1137     return p
1138 
1139 
1140 def main():
1141     args = build_parser().parse_args()
1142 
1143     chosen, why = resolve_model(args.model, args.lang, HERE)
1144     args.model = chosen
1145     if why:
1146         log(why)
1147 
1148     args.langs = [s.strip() for s in args.langs.split(",") if s.strip()]
1149     args.model_choices = [s.strip() for s in args.model_choices.split(",") if s.strip()]
1150     if args.model not in args.model_choices:
1151         args.model_choices.insert(0, args.model)
1152     if args.lang not in args.langs:
1153         args.langs.insert(0, args.lang)
1154 
1155     if args.list_devices:
1156         list_devices()
1157         return 0
1158     if args.selftest is not None:
1159         return selftest(args, args.selftest)
1160 
1161     bus = Bus(args)
1162     stop_ev = threading.Event()
1163     jobs = queue.Queue(maxsize=3)
1164 
1165     bus.mark_session()
1166     ctl = Controller(args, bus)
1167     tr = Transcriber(args, bus)
1168     ctl.tr = tr
1169     dev = resolve_device(args.audio_device, args.loopback)
1170 
1171     srv = http_server(args)
1172     loop = asyncio.new_event_loop()
1173     bus.attach(loop)
1174     ws_stop = asyncio.Event()
1175 
1176     def run_loop():
1177         asyncio.set_event_loop(loop)
1178         loop.run_until_complete(ws_server(bus, args, ctl, ws_stop))
1179 
1180     threading.Thread(target=run_loop, daemon=True).start()
1181     if not args.stream:
1182         threading.Thread(target=tr.worker, args=(jobs, stop_ev, ctl),
1183                          daemon=True).start()
1184 
1185     try:
1186         with Capture(dev, args.loopback) as cap:
1187             log("lang=%s  model=%s  -- Ctrl+C to stop" % (args.lang, args.model))
1188             if args.stream:
1189                 stream_segmenter(cap, tr, args, stop_ev, bus)
1190             else:
1191                 segmenter(cap, tr, jobs, args, stop_ev)
1192     except KeyboardInterrupt:
1193         print()
1194         log("stopping")
1195     finally:
1196         stop_ev.set()
1197         loop.call_soon_threadsafe(ws_stop.set)
1198         srv.shutdown()
1199     return 0
1200 
1201 
1202 if __name__ == "__main__":
1203     sys.exit(main())