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    speaking: bool,
58    silence_since: Option<u64>,
59}
60
61impl<V: VoiceActivityDetector> Endpointer<V> {
62    /// Resets `vad`, so state from earlier audio does not leak in.
63    pub fn new(mut vad: V, trailing_silence: SilenceDuration) -> Self {
64        // A zero-sample frame never fills, so `push` would never return.
65        const { assert!(V::FRAME_SAMPLES > 0) };
66        vad.reset();
67        Self {
68            vad,
69            trailing_silence_samples: canonical_samples(trailing_silence.get()),
70            frame: Vec::with_capacity(V::FRAME_SAMPLES),
71            scored_samples: 0,
72            speaking: false,
73            silence_since: None,
74        }
75    }
76
77    /// A partial frame waits for the next push.
78    pub fn push(&mut self, mut samples: &[i16]) -> Result<Vec<Boundary>, VoiceActivityError> {
79        let mut boundaries = Vec::new();
80        while !samples.is_empty() {
81            let copied = (V::FRAME_SAMPLES - self.frame.len()).min(samples.len());
82            self.frame.extend_from_slice(&samples[..copied]);
83            samples = &samples[copied..];
84            if self.frame.len() == V::FRAME_SAMPLES {
85                boundaries.extend(self.score_frame()?);
86                self.frame.clear();
87            }
88        }
89        Ok(boundaries)
90    }
91
92    fn score_frame(&mut self) -> Result<Option<Boundary>, VoiceActivityError> {
93        let probability = self.vad.probability(&self.frame)?;
94        // Counted only once scored, so a push retried after an error keeps
95        // its positions.
96        let frame_start = self.scored_samples;
97        self.scored_samples += V::FRAME_SAMPLES as u64;
98
99        if !self.speaking {
100            self.speaking = probability >= SPEECH_THRESHOLD;
101            return Ok(self.speaking.then_some(Boundary::Start(frame_start)));
102        }
103        if probability >= SPEECH_THRESHOLD {
104            self.silence_since = None;
105        } else if probability <= SPEECH_END_THRESHOLD && self.silence_since.is_none() {
106            self.silence_since = Some(frame_start);
107        }
108        let silence = self
109            .silence_since
110            .map_or(0, |since| self.scored_samples - since);
111        if silence < self.trailing_silence_samples {
112            return Ok(None);
113        }
114        self.speaking = false;
115        self.silence_since = None;
116        Ok(Some(Boundary::End(self.scored_samples)))
117    }
118}
119
120#[cfg(test)]
121mod tests {
122    use std::time::Duration;
123
124    use super::*;
125
126    struct Scripted {
127        scores: Vec<f32>,
128        calls: usize,
129        resets: usize,
130    }
131
132    impl VoiceActivityDetector for Scripted {
133        const FRAME_SAMPLES: usize = 160;
134
135        fn reset(&mut self) {
136            self.resets += 1;
137        }
138
139        fn probability(&mut self, _samples: &[i16]) -> Result<f32, VoiceActivityError> {
140            let score = self.scores[self.calls];
141            self.calls += 1;
142            Ok(score)
143        }
144    }
145
146    fn scripted(scores: &[f32]) -> Scripted {
147        Scripted {
148            scores: scores.to_vec(),
149            calls: 0,
150            resets: 0,
151        }
152    }
153
154    const FRAME: u64 = Scripted::FRAME_SAMPLES as u64;
155
156    fn ten_frames_of_silence() -> SilenceDuration {
157        SilenceDuration::new(Duration::from_millis(100)).expect("the duration is in range")
158    }
159
160    fn boundaries(scores: &[f32], batch: usize) -> Vec<Boundary> {
161        let mut endpointer = Endpointer::new(scripted(scores), ten_frames_of_silence());
162        let samples = vec![0; scores.len() * Scripted::FRAME_SAMPLES];
163        let mut found = Vec::new();
164        for piece in samples.chunks(batch) {
165            found.extend(endpointer.push(piece).expect("the scripted VAD scores"));
166        }
167        found
168    }
169
170    #[test]
171    fn consecutive_utterances_each_get_a_start_and_an_end() {
172        let mut scores = vec![0.0, 0.9, 0.9];
173        scores.extend([0.0; 10]);
174        scores.extend([0.9, 0.0]);
175        let ten_frames = 10 * FRAME;
176
177        for batch in [1, 100, 160, 1_024] {
178            assert_eq!(
179                boundaries(&scores, batch),
180                [
181                    Boundary::Start(FRAME),
182                    Boundary::End(3 * FRAME + ten_frames),
183                    Boundary::Start(13 * FRAME),
184                ],
185                "batches of {batch} samples"
186            );
187        }
188    }
189
190    #[test]
191    fn a_score_between_the_thresholds_does_not_end_the_utterance() {
192        let mut scores = vec![0.9];
193        scores.extend([0.4; 20]);
194
195        assert_eq!(boundaries(&scores, 160), [Boundary::Start(0)]);
196    }
197
198    #[test]
199    fn a_new_endpointer_resets_its_detector() {
200        let mut vad = scripted(&[]);
201
202        Endpointer::new(&mut vad, ten_frames_of_silence());
203
204        assert_eq!(vad.resets, 1);
205    }
206}
aria_kit/voice_activity/endpointer.rs

aria_kit/voice_activity/
endpointer.rs
