micro_wakeword/
detector.rs

1use std::collections::VecDeque;
2use std::path::{Path, PathBuf};
3
4use tflite_c::{Interpreter, InterpreterOptions, Model, TfLiteType};
5
6use crate::config::{Config, ModelMetadata};
7use crate::features::{FEATURE_COUNT, FEATURE_SCALE, Frontend};
8use crate::{AUDIO_BLOCK_SAMPLES, Error, Result, Runtime};
9
10/// A thresholded wake-word prediction.
11#[derive(Clone, Debug, PartialEq)]
12pub struct Detection {
13    pub wake_word: String,
14    /// Averaged probability in the inclusive range 0–1.
15    pub probability: f32,
16}
17
18/// Streaming detector for 16 kHz mono signed 16-bit PCM.
19pub struct Detector {
20    config: Config,
21    runtime: Runtime,
22    frontend: Frontend,
23    engine: Engine,
24    feature_rows: Vec<[u16; FEATURE_COUNT]>,
25    probabilities: VecDeque<f32>,
26}
27
28struct Engine {
29    interpreter: Interpreter,
30    input_scale: f32,
31    input_zero_point: i32,
32    model_rows: usize,
33}
34
35impl Engine {
36    fn new(config: &Config, runtime: &Runtime) -> Result<Self> {
37        let library = runtime.load()?;
38        let model = Model::from_file(&config.model_path, library.clone())?;
39        let mut options = InterpreterOptions::new(library);
40        options.num_threads(1);
41        let interpreter = Interpreter::new(model, options)?;
42
43        if interpreter.input_count() != 1 {
44            return Err(Error::IncompatibleModel(format!(
45                "expected one input tensor, got {}",
46                interpreter.input_count()
47            )));
48        }
49        if interpreter.output_count() != 1 {
50            return Err(Error::IncompatibleModel(format!(
51                "expected one output tensor, got {}",
52                interpreter.output_count()
53            )));
54        }
55        let input = interpreter.input(0)?;
56        let dims = input.dims();
57        if dims.len() != 3 || dims[0] != 1 || dims[1] <= 0 || dims[2] != FEATURE_COUNT as i32 {
58            return Err(Error::IncompatibleModel(format!(
59                "expected int8 input [1, rows, {FEATURE_COUNT}], got {:?}",
60                dims
61            )));
62        }
63        if input.dtype() != TfLiteType::Int8 {
64            return Err(Error::IncompatibleModel(format!(
65                "expected int8 input, got {:?}",
66                input.dtype()
67            )));
68        }
69        let quantization = input.quantization();
70        if !quantization.scale.is_finite() || quantization.scale <= 0.0 {
71            return Err(Error::IncompatibleModel(
72                "input tensor has invalid quantization scale".into(),
73            ));
74        }
75        let model_rows = dims[1] as usize;
76        let output = interpreter.output(0)?;
77        let output_elements = output
78            .dims()
79            .iter()
80            .try_fold(1_i32, |n, dim| n.checked_mul(*dim));
81        if output_elements != Some(1) {
82            return Err(Error::IncompatibleModel(format!(
83                "expected a single probability output, got shape {:?}",
84                output.dims()
85            )));
86        }
87        if !matches!(
88            output.dtype(),
89            TfLiteType::Int8 | TfLiteType::UInt8 | TfLiteType::Float32
90        ) {
91            return Err(Error::IncompatibleModel(format!(
92                "expected a quantized or float probability output, got {:?}",
93                output.dtype()
94            )));
95        }
96        if output.dtype() != TfLiteType::Float32 {
97            let quantization = output.quantization();
98            if !quantization.scale.is_finite() || quantization.scale <= 0.0 {
99                return Err(Error::IncompatibleModel(
100                    "output tensor has invalid quantization scale".into(),
101                ));
102            }
103        }
104
105        Ok(Self {
106            interpreter,
107            input_scale: quantization.scale,
108            input_zero_point: quantization.zero_point,
109            model_rows,
110        })
111    }
112}
113
114impl Detector {
115    pub fn from_config(path: impl AsRef<Path>) -> Result<Self> {
116        Self::from_config_with_runtime(path, Runtime::Auto)
117    }
118
119    pub fn from_config_with_runtime(path: impl AsRef<Path>, runtime: Runtime) -> Result<Self> {
120        Self::from_parts(Config::from_file(path)?, runtime)
121    }
122
123    pub fn builder(model_path: impl Into<PathBuf>) -> DetectorBuilder {
124        DetectorBuilder::new(model_path)
125    }
126
127    pub(crate) fn from_parts(config: Config, runtime: Runtime) -> Result<Self> {
128        config.validate()?;
129        let frontend = Frontend::new()?;
130        let engine = Engine::new(&config, &runtime)?;
131        let model_rows = engine.model_rows;
132        let window = config.sliding_window_size;
133        Ok(Self {
134            config,
135            runtime,
136            frontend,
137            engine,
138            feature_rows: Vec::with_capacity(model_rows),
139            probabilities: VecDeque::with_capacity(window),
140        })
141    }
142
143    pub fn config(&self) -> &Config {
144        &self.config
145    }
146
147    /// Consume exactly one 10 ms audio block (160 samples).
148    pub fn process_audio(&mut self, samples: &[i16]) -> Result<Option<Detection>> {
149        let samples: &[i16; AUDIO_BLOCK_SAMPLES] = samples.try_into().map_err(|_| {
150            Error::Audio(format!(
151                "expected exactly {AUDIO_BLOCK_SAMPLES} samples of 16 kHz mono PCM, got {}",
152                samples.len()
153            ))
154        })?;
155        let Some(features) = self.frontend.process(samples)? else {
156            return Ok(None);
157        };
158        self.process_features(features)
159    }
160
161    fn process_features(&mut self, features: [u16; FEATURE_COUNT]) -> Result<Option<Detection>> {
162        self.feature_rows.push(features);
163        if self.feature_rows.len() < self.engine.model_rows {
164            return Ok(None);
165        }
166
167        {
168            let mut input = self.engine.interpreter.input_mut(0)?;
169            let destination = input.data_mut()?;
170            for (target, feature) in destination
171                .iter_mut()
172                .zip(self.feature_rows.iter().flat_map(|row| row.iter().copied()))
173            {
174                let real = feature as f32 * FEATURE_SCALE;
175                let quantized = (real / self.engine.input_scale
176                    + self.engine.input_zero_point as f32)
177                    .round()
178                    .clamp(i8::MIN as f32, i8::MAX as f32) as i8;
179                *target = quantized as u8;
180            }
181        }
182        self.feature_rows.clear();
183        self.engine.interpreter.invoke()?;
184        let output = self.engine.interpreter.output(0)?.to_vec_f32()?;
185        let probability = output[0].clamp(0.0, 1.0);
186        self.probabilities.push_back(probability);
187        if self.probabilities.len() > self.config.sliding_window_size {
188            self.probabilities.pop_front();
189        }
190        if self.probabilities.len() < self.config.sliding_window_size {
191            return Ok(None);
192        }
193        let probability = self.probabilities.iter().sum::<f32>() / self.probabilities.len() as f32;
194        Ok(
195            passes_cutoff(probability, self.config.probability_cutoff).then(|| Detection {
196                wake_word: self.config.wake_word.clone(),
197                probability,
198            }),
199        )
200    }
201
202    /// Clear frontend history, probability history, and all model recurrent state.
203    pub fn reset(&mut self) -> Result<()> {
204        let frontend = Frontend::new()?;
205        let engine = Engine::new(&self.config, &self.runtime)?;
206        self.frontend = frontend;
207        self.engine = engine;
208        self.feature_rows.clear();
209        self.probabilities.clear();
210        Ok(())
211    }
212}
213
214fn passes_cutoff(probability: f32, cutoff: f32) -> bool {
215    probability > cutoff
216}
217
218/// Builder for detectors that do not use a model JSON file.
219pub struct DetectorBuilder {
220    model_path: PathBuf,
221    wake_word: Option<String>,
222    probability_cutoff: Option<f32>,
223    sliding_window_size: Option<usize>,
224    feature_step_size_ms: u32,
225    runtime: Runtime,
226}
227
228impl DetectorBuilder {
229    fn new(model_path: impl Into<PathBuf>) -> Self {
230        Self {
231            model_path: model_path.into(),
232            wake_word: None,
233            probability_cutoff: None,
234            sliding_window_size: None,
235            feature_step_size_ms: 10,
236            runtime: Runtime::Auto,
237        }
238    }
239
240    pub fn wake_word(mut self, value: impl Into<String>) -> Self {
241        self.wake_word = Some(value.into());
242        self
243    }
244    pub fn probability_cutoff(mut self, value: f32) -> Self {
245        self.probability_cutoff = Some(value);
246        self
247    }
248    pub fn sliding_window_size(mut self, value: usize) -> Self {
249        self.sliding_window_size = Some(value);
250        self
251    }
252    pub fn feature_step_size_ms(mut self, value: u32) -> Self {
253        self.feature_step_size_ms = value;
254        self
255    }
256    pub fn runtime(mut self, value: Runtime) -> Self {
257        self.runtime = value;
258        self
259    }
260
261    pub fn build(self) -> Result<Detector> {
262        let missing = |name| Error::InvalidConfig(format!("builder requires `{name}`"));
263        let config = Config {
264            model_path: self.model_path,
265            wake_word: self.wake_word.ok_or_else(|| missing("wake_word"))?,
266            probability_cutoff: self
267                .probability_cutoff
268                .ok_or_else(|| missing("probability_cutoff"))?,
269            sliding_window_size: self
270                .sliding_window_size
271                .ok_or_else(|| missing("sliding_window_size"))?,
272            feature_step_size_ms: self.feature_step_size_ms,
273            metadata: ModelMetadata {
274                format_version: 2,
275                ..ModelMetadata::default()
276            },
277        };
278        Detector::from_parts(config, self.runtime)
279    }
280}
281
282#[cfg(test)]
283mod tests {
284    use super::passes_cutoff;
285
286    #[test]
287    fn cutoff_boundary_is_strict_like_micro_wakeword() {
288        assert!(!passes_cutoff(0.3, 0.3));
289        assert!(passes_cutoff(0.300_001, 0.3));
290    }
291}
micro_wakeword/detector.rs

micro_wakeword/
detector.rs
