path: root/crates/asist-audio/src/vad.rs

//! Voice activity detection and splitting into utterances.
//!
//! It is written as a pure state machine: it is fed samples and returns
//! events, with no threads or channels inside. That way the behaviour that is
//! hardest to debug by ear (when a turn starts, when silence ends it, when
//! the echo of the assistant's own speaker is ignored) can be fully tested
//! with synthetic audio and no microphone.

use std::collections::VecDeque;

use asist_core::config::VadConfig;

use crate::{rms, ASR_SAMPLE_RATE};

/// A closed utterance, ready to be transcribed.
#[derive(Debug, Clone)]
pub struct Utterance {
    pub samples: Vec<f32>,
    pub sample_rate: u32,
}

impl Utterance {
    pub fn duration_secs(&self) -> f32 {
        self.samples.len() as f32 / self.sample_rate as f32
    }
}

/// What the segmenter reports to the outside.
#[derive(Debug, Clone)]
pub enum VoiceEvent {
    /// Ha empezado a hablarse.
    Started,
    /// New audio inside the current utterance, for the partial
    /// transcriptions.
    Audio(Vec<f32>),
    /// Utterance finished and long enough to be transcribed.
    Ended(Utterance),
    /// Finished but too short: a knock on the table, a cough.
    Discarded,
    /// Speech was detected while the assistant was talking, with barge-in
    /// on. The orchestrator cuts playback when it gets this.
    BargeIn,
}

/// What the segmenter does with the microphone while the speaker plays.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Gate {
    /// Nothing is playing on the speaker: listen normally.
    Open,
    /// The assistant is talking. Depending on the configuration, input is
    /// either ignored (half duplex) or needs more volume to interrupt (barge-in).
    Speaking,
}

pub struct Segmenter {
    config: VadConfig,
    frame_samples: usize,
    preroll_samples: usize,
    silence_hold_samples: usize,
    min_utterance_samples: usize,
    max_utterance_samples: usize,

    pending: Vec<f32>,
    preroll: VecDeque<f32>,
    utterance: Vec<f32>,
    /// Samples of the utterance that really were above the threshold. The
    /// minimum is measured on this and not on `utterance`, which carries the
    /// preroll: otherwise 0.1 s of a knock on the table plus 0.2 s of preroll
    /// passes for a valid utterance and triggers a whole turn.
    voiced: usize,
    noise_floor: f32,
    silence_run: usize,
    speaking: bool,
    gate: Gate,
}

impl Segmenter {
    pub fn new(config: &VadConfig) -> Self {
        let rate = ASR_SAMPLE_RATE as f32;
        Self {
            frame_samples: (rate * config.frame_seconds).max(1.0) as usize,
            preroll_samples: (rate * config.preroll_seconds) as usize,
            silence_hold_samples: (rate * config.silence_hold) as usize,
            min_utterance_samples: (rate * config.min_utterance) as usize,
            max_utterance_samples: (rate * config.max_utterance) as usize,
            config: config.clone(),
            pending: Vec::new(),
            preroll: VecDeque::new(),
            utterance: Vec::new(),
            voiced: 0,
            noise_floor: 0.0,
            silence_run: 0,
            speaking: false,
            gate: Gate::Open,
        }
    }

    /// Opens or closes the microphone depending on whether the assistant talks.
    pub fn set_gate(&mut self, gate: Gate) {
        if self.gate == gate {
            return;
        }
        self.gate = gate;
        // When reopening after an answer, what has accumulated is the tail of
        // the assistant's own speaker: starting a turn with it would be a ghost turn.
        if gate == Gate::Open {
            self.pending.clear();
            self.preroll.clear();
            self.utterance.clear();
            self.voiced = 0;
            self.silence_run = 0;
            self.speaking = false;
        }
    }

    pub fn is_speaking(&self) -> bool {
        self.speaking
    }

    pub fn noise_floor(&self) -> f32 {
        self.noise_floor
    }

    /// Threshold that separates speech from silence right now.
    pub fn threshold(&self) -> f32 {
        let base = (self.noise_floor * self.config.threshold_factor)
            .clamp(self.config.min_threshold, self.config.max_threshold);
        // With the speaker playing you must speak louder to get through: the
        // microphone is hearing itself.
        if self.gate == Gate::Speaking {
            base * self.config.barge_in_factor
        } else {
            base
        }
    }

    /// Forcibly closes the current utterance (program shutdown).
    pub fn flush(&mut self) -> Option<Utterance> {
        if !self.speaking || self.voiced < self.min_utterance_samples {
            return None;
        }
        self.speaking = false;
        self.voiced = 0;
        Some(Utterance {
            samples: std::mem::take(&mut self.utterance),
            sample_rate: ASR_SAMPLE_RATE,
        })
    }

    /// Feeds 16 kHz mono audio and collects whatever needs doing.
    pub fn push(&mut self, samples: &[f32]) -> Vec<VoiceEvent> {
        let mut events = Vec::new();
        // In half duplex the microphone is effectively off: without this, the
        // assistant transcribes itself and answers itself.
        if self.gate == Gate::Speaking && !self.config.barge_in {
            return events;
        }
        self.pending.extend_from_slice(samples);

        while self.pending.len() >= self.frame_samples {
            let frame: Vec<f32> = self.pending.drain(..self.frame_samples).collect();
            let level = rms(&frame);
            self.track_noise_floor(level);
            let threshold = self.threshold();

            if level < threshold && !self.speaking {
                self.preroll.extend(frame.iter().copied());
                while self.preroll.len() > self.preroll_samples {
                    self.preroll.pop_front();
                }
                continue;
            }

            if !self.speaking {
                self.speaking = true;
                self.utterance.clear();
                self.voiced = 0;
                self.utterance.extend(self.preroll.drain(..));
                if self.gate == Gate::Speaking {
                    events.push(VoiceEvent::BargeIn);
                }
                events.push(VoiceEvent::Started);
            }

            if level < threshold {
                self.silence_run += frame.len();
            } else {
                self.silence_run = 0;
                self.voiced += frame.len();
            }
            self.utterance.extend_from_slice(&frame);

            let ended = self.silence_run >= self.silence_hold_samples;
            let too_long = self.utterance.len() >= self.max_utterance_samples;
            if !ended && !too_long {
                events.push(VoiceEvent::Audio(frame));
                continue;
            }

            let long_enough = self.voiced >= self.min_utterance_samples;
            events.push(if long_enough {
                VoiceEvent::Ended(Utterance {
                    samples: std::mem::take(&mut self.utterance),
                    sample_rate: ASR_SAMPLE_RATE,
                })
            } else {
                VoiceEvent::Discarded
            });

            self.utterance.clear();
            self.voiced = 0;
            self.preroll.clear();
            self.silence_run = 0;
            // A length cut lands mid-sentence: keep listening as if the user had
            // not stopped talking, which is the truth.
            self.speaking = too_long && !ended;
            if self.speaking {
                events.push(VoiceEvent::Started);
            }
        }
        events
    }

    /// The noise floor only goes down: sustained speech must not be able to
    /// drag the threshold above itself and stop being detected.
    fn track_noise_floor(&mut self, level: f32) {
        if self.noise_floor == 0.0 {
            self.noise_floor = level;
        } else if level < self.noise_floor * 1.5 {
            self.noise_floor = self.noise_floor * 0.95 + level * 0.05;
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn config() -> VadConfig {
        VadConfig {
            frame_seconds: 0.1,
            preroll_seconds: 0.2,
            silence_hold: 0.3,
            min_utterance: 0.3,
            max_utterance: 2.0,
            threshold_factor: 3.0,
            min_threshold: 0.001,
            max_threshold: 0.02,
            barge_in: false,
            barge_in_factor: 4.0,
        }
    }

    fn samples(secs: f32, amplitude: f32) -> Vec<f32> {
        let n = (ASR_SAMPLE_RATE as f32 * secs) as usize;
        // Alternates sign so the RMS is the amplitude and not a DC level.
        (0..n)
            .map(|i| if i % 2 == 0 { amplitude } else { -amplitude })