add vad crate

Author: yujonglee

Cargo.lock

@@ -0,0 +1,12 @@
+[package]
+name = "vad"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+serde = { workspace = true }
+thiserror = { workspace = true }
+tracing = { workspace = true }
+
+ndarray = "0.16"
+ort = "=2.0.0-rc.9"

crates/vad/Cargo.toml

@@ -0,0 +1,20 @@
+use serde::{ser::Serializer, Serialize};
+
+#[derive(Debug, thiserror::Error)]
+pub enum Error {
+    #[error(transparent)]
+    OrtError(#[from] ort::Error),
+    #[error(transparent)]
+    ShapeError(#[from] ndarray::ShapeError),
+    #[error("Invalid or missing output from model")]
+    InvalidOutput,
+}
+
+impl Serialize for Error {
+    fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
+    where
+        S: Serializer,
+    {
+        serializer.serialize_str(self.to_string().as_ref())
+    }
+}

crates/vad/assets/model.onnx

@@ -0,0 +1,119 @@
+mod error;
+use error::*;
+
+use ndarray::{Array1, Array2, Array3, ArrayBase, Ix1, Ix3, OwnedRepr};
+use ort::session::{builder::GraphOptimizationLevel, Session};
+
+const MODEL_BYTES: &[u8] =
+    include_bytes!(concat!(env!("CARGO_MANIFEST_DIR"), "/assets/model.onnx"));
+
+const SAMPLE_RATE: i64 = 16000;
+const CHUNK_SIZE_MS: usize = 30; // 30ms chunks for processing
+
+pub struct Vad {
+    session: Session,
+    h_tensor: ArrayBase<OwnedRepr<f32>, Ix3>,
+    c_tensor: ArrayBase<OwnedRepr<f32>, Ix3>,
+    sample_rate_tensor: ArrayBase<OwnedRepr<i64>, Ix1>,
+}
+
+impl Vad {
+    pub fn new() -> Result<Self, crate::Error> {
+        let session = Session::builder()?
+            .with_optimization_level(GraphOptimizationLevel::Level3)?
+            .with_intra_threads(4)?
+            .commit_from_memory(MODEL_BYTES)?;
+
+        let h_tensor = Array3::<f32>::zeros((2, 1, 64));
+        let c_tensor = Array3::<f32>::zeros((2, 1, 64));
+        let sample_rate_tensor = Array1::from_vec(vec![SAMPLE_RATE]);
+
+        Ok(Self {
+            session,
+            h_tensor,
+            c_tensor,
+            sample_rate_tensor,
+        })
+    }
+
+    /// Process a chunk of audio samples through the model and return the speech probability
+    fn forward(&mut self, audio_chunk: &[f32]) -> Result<f32, crate::Error> {
+        let samples = audio_chunk.len();
+        let audio_tensor = Array2::from_shape_vec((1, samples), audio_chunk.to_vec())?;
+
+        let mut result = self.session.run(ort::inputs![
+            audio_tensor.view(),
+            self.sample_rate_tensor.view(),
+            self.h_tensor.view(),
+            self.c_tensor.view()
+        ]?)?;
+
+        // Update internal state tensors
+        self.h_tensor = result
+            .get("hn")
+            .ok_or(Error::InvalidOutput)?
+            .try_extract_tensor::<f32>()?
+            .to_owned()
+            .into_shape_with_order((2, 1, 64))?;
+
+        self.c_tensor = result
+            .get("cn")
+            .ok_or(Error::InvalidOutput)?
+            .try_extract_tensor::<f32>()?
+            .to_owned()
+            .into_shape_with_order((2, 1, 64))?;
+
+        let prob_tensor = result.remove("output").ok_or(Error::InvalidOutput)?;
+        let prob = *prob_tensor
+            .try_extract_tensor::<f32>()?
+            .first()
+            .ok_or(Error::InvalidOutput)?;
+
+        Ok(prob)
+    }
+
+    /// For longer audio, this will process in 30ms chunks and return the maximum probability
+    pub fn run(&mut self, audio_samples: &[f32]) -> Result<f32, crate::Error> {
+        if audio_samples.len() < 480 {
+            return self.forward(audio_samples);
+        }
+
+        let chunk_size = (CHUNK_SIZE_MS * SAMPLE_RATE as usize) / 1000;
+        let num_chunks = audio_samples.len() / chunk_size;
+
+        let mut max_prob = 0.0f32;
+
+        for i in 0..num_chunks {
+            let start = i * chunk_size;
+            let end = (start + chunk_size).min(audio_samples.len());
+            let prob = self.forward(&audio_samples[start..end])?;
+            max_prob = max_prob.max(prob);
+        }
+
+        let remaining_start = num_chunks * chunk_size;
+        if remaining_start < audio_samples.len() && audio_samples.len() - remaining_start >= 240 {
+            let prob = self.forward(&audio_samples[remaining_start..])?;
+            max_prob = max_prob.max(prob);
+        }
+
+        Ok(max_prob)
+    }
+
+    pub fn reset(&mut self) {
+        self.h_tensor = Array3::<f32>::zeros((2, 1, 64));
+        self.c_tensor = Array3::<f32>::zeros((2, 1, 64));
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_vad() {
+        let mut vad = Vad::new().unwrap();
+        let audio_samples = vec![0.0; 16000];
+        let prob = vad.run(&audio_samples).unwrap();
+        assert!(prob < 0.1);
+    }
+}