274 lines
9.3 KiB
Python
274 lines
9.3 KiB
Python
import ctypes
|
|
import json
|
|
import statistics
|
|
from collections import deque
|
|
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
|
|
|
|
|
|
class TfLiteQuantizationParams(ctypes.Structure):
|
|
_fields_ = [("scale", ctypes.c_float), ("zero_point", ctypes.c_int32)]
|
|
|
|
|
|
class MicroWakeWord:
|
|
def __init__(
|
|
self,
|
|
id: str, # pylint: disable=redefined-builtin
|
|
wake_word: str,
|
|
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],
|
|
):
|
|
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.lib = ctypes.cdll.LoadLibrary(
|
|
str(Path(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)
|
|
self.interpreter = self.lib.TfLiteInterpreterCreate(self.model, None)
|
|
self.lib.TfLiteInterpreterAllocateTensors(self.interpreter)
|
|
|
|
# Access input and output tensor
|
|
self.input_tensor = self.lib.TfLiteInterpreterGetInputTensor(
|
|
self.interpreter, 0
|
|
)
|
|
self.output_tensor = self.lib.TfLiteInterpreterGetOutputTensor(
|
|
self.interpreter, 0
|
|
)
|
|
|
|
# Get quantization parameters
|
|
input_q = self.lib.TfLiteTensorQuantizationParams(self.input_tensor)
|
|
output_q = self.lib.TfLiteTensorQuantizationParams(self.output_tensor)
|
|
|
|
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.
|
|
|
|
Parameters
|
|
----------
|
|
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:
|
|
config = json.load(config_file)
|
|
|
|
micro_config = config["micro"]
|
|
|
|
return MicroWakeWord(
|
|
id=Path(config["model"]).stem,
|
|
wake_word=config["wake_word"],
|
|
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.
|
|
|
|
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
|
|
return False
|
|
|
|
detected = False
|
|
audio_buffer_idx = 0
|
|
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
|
|
]
|
|
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)
|
|
|
|
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))
|
|
)
|
|
)
|
|
|
|
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
|