"""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