Support openWakeWord

linux_voice_assistant/__main__.py

@@ -4,15 +4,18 @@ import asyncio
 import json
 import logging
 import threading
+import time
 from pathlib import Path
 from queue import Queue
-from typing import Dict, List
+from typing import Dict, List, Optional, Union
 
+import numpy as np
 import sounddevice as sd
 
-from .microwakeword import MicroWakeWord
-from .models import AvailableWakeWord, Preferences, ServerState
+from .microwakeword import MicroWakeWord, MicroWakeWordFeatures
+from .models import AvailableWakeWord, Preferences, ServerState, WakeWordType
 from .mpv_player import MpvMediaPlayer
+from .openwakeword import OpenWakeWord, OpenWakeWordFeatures
 from .satellite import VoiceSatelliteProtocol
 from .util import get_mac, is_arm
 
@@ -20,6 +23,7 @@ _LOGGER = logging.getLogger(__name__)
 _MODULE_DIR = Path(__file__).parent
 _REPO_DIR = _MODULE_DIR.parent
 _WAKEWORDS_DIR = _REPO_DIR / "wakewords"
+_OWW_DIR = _WAKEWORDS_DIR / "openWakeWord"
 _SOUNDS_DIR = _REPO_DIR / "sounds"
 
 if is_arm():
@@ -43,13 +47,32 @@ async def main() -> None:
     parser.add_argument("--audio-output-device", help="mpv name for output device")
     parser.add_argument(
         "--wake-word-dir",
-        default=_WAKEWORDS_DIR,
+        default=[_WAKEWORDS_DIR],
+        action="append",
         help="Directory with wake word models (.tflite) and configs (.json)",
     )
     parser.add_argument(
         "--wake-model", default="okay_nabu", help="Id of active wake model"
     )
     parser.add_argument("--stop-model", default="stop", help="Id of stop model")
+    parser.add_argument(
+        "--refractory-seconds",
+        default=2.0,
+        type=float,
+        help="Seconds before wake word can be activated again",
+    )
+    #
+    parser.add_argument(
+        "--oww-melspectrogram-model",
+        default=_OWW_DIR / "melspectrogram.tflite",
+        help="Path to openWakeWord melspectrogram model",
+    )
+    parser.add_argument(
+        "--oww-embedding-model",
+        default=_OWW_DIR / "embedding_model.tflite",
+        help="Path to openWakeWord embedding model",
+    )
+    #
     parser.add_argument(
         "--wakeup-sound", default=str(_SOUNDS_DIR / "wake_word_triggered.flac")
     )
@@ -76,22 +99,26 @@ async def main() -> None:
     _LOGGER.debug(args)
 
     # Load available wake words
-    wake_word_dir = Path(args.wake_word_dir)
+    wake_word_dirs = [Path(ww_dir) for ww_dir in args.wake_word_dir]
     available_wake_words: Dict[str, AvailableWakeWord] = {}
-    for model_config_path in wake_word_dir.glob("*.json"):
-        model_id = model_config_path.stem
-        if model_id == args.stop_model:
-            # Don't show stop model as an available wake word
-            continue
 
-        with open(model_config_path, "r", encoding="utf-8") as model_config_file:
-            model_config = json.load(model_config_file)
-            available_wake_words[model_id] = AvailableWakeWord(
-                id=model_id,
-                wake_word=model_config["wake_word"],
-                trained_languages=model_config.get("trained_languages", []),
-                config_path=model_config_path,
-            )
+    for wake_word_dir in wake_word_dirs:
+        for model_config_path in wake_word_dir.glob("*.json"):
+            model_id = model_config_path.stem
+            if model_id == args.stop_model:
+                # Don't show stop model as an available wake word
+                continue
+
+            with open(model_config_path, "r", encoding="utf-8") as model_config_file:
+                model_config = json.load(model_config_file)
+                model_type = model_config["type"]
+                available_wake_words[model_id] = AvailableWakeWord(
+                    id=model_id,
+                    type=WakeWordType(model_type),
+                    wake_word=model_config["wake_word"],
+                    trained_languages=model_config.get("trained_languages", []),
+                    config_path=model_config_path,
+                )
 
     _LOGGER.debug("Available wake words: %s", list(sorted(available_wake_words.keys())))
 
@@ -109,31 +136,40 @@ async def main() -> None:
     _LOGGER.debug("libtensorflowlite_c path: %s", libtensorflowlite_c_path)
 
     # Load wake/stop models
-    wake_models: Dict[str, MicroWakeWord] = {}
-    wake_config_path = wake_word_dir / f"{args.wake_model}.json"
+    wake_models: Dict[str, Union[MicroWakeWord, OpenWakeWord]] = {}
     if preferences.active_wake_words:
         # Load preferred models
         for wake_word_id in preferences.active_wake_words:
-            wake_config_path = wake_word_dir / f"{wake_word_id}.json"
-            if not wake_config_path.exists():
+            wake_word = available_wake_words.get(wake_word_id)
+            if wake_word is None:
+                _LOGGER.warning("Unrecognized wake word id: %s", wake_word_id)
                 continue
 
-            _LOGGER.debug("Loading wake model: %s", wake_config_path)
-            wake_models[wake_word_id] = MicroWakeWord.from_config(
-                wake_config_path, libtensorflowlite_c_path
-            )
+            _LOGGER.debug("Loading wake model: %s", wake_word_id)
+            wake_models[wake_word_id] = wake_word.load(libtensorflowlite_c_path)
 
     if not wake_models:
         # Load default model
-        _LOGGER.debug("Loading wake model: %s", wake_config_path)
-        wake_word_id = wake_config_path.stem
-        wake_models[wake_word_id] = MicroWakeWord.from_config(
-            wake_config_path, libtensorflowlite_c_path
-        )
+        wake_word_id = args.wake_model
+        wake_word = available_wake_words[wake_word_id]
 
-    stop_config_path = wake_word_dir / f"{args.stop_model}.json"
-    _LOGGER.debug("Loading stop model: %s", stop_config_path)
-    stop_model = MicroWakeWord.from_config(stop_config_path, libtensorflowlite_c_path)
+        _LOGGER.debug("Loading wake model: %s", wake_word_id)
+        wake_models[wake_word_id] = wake_word.load(libtensorflowlite_c_path)
+
+    # TODO: allow openWakeWord for "stop"
+    stop_model: Optional[MicroWakeWord] = None
+    for wake_word_dir in wake_word_dirs:
+        stop_config_path = wake_word_dir / f"{args.stop_model}.json"
+        if not stop_config_path.exists():
+            continue
+
+        _LOGGER.debug("Loading stop model: %s", stop_config_path)
+        stop_model = MicroWakeWord.from_config(
+            stop_config_path, libtensorflowlite_c_path
+        )
+        break
+
+    assert stop_model is not None
 
     state = ServerState(
         name=args.name,
@@ -150,6 +186,9 @@ async def main() -> None:
         preferences=preferences,
         preferences_path=preferences_path,
         libtensorflowlite_c_path=libtensorflowlite_c_path,
+        oww_melspectrogram_path=Path(args.oww_melspectrogram_model),
+        oww_embedding_path=Path(args.oww_embedding_model),
+        refractory_seconds=args.refractory_seconds,
     )
 
     process_audio_thread = threading.Thread(
@@ -193,7 +232,15 @@ async def main() -> None:
 def process_audio(state: ServerState):
     """Process audio chunks from the microphone."""
 
-    wake_words: List[MicroWakeWord] = []
+    wake_words: List[Union[MicroWakeWord, OpenWakeWord]] = []
+    micro_features: Optional[MicroWakeWordFeatures] = None
+    micro_inputs: List[np.ndarray] = []
+
+    oww_features: Optional[OpenWakeWordFeatures] = None
+    oww_inputs: List[np.ndarray] = []
+    has_oww = False
+
+    last_active: Optional[float] = None
 
     try:
         while True:
@@ -209,19 +256,66 @@ def process_audio(state: ServerState):
                 state.wake_words_changed = False
                 wake_words = list(state.wake_words.values())
 
+                has_oww = False
+                for wake_word in wake_words:
+                    if isinstance(wake_word, OpenWakeWord):
+                        has_oww = True
+
+                if micro_features is None:
+                    micro_features = MicroWakeWordFeatures(
+                        libtensorflowlite_c_path=state.libtensorflowlite_c_path,
+                    )
+
+                if has_oww and (oww_features is None):
+                    oww_features = OpenWakeWordFeatures(
+                        melspectrogram_model=state.oww_melspectrogram_path,
+                        embedding_model=state.oww_embedding_path,
+                        libtensorflowlite_c_path=state.libtensorflowlite_c_path,
+                    )
+
             try:
                 state.satellite.handle_audio(audio_chunk)
 
+                assert micro_features is not None
+                micro_inputs.clear()
+                micro_inputs.extend(micro_features.process_streaming(audio_chunk))
+
+                if has_oww:
+                    assert oww_features is not None
+                    oww_inputs.clear()
+                    oww_inputs.extend(oww_features.process_streaming(audio_chunk))
+
                 for wake_word in wake_words:
                     if not wake_word.is_active:
                         continue
 
-                    if wake_word.process_streaming(audio_chunk):
-                        state.satellite.wakeup(wake_word)
+                    activated = False
+                    if isinstance(wake_word, MicroWakeWord):
+                        for oww_input in oww_inputs:
+                            if wake_word.process_streaming(oww_input):
+                                activated = True
+                    elif isinstance(wake_word, OpenWakeWord):
+                        for oww_input in oww_inputs:
+                            for prob in wake_word.process_streaming(oww_input):
+                                if prob > 0.5:
+                                    activated = True
 
-                if state.stop_word.is_active and state.stop_word.process_streaming(
-                    audio_chunk
-                ):
+                    if activated:
+                        # Check refractory
+                        now = time.monotonic()
+                        if (last_active is None) or (
+                            (now - last_active) > state.refractory_seconds
+                        ):
+                            state.satellite.wakeup(wake_word)
+                            last_active = now
+
+                # Always process to keep state correct
+                stopped = False
+                for micro_input in micro_inputs:
+                    if state.stop_word.process_streaming(micro_input):
+                        stopped = True
+
+                if stopped and state.stop_word.is_active:
                     state.satellite.stop()
             except Exception:
                 _LOGGER.exception("Unexpected error handling audio")