aria_kit/voice_activity/
map.rs

1//! Where anyone was speaking in one recording, recorded frame by frame from
2//! one [`Endpointer`] run, so every reader answers from the same detector.
3//!
4//! Positions are stream samples. Frames start where the detector starts
5//! and follow each other without gaps, so a frame is named by where it
6//! starts or ends.
7
8use std::ops::Range;
9use std::time::Duration;
10
11use super::{Boundary, Endpointer, UTTERANCE_GAP, VoiceActivityDetector, VoiceActivityError};
12use crate::audio::{SilenceDuration, canonical_samples, pcm16_from_f32};
13
14/// A silence inside an utterance at least this long is a pause: a gap a
15/// voice could take over across. Shorter ones fall between syllables.
16const PAUSE_SAMPLES: u64 = canonical_samples(Duration::from_millis(150));
17/// The newest pauses of an utterance a reader is shown. A reader looks back
18/// a second or two, which holds fewer than this.
19const PAUSES_KEPT: usize = 8;
20
21/// Whether more of the recording is still to come.
22#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
23pub(crate) enum Progress {
24    #[default]
25    Ongoing,
26    /// The whole recording was heard.
27    Finished,
28    /// Hearing failed: nothing after `reached` can be told apart.
29    Broken,
30}
31
32#[derive(Debug, Clone, PartialEq, Eq)]
33struct Utterance {
34    start: u64,
35    /// The end of the frame that completed its trailing silence.
36    end: Option<u64>,
37}
38
39/// A run of the trailing-silence clock inside an utterance.
40#[derive(Debug, Clone, PartialEq, Eq)]
41struct Quiet {
42    start: u64,
43    end: Option<u64>,
44    /// Speech stopped the clock, rather than the utterance ending.
45    resumed: bool,
46}
47
48/// Where anyone was speaking, as far as the detector has scored.
49#[derive(Debug, Clone, Default)]
50pub struct SpeechMap {
51    frame_samples: usize,
52    /// Where the first frame starts.
53    from: u64,
54    reached: u64,
55    utterances: Vec<Utterance>,
56    quiet: Vec<Quiet>,
57    /// Frames scored as speech, joined where they touch.
58    speech: Vec<Range<u64>>,
59    /// Speech in what was forgotten, so counts from the start stay whole.
60    forgotten_speech: u64,
61    progress: Progress,
62}
63
64impl SpeechMap {
65    fn new(frame_samples: usize, from: u64) -> Self {
66        Self {
67            frame_samples,
68            from,
69            reached: from,
70            ..Self::default()
71        }
72    }
73
74    pub(crate) fn frame_samples(&self) -> usize {
75        self.frame_samples
76    }
77
78    /// The end of the last frame scored.
79    pub fn reached(&self) -> u64 {
80        self.reached
81    }
82
83    /// The start of the frame that holds `at`.
84    pub(crate) fn frame_start(&self, at: u64) -> u64 {
85        let frame = self.frame_samples.max(1) as u64;
86        self.from + at.saturating_sub(self.from) / frame * frame
87    }
88
89    /// The detector failed: nothing after `reached` can be told apart.
90    pub fn is_broken(&self) -> bool {
91        self.progress == Progress::Broken
92    }
93
94    /// Whether the detector heard the whole recording, so nothing after
95    /// `reached` is still to come.
96    pub fn is_finished(&self) -> bool {
97        self.progress == Progress::Finished
98    }
99
100    /// Marks the map broken where it stands, as when its detector stopped
101    /// without saying why.
102    pub fn break_off(&mut self) {
103        self.progress = Progress::Broken;
104    }
105
106    /// Whether an utterance opened on the frame starting at `at`.
107    pub(crate) fn opens_at(&self, at: u64) -> bool {
108        self.utterances
109            .iter()
110            .rev()
111            .take_while(|utterance| utterance.start >= at)
112            .any(|utterance| utterance.start == at)
113    }
114
115    /// Whether an utterance closed on the frame ending at `at`.
116    pub(crate) fn closes_at(&self, at: u64) -> bool {
117        self.utterances
118            .iter()
119            .rev()
120            .take_while(|utterance| utterance.end.is_none_or(|end| end >= at))
121            .any(|utterance| utterance.end == Some(at))
122    }
123
124    /// How much silence the open utterance ends in, once the frame ending
125    /// at `at` is scored. Nothing while speech goes on, or outside an
126    /// utterance.
127    pub(crate) fn trailing_silence(&self, at: u64) -> u64 {
128        self.quiet
129            .iter()
130            .rev()
131            .find(|quiet| quiet.start < at)
132            .filter(|quiet| match quiet.end {
133                None => true,
134                // Speech scored at `end` is what stops the clock, so the
135                // frame ending there still ends in silence.
136                Some(end) if quiet.resumed => at <= end,
137                Some(end) => at < end,
138            })
139            .map_or(0, |quiet| at - quiet.start)
140    }
141
142    /// The pauses inside the utterance still open once the frame ending at
143    /// `at` is scored, newest first.
144    pub(crate) fn pauses_before(&self, at: u64) -> impl Iterator<Item = Range<u64>> + '_ {
145        let open = self
146            .utterances
147            .iter()
148            .rev()
149            .find(|utterance| utterance.start < at)
150            .filter(|utterance| utterance.end.is_none_or(|end| at < end))
151            .map(|utterance| utterance.start);
152        self.quiet
153            .iter()
154            .rev()
155            .filter_map(|quiet| match quiet.end {
156                Some(end) if quiet.resumed && end - quiet.start >= PAUSE_SAMPLES => {
157                    Some(quiet.start..end)
158                }
159                _ => None,
160            })
161            .filter(move |pause| pause.end < at)
162            .take_while(move |pause| open.is_some_and(|start| pause.start >= start))
163            .take(PAUSES_KEPT)
164    }
165
166    /// Samples of speech scored before `at`.
167    pub fn speech_before(&self, at: u64) -> u64 {
168        self.forgotten_speech + self.speech(&(0..at))
169    }
170
171    pub(crate) fn forget_before(&mut self, at: u64) {
172        let mut forgotten = 0;
173        self.speech.retain_mut(|run| {
174            if run.end <= at {
175                forgotten += run.end - run.start;
176                return false;
177            }
178            if run.start < at {
179                forgotten += at - run.start;
180                run.start = at;
181            }
182            true
183        });
184        self.forgotten_speech += forgotten;
185        self.utterances
186            .retain(|utterance| utterance.end.is_none_or(|end| end > at));
187        self.quiet
188            .retain(|quiet| quiet.end.is_none_or(|end| end > at));
189    }
190
191    /// Samples of speech scored inside `span`. Speech already forgotten does
192    /// not count.
193    pub fn speech(&self, span: &Range<u64>) -> u64 {
194        self.speech
195            .iter()
196            .map(|run| {
197                run.end
198                    .min(span.end)
199                    .saturating_sub(run.start.max(span.start))
200            })
201            .sum()
202    }
203
204    fn add_speech(&mut self, frame: Range<u64>) {
205        match self.speech.last_mut() {
206            Some(run) if run.end == frame.start => run.end = frame.end,
207            _ => self.speech.push(frame),
208        }
209    }
210
211    fn close_quiet(&mut self, end: u64, resumed: bool) {
212        if let Some(quiet) = self.quiet.last_mut().filter(|quiet| quiet.end.is_none()) {
213            quiet.end = Some(end);
214            quiet.resumed = resumed;
215        }
216    }
217}
218
219/// Scores a recording's frames and records them in a [`SpeechMap`].
220pub struct SpeechDetector<V> {
221    endpointer: Endpointer<V>,
222    frame: Vec<i16>,
223    map: SpeechMap,
224}
225
226impl<V: VoiceActivityDetector> SpeechDetector<V> {
227    /// A detector for one recording. The first sample heard sits at `at`.
228    pub fn new(vad: V, at: u64) -> Self {
229        let gap = SilenceDuration::new(UTTERANCE_GAP).expect("the utterance gap is a silence");
230        Self::with_gap(vad, gap, at)
231    }
232
233    fn with_gap(vad: V, utterance_gap: SilenceDuration, at: u64) -> Self {
234        Self {
235            endpointer: Endpointer::new(vad, utterance_gap),
236            frame: Vec::with_capacity(V::FRAME_SAMPLES),
237            map: SpeechMap::new(V::FRAME_SAMPLES, at),
238        }
239    }
240
241    pub fn map(&self) -> &SpeechMap {
242        &self.map
243    }
244
245    /// Forgets the map before `at`: no reader asks about it again. Only the
246    /// count of speech in it stays.
247    pub(crate) fn forget_before(&mut self, at: u64) {
248        self.map.forget_before(at);
249    }
250
251    /// The recording ended after the audio heard so far. A last partial
252    /// frame is never scored.
253    pub(crate) fn finish(&mut self) {
254        if self.map.progress == Progress::Ongoing {
255            self.map.progress = Progress::Finished;
256        }
257    }
258
259    /// Scores `audio`, which continues the stream where the last call left
260    /// off. A partial frame waits for the next call.
261    pub fn hear(&mut self, mut audio: &[f32]) -> Result<(), VoiceActivityError> {
262        while !audio.is_empty() && !self.map.is_broken() {
263            let taken = audio.len().min(V::FRAME_SAMPLES - self.frame.len());
264            self.frame
265                .extend(audio[..taken].iter().copied().map(pcm16_from_f32));
266            audio = &audio[taken..];
267            if self.frame.len() == V::FRAME_SAMPLES {
268                self.score_frame()?;
269                self.frame.clear();
270            }
271        }
272        Ok(())
273    }
274
275    fn score_frame(&mut self) -> Result<(), VoiceActivityError> {
276        let start = self.map.reached;
277        let speech_before = self.endpointer.speech_samples();
278        let quiet_before = self.endpointer.silence_since();
279        let boundaries = self
280            .endpointer
281            .push(&self.frame)
282            .inspect_err(|_| self.map.break_off())?;
283        let end = start + V::FRAME_SAMPLES as u64;
284        self.map.reached = end;
285        if self.endpointer.speech_samples() > speech_before {
286            self.map.add_speech(start..end);
287        }
288        match boundaries.first() {
289            Some(Boundary::Start(_)) => self.map.utterances.push(Utterance { start, end: None }),
290            Some(Boundary::End(_)) => {
291                if let Some(utterance) = self.map.utterances.last_mut() {
292                    utterance.end = Some(end);
293                }
294                self.map.close_quiet(end, false);
295            }
296            None => match (quiet_before, self.endpointer.silence_since()) {
297                (None, Some(since)) => self.map.quiet.push(Quiet {
298                    // The endpointer counts from where the map starts.
299                    start: self.map.from + since,
300                    end: None,
301                    resumed: false,
302                }),
303                (Some(_), None) => self.map.close_quiet(start, true),
304                (None, None) | (Some(_), Some(_)) => {}
305            },
306        }
307        Ok(())
308    }
309}
310
311#[cfg(test)]
312mod tests {
313    use super::*;
314
315    struct Scripted {
316        scores: Vec<f32>,
317        calls: usize,
318    }
319
320    impl VoiceActivityDetector for Scripted {
321        const FRAME_SAMPLES: usize = 160;
322
323        fn probability(&mut self, _samples: &[i16]) -> Result<f32, VoiceActivityError> {
324            let score = self.scores[self.calls];
325            self.calls += 1;
326            Ok(score)
327        }
328    }
329
330    const FRAME: u64 = Scripted::FRAME_SAMPLES as u64;
331
332    fn detector(scores: &[f32], gap_frames: u64, at: u64) -> SpeechDetector<Scripted> {
333        let gap = SilenceDuration::new(Duration::from_millis(gap_frames * 10))
334            .expect("the duration is in range");
335        SpeechDetector::with_gap(
336            Scripted {
337                scores: scores.to_vec(),
338                calls: 0,
339            },
340            gap,
341            at,
342        )
343    }
344
345    fn heard(scores: &[f32], gap_frames: u64, piece: usize) -> SpeechMap {
346        let mut detector = detector(scores, gap_frames, 0);
347        let audio = vec![0.0; scores.len() * Scripted::FRAME_SAMPLES];
348        for chunk in audio.chunks(piece) {
349            detector.hear(chunk).expect("the scripted VAD scores");
350        }
351        detector.map().clone()
352    }
353
354    #[test]
355    fn a_pause_is_a_silence_too_short_to_end_the_utterance() {
356        let mut scores = vec![0.9];
357        scores.extend([0.0; 15]);
358        scores.push(0.9);
359        scores.extend([0.0; 5]);
360        scores.push(0.9);
361        scores.extend([0.0; 30]);
362        let map = heard(&scores, 30, 160);
363        let pause = FRAME..16 * FRAME;
364
365        assert_eq!(map.pauses_before(22 * FRAME).collect::<Vec<_>>(), [pause]);
366        assert_eq!(map.trailing_silence(22 * FRAME), 5 * FRAME);
367        assert_eq!(map.trailing_silence(23 * FRAME), 0);
368        assert!(map.closes_at(53 * FRAME));
369        assert_eq!(map.pauses_before(53 * FRAME).count(), 0);
370    }
371
372    #[test]
373    fn utterances_open_and_close_on_their_frames() {
374        let mut scores = vec![0.0, 0.9, 0.9];
375        scores.extend([0.0; 10]);
376        scores.extend([0.9, 0.0]);
377
378        for piece in [1, 100, 160, 1_024] {
379            let map = heard(&scores, 10, piece);
380            assert!(map.opens_at(FRAME), "pieces of {piece}");
381            assert!(map.closes_at(13 * FRAME), "pieces of {piece}");
382            assert!(map.opens_at(13 * FRAME), "pieces of {piece}");
383            assert!(!map.opens_at(0) && !map.closes_at(3 * FRAME));
384        }
385    }
386
387    #[test]
388    fn speech_counts_only_frames_scored_as_speech() {
389        let map = heard(&[0.0, 0.9, 0.4, 0.9, 0.0], 10, 160);
390
391        assert_eq!(map.speech_before(5 * FRAME), 2 * FRAME);
392        assert_eq!(map.speech(&(FRAME..2 * FRAME)), FRAME);
393        assert_eq!(map.speech(&(2 * FRAME..3 * FRAME)), 0);
394    }
395
396    #[test]
397    fn positions_count_from_where_the_detector_starts() {
398        let mut detector = detector(&[0.9], 10, 1_000);
399        detector.hear(&[0.0; 160]).unwrap();
400
401        let map = detector.map();
402        assert_eq!(map.reached(), 1_000 + FRAME);
403        assert!(map.opens_at(1_000));
404        assert_eq!(map.speech_before(map.reached()), FRAME);
405    }
406
407    #[test]
408    fn forgetting_keeps_what_comes_after_and_the_count_before() {
409        let mut scores = vec![0.9; 5];
410        scores.extend([0.0; 15]);
411        scores.extend([0.9; 5]);
412        scores.extend([0.0; 15]);
413        let mut map = heard(&scores, 10, 160);
414        let before = map.speech_before(map.reached());
415
416        map.forget_before(20 * FRAME);
417
418        assert_eq!(map.speech_before(map.reached()), before);
419        assert_eq!(map.speech(&(20 * FRAME..25 * FRAME)), 5 * FRAME);
420        assert!(map.opens_at(20 * FRAME));
421        assert!(!map.closes_at(15 * FRAME));
422    }
423}
aria_kit/voice_activity/map.rs

aria_kit/voice_activity/
map.rs
