186 lines
7.0 KiB
Python
186 lines
7.0 KiB
Python
"""Wake-word detection via a custom-trained openWakeWord model.
|
|
|
|
Uses `thunderbolt.onnx` — trained specifically for "thunderbolt", the same
|
|
way the main repo's `desk_client/bolt_desk.py` uses `bolt.onnx` for "hey
|
|
bolt". Same runtime (openWakeWord, ONNX inference), same per-frame
|
|
predict()/reset() pattern; the only difference is the model file
|
|
(WAKE_MODEL_FILE) and threshold (WAKE_WORD_THRESHOLD), both configurable via
|
|
.env if a differently-trained model is swapped in later.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections import deque
|
|
from typing import Callable, Optional, Protocol, Union
|
|
|
|
import numpy as np
|
|
|
|
from .. import config
|
|
|
|
|
|
class WakeModel(Protocol):
|
|
def predict(self, frame: np.ndarray) -> dict: ...
|
|
def reset(self) -> None: ...
|
|
|
|
|
|
class NearMissLog:
|
|
"""Rolling record of frames that *almost* fired the wake word.
|
|
|
|
WAKE_WORD_THRESHOLD is otherwise tuned by guessing at a number in .env
|
|
and seeing whether the pet ignores you. Keeping the near misses (scores
|
|
within WAKE_NEAR_MISS_MARGIN below the threshold) turns that into
|
|
evidence: the tray's wake-word tuner shows what your actual "thunderbolt"
|
|
scores, so you can set the threshold just under it.
|
|
|
|
Pure bookkeeping — the caller supplies timestamps, so it's testable.
|
|
"""
|
|
|
|
def __init__(self, limit: int = None, margin: float = None):
|
|
self._entries: deque[tuple[float, float, float]] = deque( # (timestamp, score, threshold)
|
|
maxlen=max(1, config.WAKE_NEAR_MISS_LIMIT if limit is None else limit)
|
|
)
|
|
self._margin = config.WAKE_NEAR_MISS_MARGIN if margin is None else margin
|
|
self._peak = 0.0
|
|
|
|
@property
|
|
def peak(self) -> float:
|
|
"""Highest score seen since the last reset — the "how close did I
|
|
get?" readout while you test the wake phrase."""
|
|
return self._peak
|
|
|
|
def observe(self, score: float, threshold: float, timestamp: float) -> bool:
|
|
"""Record *score*; returns True if it counted as a near miss."""
|
|
self._peak = max(self._peak, score)
|
|
if score >= threshold or score < threshold - self._margin:
|
|
return False
|
|
self._entries.append((timestamp, score, threshold))
|
|
return True
|
|
|
|
def entries(self) -> list[tuple[float, float, float]]:
|
|
return list(self._entries)
|
|
|
|
def clear(self) -> None:
|
|
self._entries.clear()
|
|
self._peak = 0.0
|
|
|
|
|
|
def _construct_model(model_cls, model_path: str):
|
|
"""openwakeword's Model() constructor keyword has drifted across
|
|
releases (wakeword_models -> wakeword_model_paths) and some builds
|
|
reject inference_framework entirely — the main repo's ai/wake_word.py
|
|
hit the same drift and works around it the same way: try each known
|
|
calling convention in turn."""
|
|
attempts = [
|
|
lambda: model_cls(wakeword_model_paths=[model_path], inference_framework="onnx"),
|
|
lambda: model_cls(wakeword_model_paths=[model_path]),
|
|
lambda: model_cls(wakeword_models=[model_path], inference_framework="onnx"),
|
|
lambda: model_cls(wakeword_models=[model_path]),
|
|
lambda: model_cls([model_path]),
|
|
]
|
|
last_exc: Optional[TypeError] = None
|
|
for attempt in attempts:
|
|
try:
|
|
return attempt()
|
|
except TypeError as exc:
|
|
last_exc = exc
|
|
raise RuntimeError(
|
|
f"Could not construct openwakeword.Model with any known calling convention "
|
|
f"(last error: {last_exc})"
|
|
)
|
|
|
|
|
|
class _OpenWakeWordModel:
|
|
"""Lazily loads the ONNX model on first use so importing this module
|
|
(and unit-testing listen_for_wake_word with a fake model) never requires
|
|
onnxruntime/openwakeword or the model file to be present."""
|
|
|
|
def __init__(self):
|
|
self._model = None
|
|
|
|
def _ensure_model(self):
|
|
if self._model is None:
|
|
from openwakeword.model import Model
|
|
# from openwakeword.utils import download_models
|
|
|
|
# # The pip package doesn't bundle its melspectrogram/embedding
|
|
# # feature-extraction sub-models — fetch them once on first use
|
|
# # (no-op if already cached in openwakeword's own resources dir).
|
|
# # A non-empty, non-matching model_names list keeps this from
|
|
# # also pulling every official pretrained wakeword model.
|
|
# download_models(model_names=["thunderbolt"])
|
|
|
|
self._model = _construct_model(Model, config.WAKE_MODEL_PATH)
|
|
return self._model
|
|
|
|
def predict(self, frame: np.ndarray) -> dict:
|
|
return self._ensure_model().predict(frame)
|
|
|
|
def reset(self) -> None:
|
|
if self._model is not None:
|
|
self._model.reset()
|
|
|
|
|
|
_default_model = _OpenWakeWordModel()
|
|
|
|
|
|
def listen_for_wake_word(
|
|
stream,
|
|
should_continue=lambda: True,
|
|
model: Optional[WakeModel] = None,
|
|
threshold: Union[float, Callable[[], float], None] = None,
|
|
on_tick=None,
|
|
on_score: Optional[Callable[[float, float], None]] = None,
|
|
) -> bool:
|
|
"""Block until the wake word fires (returns True) or *should_continue*
|
|
goes false (returns False).
|
|
|
|
Feeds every frame to *model* (the thunderbolt openWakeWord model by
|
|
default) and treats any class score >= *threshold* as a detection,
|
|
resetting the model's internal state afterward so the next call starts
|
|
clean — same pattern as desk_client/bolt_desk.py's main loop.
|
|
|
|
*threshold* may be a number or a zero-argument callable. The callable
|
|
form exists because this function blocks for minutes at a time: the
|
|
tray's wake-word tuner slider has to be able to change sensitivity
|
|
*during* a listen, not only at the start of the next one.
|
|
|
|
*on_tick*, if given, is called once per ``WAKE_CHECK_INTERVAL_SECONDS``
|
|
(not every frame — prediction is cheap enough to run on every frame, but
|
|
this is the only point control returns to the caller while otherwise
|
|
blocked here for a possibly long time, so it's how a caller drives
|
|
periodic work, e.g. the heartbeat/announcement poll in controller.py,
|
|
during quiet stretches with no wake word).
|
|
|
|
*on_score*, if given, gets ``(best_score, threshold)`` every frame — used
|
|
to log near misses for threshold tuning.
|
|
"""
|
|
model = model or _default_model
|
|
if threshold is None:
|
|
threshold = config.WAKE_WORD_THRESHOLD
|
|
resolve_threshold = threshold if callable(threshold) else (lambda: threshold)
|
|
|
|
frame_len = config.FRAME_LEN
|
|
check_every_frames = max(1, int(config.WAKE_CHECK_INTERVAL_SECONDS * config.SAMPLE_RATE / frame_len))
|
|
frames_since_tick = 0
|
|
|
|
while should_continue():
|
|
chunk, _ = stream.read(frame_len)
|
|
frame = np.asarray(chunk)[:, 0]
|
|
scores = model.predict(frame)
|
|
current_threshold = resolve_threshold()
|
|
best = max(scores.values()) if scores else 0.0
|
|
|
|
frames_since_tick += 1
|
|
if frames_since_tick >= check_every_frames:
|
|
frames_since_tick = 0
|
|
if on_tick is not None:
|
|
on_tick()
|
|
|
|
if on_score is not None:
|
|
on_score(best, current_threshold)
|
|
|
|
if scores and best >= current_threshold:
|
|
model.reset()
|
|
return True
|
|
return False
|