aria_kit/voice_activity/
endpointer.rs

1use crate::audio::{SilenceDuration, canonical_samples};
2
3use super::{SPEECH_THRESHOLD, VoiceActivityDetector, VoiceActivityError};
4
5/// At or below this a frame starts the trailing-silence clock. Between it and
6/// [`SPEECH_THRESHOLD`] the clock neither starts nor resets, so a soft
7/// syllable does not end an utterance.
8const SPEECH_END_THRESHOLD: f32 = 0.35;
9
10/// A speech boundary, in canonical samples counted from the first
11/// [`Endpointer::push`].
12#[derive(Debug, Clone, Copy, PartialEq, Eq)]
13pub enum Boundary {
14    /// The first sample of the frame where speech opened.
15    Start(u64),
16    /// The end of the frame where the trailing silence completed, so the
17    /// utterance before it includes that silence.
18    End(u64),
19}
20
21/// Finds utterances in a stream. Speech opens on a frame scored at or above
22/// [`SPEECH_THRESHOLD`] and closes once silence has lasted `trailing_silence`.
23/// After an `End` it keeps listening, so the next `Start` opens the next
24/// utterance. An utterance still open when input ends gets no `End`.
25///
26/// Positions count from the first push. To land them on a
27/// [`RecordingReader`](crate::audio::RecordingReader) that did not start at
28/// zero, add its `held_from()` at that push. `Start` has no lead-in: open an
29/// utterance [`LEAD_IN_FRAMES`](super::LEAD_IN_FRAMES) frames earlier to cut
30/// as the microphone does.
31///
32/// ```
33/// # use aria_kit::audio::SilenceDuration;
34/// # use aria_kit::voice_activity::{Boundary, Endpointer, VoiceActivityDetector, VoiceActivityError};
35/// fn utterances<V: VoiceActivityDetector>(
36///     vad: &mut V,
37///     batches: &[Vec<i16>],
38/// ) -> Result<(), VoiceActivityError> {
39///     // `&mut V` leaves the loaded detector with the caller for the next run.
40///     let mut endpointer = Endpointer::new(vad, SilenceDuration::default());
41///     for batch in batches {
42///         for boundary in endpointer.push(batch)? {
43///             match boundary {
44///                 Boundary::Start(at) => println!("speech from sample {at}"),
45///                 Boundary::End(at) => println!("utterance ends at sample {at}"),
46///             }
47///         }
48///     }
49///     Ok(())
50/// }
51/// ```
52pub struct Endpointer<V> {
53    vad: V,
54    trailing_silence_samples: u64,
55    frame: Vec<i16>,
56    scored_samples: u64,
57    speech_samples: u64,
58    speaking: bool,
59    silence_since: Option<u64>,
60}
61
62impl<V: VoiceActivityDetector> Endpointer<V> {
63    /// Resets `vad`, so state from earlier audio does not leak in.
64    pub fn new(mut vad: V, trailing_silence: SilenceDuration) -> Self {
65        // A zero-sample frame never fills, so `push` would never return.
66        const { assert!(V::FRAME_SAMPLES > 0) };
67        vad.reset();
68        Self {
69            vad,
70            trailing_silence_samples: canonical_samples(trailing_silence.get()),
71            frame: Vec::with_capacity(V::FRAME_SAMPLES),
72            scored_samples: 0,
73            speech_samples: 0,
74            speaking: false,
75            silence_since: None,
76        }
77    }
78
79    /// Where the open utterance's trailing silence began: `None` while
80    /// speech goes on, or outside an utterance.
81    #[cfg(all(feature = "silero", not(target_family = "wasm")))]
82    pub(crate) fn silence_since(&self) -> Option<u64> {
83        self.silence_since
84    }
85
86    #[cfg(all(feature = "silero", not(target_family = "wasm")))]
87    pub(crate) fn speech_samples(&self) -> u64 {
88        self.speech_samples
89    }
90
91    /// A partial frame waits for the next push.
92    pub fn push(&mut self, mut samples: &[i16]) -> Result<Vec<Boundary>, VoiceActivityError> {
93        let mut boundaries = Vec::new();
94        while !samples.is_empty() {
95            let copied = (V::FRAME_SAMPLES - self.frame.len()).min(samples.len());
96            self.frame.extend_from_slice(&samples[..copied]);
97            samples = &samples[copied..];
98            if self.frame.len() == V::FRAME_SAMPLES {
99                boundaries.extend(self.score_frame()?);
100                self.frame.clear();
101            }
102        }
103        Ok(boundaries)
104    }
105
106    fn score_frame(&mut self) -> Result<Option<Boundary>, VoiceActivityError> {
107        let probability = self.vad.probability(&self.frame)?;
108        // Counted only once scored, so a push retried after an error keeps
109        // its positions.
110        let frame_start = self.scored_samples;
111        self.scored_samples += V::FRAME_SAMPLES as u64;
112        if probability >= SPEECH_THRESHOLD {
113            self.speech_samples += V::FRAME_SAMPLES as u64;
114        }
115
116        if !self.speaking {
117            self.speaking = probability >= SPEECH_THRESHOLD;
118            return Ok(self.speaking.then_some(Boundary::Start(frame_start)));
119        }
120        if probability >= SPEECH_THRESHOLD {
121            self.silence_since = None;
122        } else if probability <= SPEECH_END_THRESHOLD && self.silence_since.is_none() {
123            self.silence_since = Some(frame_start);
124        }
125        let silence = self
126            .silence_since
127            .map_or(0, |since| self.scored_samples - since);
128        if silence < self.trailing_silence_samples {
129            return Ok(None);
130        }
131        self.speaking = false;
132        self.silence_since = None;
133        Ok(Some(Boundary::End(self.scored_samples)))
134    }
135}
136
137#[cfg(test)]
138mod tests {
139    use std::time::Duration;
140
141    use super::*;
142
143    struct Scripted {
144        scores: Vec<f32>,
145        calls: usize,
146        resets: usize,
147    }
148
149    impl VoiceActivityDetector for Scripted {
150        const FRAME_SAMPLES: usize = 160;
151
152        fn reset(&mut self) {
153            self.resets += 1;
154        }
155
156        fn probability(&mut self, _samples: &[i16]) -> Result<f32, VoiceActivityError> {
157            let score = self.scores[self.calls];
158            self.calls += 1;
159            Ok(score)
160        }
161    }
162
163    fn scripted(scores: &[f32]) -> Scripted {
164        Scripted {
165            scores: scores.to_vec(),
166            calls: 0,
167            resets: 0,
168        }
169    }
170
171    const FRAME: u64 = Scripted::FRAME_SAMPLES as u64;
172
173    fn ten_frames_of_silence() -> SilenceDuration {
174        SilenceDuration::new(Duration::from_millis(100)).expect("the duration is in range")
175    }
176
177    fn boundaries(scores: &[f32], batch: usize) -> Vec<Boundary> {
178        let mut endpointer = Endpointer::new(scripted(scores), ten_frames_of_silence());
179        let samples = vec![0; scores.len() * Scripted::FRAME_SAMPLES];
180        let mut found = Vec::new();
181        for piece in samples.chunks(batch) {
182            found.extend(endpointer.push(piece).expect("the scripted VAD scores"));
183        }
184        found
185    }
186
187    #[test]
188    fn consecutive_utterances_each_get_a_start_and_an_end() {
189        let mut scores = vec![0.0, 0.9, 0.9];
190        scores.extend([0.0; 10]);
191        scores.extend([0.9, 0.0]);
192        let ten_frames = 10 * FRAME;
193
194        for batch in [1, 100, 160, 1_024] {
195            assert_eq!(
196                boundaries(&scores, batch),
197                [
198                    Boundary::Start(FRAME),
199                    Boundary::End(3 * FRAME + ten_frames),
200                    Boundary::Start(13 * FRAME),
201                ],
202                "batches of {batch} samples"
203            );
204        }
205    }
206
207    #[test]
208    fn a_score_between_the_thresholds_does_not_end_the_utterance() {
209        let mut scores = vec![0.9];
210        scores.extend([0.4; 20]);
211
212        assert_eq!(boundaries(&scores, 160), [Boundary::Start(0)]);
213    }
214
215    #[test]
216    fn a_new_endpointer_resets_its_detector() {
217        let mut vad = scripted(&[]);
218
219        Endpointer::new(&mut vad, ten_frames_of_silence());
220
221        assert_eq!(vad.resets, 1);
222    }
223}
aria_kit/voice_activity/endpointer.rs

aria_kit/voice_activity/
endpointer.rs
