90 lines
2.9 KiB
Python
90 lines
2.9 KiB
Python
"""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
|