diff --git a/CHANGELOG.md b/CHANGELOG.md index 26ae6a4..fcd2098 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,7 @@ ## Unreleased +- Support openWakeWord - Support multiple wake words - Save active wake words to preferences JSON file - Refactor main into separate files diff --git a/linux_voice_assistant/__main__.py b/linux_voice_assistant/__main__.py index b9b26db..f49c431 100644 --- a/linux_voice_assistant/__main__.py +++ b/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") diff --git a/linux_voice_assistant/microwakeword.py b/linux_voice_assistant/microwakeword.py index a709311..79b4067 100644 --- a/linux_voice_assistant/microwakeword.py +++ b/linux_voice_assistant/microwakeword.py @@ -2,26 +2,24 @@ import ctypes import json import statistics from collections import deque +from collections.abc import Iterable from pathlib import Path from typing import Deque, List, Union import numpy as np from pymicro_features import MicroFrontend -_SAMPLES_PER_SECOND = 16000 -_SAMPLES_PER_CHUNK = 160 # 10ms -_BYTES_PER_SAMPLE = 2 # 16-bit -_BYTES_PER_CHUNK = _SAMPLES_PER_CHUNK * _BYTES_PER_SAMPLE -_SECONDS_PER_CHUNK = _SAMPLES_PER_CHUNK / _SAMPLES_PER_SECOND -_STRIDE = 3 -_DEFAULT_REFRACTORY = 2 # seconds +from .wakeword import TfLiteWakeWord + +SAMPLES_PER_SECOND = 16000 +SAMPLES_PER_CHUNK = 160 # 10ms +BYTES_PER_SAMPLE = 2 # 16-bit +BYTES_PER_CHUNK = SAMPLES_PER_CHUNK * BYTES_PER_SAMPLE +SECONDS_PER_CHUNK = SAMPLES_PER_CHUNK / SAMPLES_PER_SECOND +STRIDE = 3 -class TfLiteQuantizationParams(ctypes.Structure): - _fields_ = [("scale", ctypes.c_float), ("zero_point", ctypes.c_int32)] - - -class MicroWakeWord: +class MicroWakeWord(TfLiteWakeWord): def __init__( self, id: str, # pylint: disable=redefined-builtin @@ -29,76 +27,27 @@ class MicroWakeWord: tflite_model: Union[str, Path], probability_cutoff: float, sliding_window_size: int, - refractory_seconds: float, trained_languages: List[str], libtensorflowlite_c_path: Union[str, Path], - ): + ) -> None: + TfLiteWakeWord.__init__(self, libtensorflowlite_c_path) + self.id = id self.wake_word = wake_word self.tflite_model = tflite_model self.probability_cutoff = probability_cutoff self.sliding_window_size = sliding_window_size - self.refractory_seconds = refractory_seconds self.trained_languages = trained_languages self.is_active = True - # Load the shared library - self.libtensorflowlite_c_path = Path(libtensorflowlite_c_path) - self.lib = ctypes.cdll.LoadLibrary(str(self.libtensorflowlite_c_path.resolve())) - - # Define required argument/return types for C API - self.lib.TfLiteModelCreateFromFile.argtypes = [ctypes.c_char_p] - self.lib.TfLiteModelCreateFromFile.restype = ctypes.c_void_p - - self.lib.TfLiteInterpreterCreate.argtypes = [ctypes.c_void_p, ctypes.c_void_p] - self.lib.TfLiteInterpreterCreate.restype = ctypes.c_void_p - - self.lib.TfLiteInterpreterAllocateTensors.argtypes = [ctypes.c_void_p] - self.lib.TfLiteInterpreterInvoke.argtypes = [ctypes.c_void_p] - - self.lib.TfLiteInterpreterGetInputTensor.argtypes = [ - ctypes.c_void_p, - ctypes.c_int, - ] - self.lib.TfLiteInterpreterGetInputTensor.restype = ctypes.c_void_p - - self.lib.TfLiteInterpreterGetOutputTensor.argtypes = [ - ctypes.c_void_p, - ctypes.c_int, - ] - self.lib.TfLiteInterpreterGetOutputTensor.restype = ctypes.c_void_p - - self.lib.TfLiteTensorCopyFromBuffer.argtypes = [ - ctypes.c_void_p, - ctypes.c_void_p, - ctypes.c_size_t, - ] - self.lib.TfLiteTensorCopyToBuffer.argtypes = [ - ctypes.c_void_p, - ctypes.c_void_p, - ctypes.c_size_t, - ] - - self.lib.TfLiteTensorType.restype = ctypes.c_int - - self.lib.TfLiteTensorByteSize.argtypes = [ctypes.c_void_p] - self.lib.TfLiteTensorByteSize.restype = ctypes.c_size_t - - self.lib.TfLiteTensorQuantizationParams.restype = ctypes.Structure - - self.lib.TfLiteTensorQuantizationParams.argtypes = [ctypes.c_void_p] - self.lib.TfLiteTensorQuantizationParams.restype = TfLiteQuantizationParams - # Load the model and create interpreter self.model_path = str(Path(tflite_model).resolve()).encode("utf-8") self._load_model() - self._frontend = MicroFrontend() self._features: List[np.ndarray] = [] self._probabilities: Deque[float] = deque(maxlen=self.sliding_window_size) self._audio_buffer = bytes() - self._ignore_seconds: float = 0 def _load_model(self) -> None: self.model = self.lib.TfLiteModelCreateFromFile(self.model_path) @@ -120,41 +69,10 @@ class MicroWakeWord: self.input_scale, self.input_zero_point = input_q.scale, input_q.zero_point self.output_scale, self.output_zero_point = output_q.scale, output_q.zero_point - def reset(self) -> None: - """Reload model and clear state. - - This must be done between audio clips when not streaming. - """ - self._audio_buffer = bytes() - self._features.clear() - self._probabilities.clear() - self._ignore_seconds = 0 - - # Clear out residual features - self._frontend = MicroFrontend() - - # Need to reload model to reset intermediary results - # reset_all_variables() doesn't work. - self._load_model() - - @property - def samples_per_chunk(self) -> int: - """Number of samples in a streaming audio chunk.""" - return _SAMPLES_PER_CHUNK - - @property - def bytes_per_chunk(self) -> int: - """Number of bytes in a streaming audio chunk. - - Assumes 16-bit mono samples at 16Khz. - """ - return _BYTES_PER_CHUNK - @staticmethod def from_config( config_path: Union[str, Path], libtensorflowlite_c_path: Union[str, Path], - refractory_seconds: float = _DEFAULT_REFRACTORY, ) -> "MicroWakeWord": """Load a microWakeWord model from a JSON config file. @@ -162,8 +80,6 @@ class MicroWakeWord: ---------- config_path: str or Path Path to JSON configuration file - refractory_seconds: float - Number of seconds to ignore after detection """ config_path = Path(config_path) with open(config_path, "r", encoding="utf-8") as config_file: @@ -177,96 +93,97 @@ class MicroWakeWord: tflite_model=config_path.parent / config["model"], probability_cutoff=micro_config["probability_cutoff"], sliding_window_size=micro_config["sliding_window_size"], - refractory_seconds=refractory_seconds, trained_languages=micro_config.get("trained_languages", []), libtensorflowlite_c_path=libtensorflowlite_c_path, ) - def process_streaming(self, audio_bytes: bytes) -> bool: - """Process a chunk of audio in streaming mode. + def process_streaming(self, features: np.ndarray) -> bool: + self._features.append(features) - Parameters - ---------- - audio_bytes: bytes - Raw 16-bit mono audio samples at 16Khz - - Returns True if the wake word was detected. - """ - self._audio_buffer += audio_bytes - - if len(self._audio_buffer) < _BYTES_PER_CHUNK: - # Not enough audio to get features + if len(self._features) < STRIDE: + # Not enough windows return False - detected = False + # Allocate and quantize input data + quant_features = np.round( + np.concatenate(self._features, axis=1) / self.input_scale + + self.input_zero_point + ).astype(np.uint8) + + # Stride instead of rolling + self._features.clear() + + # Set tensor + quant_ptr = quant_features.ctypes.data_as(ctypes.c_void_p) + self.lib.TfLiteTensorCopyFromBuffer( + self.input_tensor, quant_ptr, quant_features.nbytes + ) + + # Run inference + self.lib.TfLiteInterpreterInvoke(self.interpreter) + + # Read output + output_bytes = self.lib.TfLiteTensorByteSize(self.output_tensor) + output_data = np.empty(output_bytes, dtype=np.uint8) + self.lib.TfLiteTensorCopyToBuffer( + self.output_tensor, + output_data.ctypes.data_as(ctypes.c_void_p), + output_bytes, + ) + + # Dequantize output + result = ( + output_data.astype(np.float32) - self.output_zero_point + ) * self.output_scale + + self._probabilities.append(result.item()) + + if len(self._probabilities) < self.sliding_window_size: + # Not enough probabilities + return False + + if statistics.mean(self._probabilities) > self.probability_cutoff: + return True + + return False + + +# ----------------------------------------------------------------------------- + + +class MicroWakeWordFeatures(TfLiteWakeWord): + def __init__( + self, + libtensorflowlite_c_path: Union[str, Path], + ) -> None: + TfLiteWakeWord.__init__(self, libtensorflowlite_c_path) + + self._audio_buffer = bytes() + self._frontend = MicroFrontend() + + def process_streaming(self, audio_bytes: bytes) -> Iterable[np.ndarray]: + self._audio_buffer += audio_bytes + + if len(self._audio_buffer) < BYTES_PER_CHUNK: + # Not enough audio to get features + return + audio_buffer_idx = 0 - while (audio_buffer_idx + _BYTES_PER_CHUNK) <= len(self._audio_buffer): + while (audio_buffer_idx + BYTES_PER_CHUNK) <= len(self._audio_buffer): # Process chunk chunk_bytes = self._audio_buffer[ - audio_buffer_idx : audio_buffer_idx + _BYTES_PER_CHUNK + audio_buffer_idx : audio_buffer_idx + BYTES_PER_CHUNK ] frontend_result = self._frontend.ProcessSamples(chunk_bytes) - audio_buffer_idx += frontend_result.samples_read * _BYTES_PER_SAMPLE - self._ignore_seconds = max(0, self._ignore_seconds - _SECONDS_PER_CHUNK) + audio_buffer_idx += frontend_result.samples_read * BYTES_PER_SAMPLE if not frontend_result.features: # Not enough audio for a full window continue - self._features.append( - np.array(frontend_result.features).reshape( - (1, 1, len(frontend_result.features)) - ) + yield np.array(frontend_result.features).reshape( + (1, 1, len(frontend_result.features)) ) - if len(self._features) < _STRIDE: - # Not enough windows - continue - - # Allocate and quantize input data - quant_features = np.round( - np.concatenate(self._features, axis=1) / self.input_scale - + self.input_zero_point - ).astype(np.uint8) - - # Stride instead of rolling - self._features.clear() - - # Set tensor - quant_ptr = quant_features.ctypes.data_as(ctypes.c_void_p) - self.lib.TfLiteTensorCopyFromBuffer( - self.input_tensor, quant_ptr, quant_features.nbytes - ) - - # Run inference - self.lib.TfLiteInterpreterInvoke(self.interpreter) - - # Read output - output_bytes = self.lib.TfLiteTensorByteSize(self.output_tensor) - output_data = np.empty(output_bytes, dtype=np.uint8) - self.lib.TfLiteTensorCopyToBuffer( - self.output_tensor, - output_data.ctypes.data_as(ctypes.c_void_p), - output_bytes, - ) - - # Dequantize output - result = ( - output_data.astype(np.float32) - self.output_zero_point - ) * self.output_scale - - self._probabilities.append(result.item()) - - if len(self._probabilities) < self.sliding_window_size: - # Not enough probabilities - continue - - if statistics.mean(self._probabilities) > self.probability_cutoff: - if self._ignore_seconds <= 0: - detected = True - self._ignore_seconds = self.refractory_seconds - # Remove processed audio self._audio_buffer = self._audio_buffer[audio_buffer_idx:] - - return detected diff --git a/linux_voice_assistant/models.py b/linux_voice_assistant/models.py index f6c47ad..4ca4f9e 100644 --- a/linux_voice_assistant/models.py +++ b/linux_voice_assistant/models.py @@ -3,26 +3,55 @@ import json import logging from dataclasses import asdict, dataclass, field +from enum import Enum from pathlib import Path from queue import Queue -from typing import TYPE_CHECKING, Dict, List, Optional +from typing import TYPE_CHECKING, Dict, List, Optional, Union if TYPE_CHECKING: from .entity import ESPHomeEntity, MediaPlayerEntity from .microwakeword import MicroWakeWord from .mpv_player import MpvMediaPlayer + from .openwakeword import OpenWakeWord from .satellite import VoiceSatelliteProtocol _LOGGER = logging.getLogger(__name__) +class WakeWordType(str, Enum): + MICRO_WAKE_WORD = "micro" + OPEN_WAKE_WORD = "openWakeWord" + + @dataclass class AvailableWakeWord: id: str + type: WakeWordType wake_word: str trained_languages: List[str] config_path: Path + def load( + self, libtensorflowlite_c_path: Path + ) -> "Union[MicroWakeWord, OpenWakeWord]": + if self.type == WakeWordType.MICRO_WAKE_WORD: + from .microwakeword import MicroWakeWord + + return MicroWakeWord.from_config( + config_path=self.config_path, + libtensorflowlite_c_path=libtensorflowlite_c_path, + ) + + if self.type == WakeWordType.OPEN_WAKE_WORD: + from .openwakeword import OpenWakeWord + + return OpenWakeWord.from_config( + config_path=self.config_path, + libtensorflowlite_c_path=libtensorflowlite_c_path, + ) + + raise ValueError(f"Unexpected wake word type: {self.type}") + @dataclass class Preferences: @@ -36,7 +65,7 @@ class ServerState: audio_queue: "Queue[Optional[bytes]]" entities: "List[ESPHomeEntity]" available_wake_words: "Dict[str, AvailableWakeWord]" - wake_words: "Dict[str, MicroWakeWord]" + wake_words: "Dict[str, Union[MicroWakeWord, OpenWakeWord]]" stop_word: "MicroWakeWord" music_player: "MpvMediaPlayer" tts_player: "MpvMediaPlayer" @@ -45,9 +74,15 @@ class ServerState: preferences: Preferences preferences_path: Path libtensorflowlite_c_path: Path + + # openWakeWord + oww_melspectrogram_path: Path + oww_embedding_path: Path + media_player_entity: "Optional[MediaPlayerEntity]" = None satellite: "Optional[VoiceSatelliteProtocol]" = None wake_words_changed: bool = False + refractory_seconds: float = 2.0 def save_preferences(self) -> None: """Save preferences as JSON.""" diff --git a/linux_voice_assistant/openwakeword.py b/linux_voice_assistant/openwakeword.py new file mode 100644 index 0000000..317ea74 --- /dev/null +++ b/linux_voice_assistant/openwakeword.py @@ -0,0 +1,293 @@ +"""openWakeWord implementation.""" + +import ctypes as C +import json +from collections.abc import Iterable +from pathlib import Path +from typing import Final, Union + +import numpy as np + +from .wakeword import TfLiteWakeWord + +c_void_p = C.c_void_p +c_int32 = C.c_int32 +c_size_t = C.c_size_t + +BATCH_SIZE: Final = 1 + +AUTOFILL_SECONDS: Final = 8 +MAX_SECONDS: Final = 10 + +SAMPLE_RATE: Final = 16000 # 16Khz +_MAX_SAMPLES: Final = MAX_SECONDS * SAMPLE_RATE + +SAMPLES_PER_CHUNK: Final = 1280 # 80 ms @ 16Khz +MS_PER_CHUNK: Final = SAMPLES_PER_CHUNK // SAMPLE_RATE + +# window = 400, hop length = 160 +MELS_PER_SECOND: Final = 97 +MAX_MELS: Final = MAX_SECONDS * MELS_PER_SECOND +MEL_SAMPLES: Final = 1760 +NUM_MELS: Final = 32 + +EMB_FEATURES: Final = 76 # 775 ms +EMB_STEP: Final = 8 +MAX_EMB: Final = MAX_SECONDS * EMB_STEP +WW_FEATURES: Final = 96 + +MEL_SHAPE: Final = (BATCH_SIZE, MEL_SAMPLES) +EMB_SHAPE: Final = (BATCH_SIZE, EMB_FEATURES, NUM_MELS, 1) + +# melspec = [batch x samples (min: 1280)] => [batch x 1 x window x mels (32)] +# stft window size: 25ms (400) +# stft window step: 10ms (160) +# mel band limits: 60Hz - 3800Hz +# mel frequency bins: 32 +# +# embedding = [batch x window x mels (32) x 1] => [batch x 1 x 1 x features (96)] +# ww = [batch x window x features (96)] => [batch x probability] + + +class OpenWakeWord(TfLiteWakeWord): + def __init__( + self, + id: str, # pylint: disable=redefined-builtin + wake_word: str, + tflite_model: Union[str, Path], + libtensorflowlite_c_path: Union[str, Path], + ): + TfLiteWakeWord.__init__(self, libtensorflowlite_c_path) + + self.id = id + self.wake_word = wake_word + self.tflite_model = tflite_model + + self.is_active = True + + # Load the model and create interpreter + self.model_path = str(Path(tflite_model).resolve()).encode("utf-8") + self.model = self.lib.TfLiteModelCreateFromFile(self.model_path) + self.interpreter = self.lib.TfLiteInterpreterCreate(self.model, None) + self.lib.TfLiteInterpreterAllocateTensors(self.interpreter) + + self.input_tensor = self.lib.TfLiteInterpreterGetInputTensor( + self.interpreter, c_int32(0) + ) + self.output_tensor = self.lib.TfLiteInterpreterGetOutputTensor( + self.interpreter, c_int32(0) + ) + + num_input_dims = self.lib.TfLiteTensorNumDims(self.input_tensor) + input_shape = [ + self.lib.TfLiteTensorDim(self.input_tensor, i) + for i in range(num_input_dims) + ] + self.input_windows = input_shape[1] + + self.new_embeddings: int = 0 + self.embeddings: np.ndarray = np.zeros( + shape=(MAX_EMB, WW_FEATURES), dtype=np.float32 + ) + + def process_streaming(self, embeddings: np.ndarray) -> Iterable[float]: + """Generate probabilities from embeddings.""" + num_embedding_windows = embeddings.shape[2] + + # Shift + self.embeddings[:-num_embedding_windows] = self.embeddings[ + num_embedding_windows: + ] + + # Overwrite + self.embeddings[-num_embedding_windows:] = embeddings[0, 0, :, :] + self.new_embeddings = min( + len(self.embeddings), + self.new_embeddings + num_embedding_windows, + ) + + while self.new_embeddings >= self.input_windows: + emb_tensor = np.zeros( + shape=(1, self.input_windows, WW_FEATURES), + dtype=np.float32, + ) + + emb_tensor[0, :] = self.embeddings[ + -self.new_embeddings : len(self.embeddings) + - self.new_embeddings + + self.input_windows + ] + self.new_embeddings = max(0, self.new_embeddings - 1) + + # Run inference + emb_ptr = emb_tensor.ctypes.data_as(c_void_p) + self.lib.TfLiteTensorCopyFromBuffer( + self.input_tensor, emb_ptr, c_size_t(emb_tensor.nbytes) + ) + self.lib.TfLiteInterpreterInvoke(self.interpreter) + + output_bytes = self.lib.TfLiteTensorByteSize(self.output_tensor) + probs = np.empty( + output_bytes // np.dtype(np.float32).itemsize, dtype=np.float32 + ) + self.lib.TfLiteTensorCopyToBuffer( + self.output_tensor, + probs.ctypes.data_as(c_void_p), + c_size_t(output_bytes), + ) + + yield probs.item() + + @staticmethod + def from_config( + config_path: Union[str, Path], + libtensorflowlite_c_path: Union[str, Path], + ) -> "OpenWakeWord": + config_path = Path(config_path) + with open(config_path, "r", encoding="utf-8") as config_file: + config = json.load(config_file) + + return OpenWakeWord( + id=Path(config["model"]).stem, + wake_word=config["wake_word"], + tflite_model=config_path.parent / config["model"], + libtensorflowlite_c_path=libtensorflowlite_c_path, + ) + + +# ----------------------------------------------------------------------------- + + +class OpenWakeWordFeatures(TfLiteWakeWord): + def __init__( + self, + melspectrogram_model: Union[str, Path], + embedding_model: Union[str, Path], + libtensorflowlite_c_path: Union[str, Path], + ) -> None: + TfLiteWakeWord.__init__(self, libtensorflowlite_c_path) + + # Melspectrogram + mel_path = str(Path(melspectrogram_model).resolve()).encode("utf-8") + self.mel_model = self.lib.TfLiteModelCreateFromFile(mel_path) + self.mel_interpreter = self.lib.TfLiteInterpreterCreate(self.mel_model, None) + + mels_dims = (c_int32 * len(MEL_SHAPE))(*MEL_SHAPE) + self.lib.TfLiteInterpreterResizeInputTensor( + self.mel_interpreter, + c_int32(0), + mels_dims, + c_int32(len(MEL_SHAPE)), + ) + self.lib.TfLiteInterpreterAllocateTensors(self.mel_interpreter) + self.mel_input_tensor = self.lib.TfLiteInterpreterGetInputTensor( + self.mel_interpreter, c_int32(0) + ) + self.mel_output_tensor = self.lib.TfLiteInterpreterGetOutputTensor( + self.mel_interpreter, c_int32(0) + ) + + # Embedding + emb_path = str(Path(embedding_model).resolve()).encode("utf-8") + self.emb_model = self.lib.TfLiteModelCreateFromFile(emb_path) + self.emb_interpreter = self.lib.TfLiteInterpreterCreate(self.emb_model, None) + emb_dims = (c_int32 * len(EMB_SHAPE))(*EMB_SHAPE) + self.lib.TfLiteInterpreterResizeInputTensor( + self.emb_interpreter, + c_int32(0), + emb_dims, + c_int32(len(EMB_SHAPE)), + ) + self.lib.TfLiteInterpreterAllocateTensors(self.emb_interpreter) + self.emb_input_tensor = self.lib.TfLiteInterpreterGetInputTensor( + self.emb_interpreter, c_int32(0) + ) + self.emb_output_tensor = self.lib.TfLiteInterpreterGetOutputTensor( + self.emb_interpreter, c_int32(0) + ) + + # State + self.new_audio_samples: int = AUTOFILL_SECONDS * SAMPLE_RATE + self.audio: np.ndarray = np.zeros(shape=(_MAX_SAMPLES,), dtype=np.float32) + self.new_mels: int = 0 + self.mels: np.ndarray = np.zeros(shape=(MAX_MELS, NUM_MELS), dtype=np.float32) + + def process_streaming(self, audio_chunk: bytes) -> Iterable[np.ndarray]: + """Generate embeddings from audio.""" + chunk_array = np.frombuffer(audio_chunk, dtype=np.int16).astype(np.float32) + + # Shift samples left + self.audio[: -len(chunk_array)] = self.audio[len(chunk_array) :] + + # Add new samples to end + self.audio[-len(chunk_array) :] = chunk_array + self.new_audio_samples = min( + len(self.audio), + self.new_audio_samples + len(chunk_array), + ) + + while self.new_audio_samples >= MEL_SAMPLES: + audio_tensor = np.zeros(shape=(BATCH_SIZE, MEL_SAMPLES), dtype=np.float32) + audio_tensor[0, :] = self.audio[ + -self.new_audio_samples : len(self.audio) + - self.new_audio_samples + + MEL_SAMPLES + ] + audio_tensor = np.ascontiguousarray(audio_tensor) + self.new_audio_samples = max(0, self.new_audio_samples - SAMPLES_PER_CHUNK) + + # Run inference (mels) + audio_ptr = audio_tensor.ctypes.data_as(c_void_p) + self.lib.TfLiteTensorCopyFromBuffer( + self.mel_input_tensor, audio_ptr, c_size_t(audio_tensor.nbytes) + ) + self.lib.TfLiteInterpreterInvoke(self.mel_interpreter) + + mels_output_bytes = self.lib.TfLiteTensorByteSize(self.mel_output_tensor) + mels = np.empty( + mels_output_bytes // np.dtype(np.float32).itemsize, dtype=np.float32 + ) + self.lib.TfLiteTensorCopyToBuffer( + self.mel_output_tensor, + mels.ctypes.data_as(c_void_p), + c_size_t(mels_output_bytes), + ) + + mels = (mels / 10) + 2 # transform to fit embedding + mels = mels.reshape((1, 1, -1, NUM_MELS)) + + # Shift left + num_mel_windows = mels.shape[2] + self.mels[:-num_mel_windows] = self.mels[num_mel_windows:] + + # Overwrite + self.mels[-num_mel_windows:] = mels[0, 0, :, :] + self.new_mels = min(len(self.mels), self.new_mels + num_mel_windows) + + while self.new_mels >= EMB_FEATURES: + mels_tensor = np.ascontiguousarray( + np.zeros(shape=EMB_SHAPE, dtype=np.float32) + ) + mels_tensor[0, :, :, 0] = self.mels[ + -self.new_mels : len(self.mels) - self.new_mels + EMB_FEATURES, : + ] + self.new_mels = max(0, self.new_mels - EMB_STEP) + + # Run inference (embedding) + mels_ptr = mels_tensor.ctypes.data_as(c_void_p) + self.lib.TfLiteTensorCopyFromBuffer( + self.emb_input_tensor, mels_ptr, c_size_t(mels_tensor.nbytes) + ) + self.lib.TfLiteInterpreterInvoke(self.emb_interpreter) + + emb_output_bytes = self.lib.TfLiteTensorByteSize(self.emb_output_tensor) + emb = np.empty( + emb_output_bytes // np.dtype(np.float32).itemsize, dtype=np.float32 + ) + self.lib.TfLiteTensorCopyToBuffer( + self.emb_output_tensor, + emb.ctypes.data_as(c_void_p), + c_size_t(emb_output_bytes), + ) + emb = emb.reshape((1, 1, -1, WW_FEATURES)) + yield emb diff --git a/linux_voice_assistant/satellite.py b/linux_voice_assistant/satellite.py index 1537d0d..e43b168 100644 --- a/linux_voice_assistant/satellite.py +++ b/linux_voice_assistant/satellite.py @@ -3,7 +3,7 @@ import logging import time from collections.abc import Iterable -from typing import Dict, Optional, Set +from typing import Dict, Optional, Set, Union # pylint: disable=no-name-in-module from aioesphomeapi.api_pb2 import ( # type: ignore[attr-defined] @@ -35,6 +35,7 @@ from .api_server import APIServer from .entity import MediaPlayerEntity from .microwakeword import MicroWakeWord from .models import ServerState +from .openwakeword import OpenWakeWord from .util import call_all _LOGGER = logging.getLogger(__name__) @@ -194,9 +195,8 @@ class VoiceSatelliteProtocol(APIServer): continue _LOGGER.debug("Loading wake word: %s", model_info.config_path) - self.state.wake_words[wake_word_id] = MicroWakeWord.from_config( - model_info.config_path, - self.state.libtensorflowlite_c_path, + self.state.wake_words[wake_word_id] = model_info.load( + self.state.libtensorflowlite_c_path ) _LOGGER.info("Wake word set: %s", wake_word_id) @@ -219,7 +219,7 @@ class VoiceSatelliteProtocol(APIServer): self.send_messages([VoiceAssistantAudio(data=audio_chunk)]) - def wakeup(self, wake_word: MicroWakeWord) -> None: + def wakeup(self, wake_word: Union[MicroWakeWord, OpenWakeWord]) -> None: if self._timer_finished: # Stop timer instead self._timer_finished = False diff --git a/linux_voice_assistant/wakeword.py b/linux_voice_assistant/wakeword.py new file mode 100644 index 0000000..ae9f40e --- /dev/null +++ b/linux_voice_assistant/wakeword.py @@ -0,0 +1,89 @@ +"""Base class for wake words.""" + +import ctypes as C +import os +from pathlib import Path +from typing import Union + +TfLiteStatus = C.c_int # kTfLiteOk == 0 + + +class TfLiteQuantizationParams(C.Structure): + _fields_ = [("scale", C.c_float), ("zero_point", C.c_int32)] + + +c_void_p = C.c_void_p +c_int32 = C.c_int32 +c_size_t = C.c_size_t + + +def _not_null(result, func, args): + if not result: + raise RuntimeError(f"{func.__name__} returned NULL") + return result + + +class TfLiteWakeWord: + def __init__(self, libtensorflowlite_c_path: Union[str, Path]): + self.libtensorflowlite_c_path = Path(libtensorflowlite_c_path) + # Load ONCE with RTLD_GLOBAL + self.lib = C.CDLL( + str(self.libtensorflowlite_c_path.resolve()), mode=os.RTLD_GLOBAL + ) + + lib = self.lib + + # Model / Interpreter + lib.TfLiteModelCreateFromFile.argtypes = [C.c_char_p] + lib.TfLiteModelCreateFromFile.restype = c_void_p + + lib.TfLiteInterpreterCreate.argtypes = [c_void_p, c_void_p] + lib.TfLiteInterpreterCreate.restype = c_void_p + + lib.TfLiteInterpreterAllocateTensors.argtypes = [c_void_p] + lib.TfLiteInterpreterAllocateTensors.restype = TfLiteStatus + + lib.TfLiteInterpreterInvoke.argtypes = [c_void_p] + lib.TfLiteInterpreterInvoke.restype = TfLiteStatus + + # Tensors: get + lib.TfLiteInterpreterGetInputTensor.argtypes = [c_void_p, c_int32] + lib.TfLiteInterpreterGetInputTensor.restype = c_void_p + lib.TfLiteInterpreterGetInputTensor.errcheck = _not_null + + lib.TfLiteInterpreterGetOutputTensor.argtypes = [c_void_p, c_int32] + lib.TfLiteInterpreterGetOutputTensor.restype = c_void_p + lib.TfLiteInterpreterGetOutputTensor.errcheck = _not_null + + # Sizes / dims + lib.TfLiteTensorByteSize.argtypes = [c_void_p] + lib.TfLiteTensorByteSize.restype = c_size_t + + lib.TfLiteTensorNumDims.argtypes = [c_void_p] + lib.TfLiteTensorNumDims.restype = c_int32 + + lib.TfLiteTensorDim.argtypes = [c_void_p, c_int32] + lib.TfLiteTensorDim.restype = c_int32 + + # Type & quant params + lib.TfLiteTensorType.argtypes = [c_void_p] + lib.TfLiteTensorType.restype = c_int32 + + lib.TfLiteTensorQuantizationParams.argtypes = [c_void_p] + lib.TfLiteTensorQuantizationParams.restype = TfLiteQuantizationParams + + # Resize + lib.TfLiteInterpreterResizeInputTensor.argtypes = [ + c_void_p, + c_int32, + C.POINTER(c_int32), + c_int32, + ] + lib.TfLiteInterpreterResizeInputTensor.restype = TfLiteStatus + + # Copy buffers + lib.TfLiteTensorCopyFromBuffer.argtypes = [c_void_p, c_void_p, c_size_t] + lib.TfLiteTensorCopyFromBuffer.restype = TfLiteStatus + + lib.TfLiteTensorCopyToBuffer.argtypes = [c_void_p, c_void_p, c_size_t] + lib.TfLiteTensorCopyToBuffer.restype = TfLiteStatus diff --git a/script/test b/script/test new file mode 100755 index 0000000..98a3068 --- /dev/null +++ b/script/test @@ -0,0 +1,18 @@ +#!/usr/bin/env python3 +import subprocess +import sys +import venv +from pathlib import Path + +_DIR = Path(__file__).parent +_PROGRAM_DIR = _DIR.parent +_VENV_DIR = _PROGRAM_DIR / ".venv" +_TEST_DIR = _PROGRAM_DIR / "tests" + +if _VENV_DIR.exists(): + context = venv.EnvBuilder().ensure_directories(_VENV_DIR) + python_exe = context.env_exe +else: + python_exe = "python3" + +subprocess.check_call([python_exe, "-m", "pytest", _TEST_DIR] + sys.argv[1:]) diff --git a/tests/ok_nabu.wav b/tests/ok_nabu.wav new file mode 100644 index 0000000..185fa75 Binary files /dev/null and b/tests/ok_nabu.wav differ diff --git a/tests/test_microwakeword.py b/tests/test_microwakeword.py new file mode 100644 index 0000000..29b240b --- /dev/null +++ b/tests/test_microwakeword.py @@ -0,0 +1,43 @@ +"""Tests for microWakeWord.""" + +import wave +from pathlib import Path + +from linux_voice_assistant.microwakeword import MicroWakeWord, MicroWakeWordFeatures +from linux_voice_assistant.util import is_arm + +_TESTS_DIR = Path(__file__).parent +_REPO_DIR = _TESTS_DIR.parent +_MICRO_DIR = _REPO_DIR / "wakewords" + +if is_arm(): + _LIB_DIR = _REPO_DIR / "lib" / "linux_arm64" +else: + _LIB_DIR = _REPO_DIR / "lib" / "linux_amd64" + + +libtensorflowlite_c_path = _LIB_DIR / "libtensorflowlite_c.so" + + +def test_features() -> None: + features = MicroWakeWordFeatures( + libtensorflowlite_c_path=libtensorflowlite_c_path, + ) + ww = MicroWakeWord.from_config( + config_path=_MICRO_DIR / "okay_nabu.json", + libtensorflowlite_c_path=libtensorflowlite_c_path, + ) + + detected = False + with wave.open(str(_TESTS_DIR / "ok_nabu.wav"), "rb") as wav_file: + assert wav_file.getframerate() == 16000 + assert wav_file.getsampwidth() == 2 + assert wav_file.getnchannels() == 1 + + for micro_input in features.process_streaming( + wav_file.readframes(wav_file.getnframes()) + ): + if ww.process_streaming(micro_input): + detected = True + + assert detected diff --git a/tests/test_openwakeword.py b/tests/test_openwakeword.py new file mode 100644 index 0000000..b2594ec --- /dev/null +++ b/tests/test_openwakeword.py @@ -0,0 +1,47 @@ +"""Tests for openWakeWord.""" + +import wave +from pathlib import Path + +from linux_voice_assistant.openwakeword import OpenWakeWordFeatures, OpenWakeWord +from linux_voice_assistant.util import is_arm + +_TESTS_DIR = Path(__file__).parent +_REPO_DIR = _TESTS_DIR.parent +_OWW_DIR = _REPO_DIR / "wakewords" / "openWakeWord" + +if is_arm(): + _LIB_DIR = _REPO_DIR / "lib" / "linux_arm64" +else: + _LIB_DIR = _REPO_DIR / "lib" / "linux_amd64" + + +libtensorflowlite_c_path = _LIB_DIR / "libtensorflowlite_c.so" + + +def test_features() -> None: + features = OpenWakeWordFeatures( + melspectrogram_model=_OWW_DIR / "melspectrogram.tflite", + embedding_model=_OWW_DIR / "embedding_model.tflite", + libtensorflowlite_c_path=libtensorflowlite_c_path, + ) + ww = OpenWakeWord( + id="ok_nabu", + wake_word="okay nabu", + tflite_model=_OWW_DIR / "ok_nabu_v0.1.tflite", + libtensorflowlite_c_path=libtensorflowlite_c_path, + ) + + max_prob = 0.0 + with wave.open(str(_TESTS_DIR / "ok_nabu.wav"), "rb") as wav_file: + assert wav_file.getframerate() == 16000 + assert wav_file.getsampwidth() == 2 + assert wav_file.getnchannels() == 1 + + for embeddings in features.process_streaming( + wav_file.readframes(wav_file.getnframes()) + ): + for prob in ww.process_streaming(embeddings): + max_prob = max(max_prob, prob) + + assert max_prob > 0.5 diff --git a/wakewords/openWakeWord/alexa_v0.1.json b/wakewords/openWakeWord/alexa_v0.1.json new file mode 100644 index 0000000..c7a4d2e --- /dev/null +++ b/wakewords/openWakeWord/alexa_v0.1.json @@ -0,0 +1,5 @@ +{ + "type": "openWakeWord", + "wake_word": "Alexa", + "model": "alexa_v0.1.tflite" +} diff --git a/wakewords/openWakeWord/alexa_v0.1.tflite b/wakewords/openWakeWord/alexa_v0.1.tflite new file mode 100644 index 0000000..de9ff43 Binary files /dev/null and b/wakewords/openWakeWord/alexa_v0.1.tflite differ diff --git a/wakewords/openWakeWord/embedding_model.tflite b/wakewords/openWakeWord/embedding_model.tflite new file mode 100644 index 0000000..d3c3bed Binary files /dev/null and b/wakewords/openWakeWord/embedding_model.tflite differ diff --git a/wakewords/openWakeWord/hey_jarvis_v0.1.json b/wakewords/openWakeWord/hey_jarvis_v0.1.json new file mode 100644 index 0000000..b7d632e --- /dev/null +++ b/wakewords/openWakeWord/hey_jarvis_v0.1.json @@ -0,0 +1,5 @@ +{ + "type": "openWakeWord", + "wake_word": "Hey Jarvis", + "model": "hey_jarvis_v0.1.tflite" +} diff --git a/wakewords/openWakeWord/hey_jarvis_v0.1.tflite b/wakewords/openWakeWord/hey_jarvis_v0.1.tflite new file mode 100644 index 0000000..aafb043 Binary files /dev/null and b/wakewords/openWakeWord/hey_jarvis_v0.1.tflite differ diff --git a/wakewords/openWakeWord/hey_mycroft_v0.1.json b/wakewords/openWakeWord/hey_mycroft_v0.1.json new file mode 100644 index 0000000..abd5118 --- /dev/null +++ b/wakewords/openWakeWord/hey_mycroft_v0.1.json @@ -0,0 +1,5 @@ +{ + "type": "openWakeWord", + "wake_word": "Hey Mycroft", + "model": "hey_mycroft_v0.1.tflite" +} diff --git a/wakewords/openWakeWord/hey_mycroft_v0.1.tflite b/wakewords/openWakeWord/hey_mycroft_v0.1.tflite new file mode 100644 index 0000000..d8a644a Binary files /dev/null and b/wakewords/openWakeWord/hey_mycroft_v0.1.tflite differ diff --git a/wakewords/openWakeWord/hey_rhasspy_v0.1.json b/wakewords/openWakeWord/hey_rhasspy_v0.1.json new file mode 100644 index 0000000..1ecc454 --- /dev/null +++ b/wakewords/openWakeWord/hey_rhasspy_v0.1.json @@ -0,0 +1,5 @@ +{ + "type": "openWakeWord", + "wake_word": "Hey Rhasspy", + "model": "hey_rhasspy_v0.1.tflite" +} diff --git a/wakewords/openWakeWord/hey_rhasspy_v0.1.tflite b/wakewords/openWakeWord/hey_rhasspy_v0.1.tflite new file mode 100644 index 0000000..3f8d3ef Binary files /dev/null and b/wakewords/openWakeWord/hey_rhasspy_v0.1.tflite differ diff --git a/wakewords/openWakeWord/melspectrogram.tflite b/wakewords/openWakeWord/melspectrogram.tflite new file mode 100644 index 0000000..e290399 Binary files /dev/null and b/wakewords/openWakeWord/melspectrogram.tflite differ diff --git a/wakewords/openWakeWord/ok_nabu_v0.1.json b/wakewords/openWakeWord/ok_nabu_v0.1.json new file mode 100644 index 0000000..3fae45e --- /dev/null +++ b/wakewords/openWakeWord/ok_nabu_v0.1.json @@ -0,0 +1,5 @@ +{ + "type": "openWakeWord", + "wake_word": "Okay Nabu", + "model": "ok_nabu_v0.1.tflite" +} diff --git a/wakewords/openWakeWord/ok_nabu_v0.1.tflite b/wakewords/openWakeWord/ok_nabu_v0.1.tflite new file mode 100644 index 0000000..ea36458 Binary files /dev/null and b/wakewords/openWakeWord/ok_nabu_v0.1.tflite differ