453 lines
16 KiB
Rust
453 lines
16 KiB
Rust
//! TEN VAD ONNX backend.
|
|
//!
|
|
//! TEN's ONNX graph does not accept raw PCM. It expects the same feature
|
|
//! stack produced by TEN's `AUP_Aed_aivad_proc`: three context frames of
|
|
//! 40 log-mel powers plus one pitch feature, followed by four recurrent
|
|
//! state tensors. This module ports that preprocessing path to Rust and
|
|
//! keeps ONNX Runtime off the realtime callback where possible.
|
|
|
|
use crate::frame::f32_to_i16;
|
|
|
|
use rustfft::{num_complex::Complex32, FftPlanner};
|
|
|
|
use super::resampler::{Downsampler48to16, INPUT_FRAME_10MS};
|
|
use super::{VadOutput, VoiceActivityDetector};
|
|
|
|
const SAMPLE_RATE_16K: f32 = 16_000.0;
|
|
const HOP_16K: usize = 256;
|
|
const WINDOW_16K: usize = 768;
|
|
const FFT_SIZE: usize = 1024;
|
|
const N_BINS: usize = FFT_SIZE / 2 + 1;
|
|
const MEL_BANDS: usize = 40;
|
|
const FEATURE_LEN: usize = 41;
|
|
const CONTEXT: usize = 3;
|
|
const HIDDEN: usize = 64;
|
|
const POWER_NORMALIZER: f32 = 32768.0 * 32768.0;
|
|
const EPS: f32 = 1.0e-20;
|
|
|
|
const FEATURE_MEANS: [f32; FEATURE_LEN] = [
|
|
-8.198236, -6.2657166, -5.4838185, -4.7586913, -4.417089, -4.142893, -3.9128504, -3.845928,
|
|
-3.6570904, -3.7234187, -3.8761342, -3.843891, -3.6904051, -3.7560658, -3.6986961, -3.650463,
|
|
-3.7004688, -3.5673213, -3.4989002, -3.477807, -3.458816, -3.4449239, -3.4013286, -3.3062613,
|
|
-3.2785568, -3.2332509, -3.198616, -3.2045264, -3.2087986, -3.257838, -3.3813767, -3.5340214,
|
|
-3.640868, -3.7268589, -3.773731, -3.8046672, -3.832901, -3.8711205, -3.990593, -4.4802895,
|
|
92.3569,
|
|
];
|
|
|
|
const FEATURE_STDS: [f32; FEATURE_LEN] = [
|
|
5.166064, 4.9772096, 4.698896, 4.6306214, 4.634348, 4.641156, 4.6406765, 4.666367, 4.6505346,
|
|
4.640021, 4.6374, 4.620099, 4.5963163, 4.562655, 4.5543604, 4.5669107, 4.56249, 4.5624127,
|
|
4.5852995, 4.6001797, 4.592846, 4.5859227, 4.5834966, 4.626093, 4.626958, 4.6262894, 4.637006,
|
|
4.683016, 4.726814, 4.7342896, 4.753227, 4.849723, 4.869435, 4.884483, 4.921327, 4.9592123,
|
|
4.996619, 5.0448236, 5.072217, 5.0964394, 115.21369,
|
|
];
|
|
|
|
/// TEN VAD using ONNX Runtime and Rust-ported TEN feature preprocessing.
|
|
pub struct TenOnnxVad {
|
|
session: ort::session::Session,
|
|
downsampler: Downsampler48to16,
|
|
hop_accum: Vec<f32>,
|
|
sample_fifo: Vec<f32>,
|
|
feature_stack: [[f32; FEATURE_LEN]; CONTEXT],
|
|
states: [[f32; HIDDEN]; 4],
|
|
mel_filters: Vec<[f32; N_BINS]>,
|
|
fft: std::sync::Arc<dyn rustfft::Fft<f32>>,
|
|
fft_buffer: Vec<Complex32>,
|
|
last_probability: f32,
|
|
last_speech: bool,
|
|
}
|
|
|
|
unsafe impl Send for TenOnnxVad {}
|
|
|
|
impl TenOnnxVad {
|
|
/// Load TEN VAD ONNX model.
|
|
pub fn try_new(model_path: &str) -> Option<Self> {
|
|
if !std::path::Path::new(model_path).exists() {
|
|
tracing::warn!(target: "chanora_audio", path = model_path, "TEN VAD ONNX model not found");
|
|
return None;
|
|
}
|
|
let session = match std::panic::catch_unwind(|| {
|
|
ort::session::Session::builder().and_then(|mut b| b.commit_from_file(model_path))
|
|
}) {
|
|
Ok(Ok(session)) => session,
|
|
Ok(Err(error)) => {
|
|
tracing::warn!(target: "chanora_audio", %error, path = model_path, "TEN VAD ONNX model load failed");
|
|
return None;
|
|
}
|
|
Err(_) => {
|
|
tracing::warn!(target: "chanora_audio", path = model_path, "TEN VAD ONNX Runtime panicked during load");
|
|
return None;
|
|
}
|
|
};
|
|
let mut fft_planner = FftPlanner::<f32>::new();
|
|
let fft = fft_planner.plan_fft_forward(FFT_SIZE);
|
|
tracing::info!(target: "chanora_audio", path = model_path, "TEN VAD ONNX model loaded");
|
|
Some(Self {
|
|
session,
|
|
downsampler: Downsampler48to16::default(),
|
|
hop_accum: Vec::with_capacity(HOP_16K + super::resampler::OUTPUT_FRAME_10MS),
|
|
sample_fifo: Vec::with_capacity(WINDOW_16K + HOP_16K),
|
|
feature_stack: [[0.0; FEATURE_LEN]; CONTEXT],
|
|
states: [[0.0; HIDDEN]; 4],
|
|
mel_filters: build_mel_filters(),
|
|
fft,
|
|
fft_buffer: vec![Complex32::ZERO; FFT_SIZE],
|
|
last_probability: 0.0,
|
|
last_speech: false,
|
|
})
|
|
}
|
|
|
|
fn process_hop(&mut self, hop: &[f32]) {
|
|
self.sample_fifo.extend_from_slice(hop);
|
|
let frame = if self.sample_fifo.len() >= WINDOW_16K {
|
|
let start = self.sample_fifo.len() - WINDOW_16K;
|
|
self.sample_fifo[start..].to_vec()
|
|
} else {
|
|
let mut padded = vec![0.0; WINDOW_16K - self.sample_fifo.len()];
|
|
padded.extend_from_slice(&self.sample_fifo);
|
|
padded
|
|
};
|
|
if self.sample_fifo.len() > WINDOW_16K {
|
|
let excess = self.sample_fifo.len() - WINDOW_16K;
|
|
self.sample_fifo.drain(..excess);
|
|
}
|
|
|
|
let feature = compute_feature(
|
|
&self.mel_filters,
|
|
self.fft.as_ref(),
|
|
&mut self.fft_buffer,
|
|
&frame,
|
|
);
|
|
self.feature_stack.copy_within(1..CONTEXT, 0);
|
|
self.feature_stack[CONTEXT - 1] = feature;
|
|
self.run_onnx();
|
|
}
|
|
|
|
fn run_onnx(&mut self) {
|
|
use ndarray::{Array, IxDyn};
|
|
use ort::value::Value;
|
|
|
|
let input: Vec<f32> = self.feature_stack.iter().flatten().copied().collect();
|
|
let input_arr = match Array::from_shape_vec(IxDyn(&[1, CONTEXT, FEATURE_LEN]), input) {
|
|
Ok(v) => v,
|
|
Err(_) => return,
|
|
};
|
|
let state_arrs = [0, 1, 2, 3]
|
|
.map(|idx| Array::from_shape_vec(IxDyn(&[1, HIDDEN]), self.states[idx].to_vec()));
|
|
let input_val = match Value::from_array(input_arr) {
|
|
Ok(v) => v,
|
|
Err(error) => {
|
|
tracing::warn!(target: "chanora_audio", %error, "TEN VAD input tensor error");
|
|
return;
|
|
}
|
|
};
|
|
let state_vals = match state_arrs {
|
|
[Ok(a), Ok(b), Ok(c), Ok(d)] => [a, b, c, d],
|
|
_ => return,
|
|
};
|
|
let state_vals = match state_vals.map(Value::from_array) {
|
|
[Ok(a), Ok(b), Ok(c), Ok(d)] => [a, b, c, d],
|
|
_ => return,
|
|
};
|
|
|
|
let outputs = match self.session.run([
|
|
(&input_val).into(),
|
|
(&state_vals[0]).into(),
|
|
(&state_vals[1]).into(),
|
|
(&state_vals[2]).into(),
|
|
(&state_vals[3]).into(),
|
|
]) {
|
|
Ok(outputs) => outputs,
|
|
Err(error) => {
|
|
tracing::warn!(target: "chanora_audio", %error, "TEN VAD ONNX inference failed");
|
|
return;
|
|
}
|
|
};
|
|
|
|
if let Ok((_, prob)) = outputs["output_1"].try_extract_tensor::<f32>() {
|
|
if let Some(&p) = prob.first() {
|
|
self.last_probability = p.clamp(0.0, 1.0);
|
|
self.last_speech = self.last_probability >= 0.5;
|
|
}
|
|
}
|
|
for (idx, name) in ["output_2", "output_3", "output_6", "output_7"]
|
|
.iter()
|
|
.enumerate()
|
|
{
|
|
if let Ok((_, state)) = outputs[*name].try_extract_tensor::<f32>() {
|
|
let copy_len = state.len().min(HIDDEN);
|
|
self.states[idx][..copy_len].copy_from_slice(&state[..copy_len]);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
fn compute_feature(
|
|
mel_filters: &[[f32; N_BINS]],
|
|
fft: &dyn rustfft::Fft<f32>,
|
|
fft_buffer: &mut [Complex32],
|
|
frame: &[f32],
|
|
) -> [f32; FEATURE_LEN] {
|
|
let power = power_spectrum(fft, fft_buffer, frame);
|
|
let mut feature = [0.0; FEATURE_LEN];
|
|
for band in 0..MEL_BANDS {
|
|
let energy = mel_filters[band]
|
|
.iter()
|
|
.zip(power.iter())
|
|
.map(|(w, p)| w * p)
|
|
.sum::<f32>()
|
|
/ POWER_NORMALIZER;
|
|
let log_energy = (energy + EPS).ln();
|
|
feature[band] = (log_energy - FEATURE_MEANS[band]) / (FEATURE_STDS[band] + EPS);
|
|
}
|
|
let pitch_hz = estimate_pitch_hz(frame);
|
|
feature[MEL_BANDS] = (pitch_hz - FEATURE_MEANS[MEL_BANDS]) / (FEATURE_STDS[MEL_BANDS] + EPS);
|
|
feature
|
|
}
|
|
|
|
impl VoiceActivityDetector for TenOnnxVad {
|
|
fn process_10ms(&mut self, samples: &[f32]) -> VadOutput {
|
|
debug_assert_eq!(samples.len(), INPUT_FRAME_10MS);
|
|
let mut input = [0.0_f32; INPUT_FRAME_10MS];
|
|
input.copy_from_slice(samples);
|
|
let downsampled = self.downsampler.process_frame_10ms(&input);
|
|
self.hop_accum.extend_from_slice(&downsampled);
|
|
while self.hop_accum.len() >= HOP_16K {
|
|
let hop: Vec<f32> = self.hop_accum[..HOP_16K].to_vec();
|
|
self.hop_accum.drain(..HOP_16K);
|
|
self.process_hop(&hop);
|
|
}
|
|
VadOutput {
|
|
probability: self.last_probability,
|
|
speech: self.last_speech,
|
|
}
|
|
}
|
|
}
|
|
|
|
fn hz_to_mel(hz: f32) -> f32 {
|
|
2595.0 * (1.0 + hz / 700.0).log10()
|
|
}
|
|
|
|
fn mel_to_hz(mel: f32) -> f32 {
|
|
700.0 * (10.0_f32.powf(mel / 2595.0) - 1.0)
|
|
}
|
|
|
|
fn build_mel_filters() -> Vec<[f32; N_BINS]> {
|
|
let low_mel = hz_to_mel(0.0);
|
|
let high_mel = hz_to_mel(8000.0);
|
|
let mut bins = [0_usize; MEL_BANDS + 2];
|
|
for idx in 0..bins.len() {
|
|
let mel = idx as f32 * (high_mel - low_mel) / (MEL_BANDS as f32 + 1.0) + low_mel;
|
|
let hz = mel_to_hz(mel);
|
|
let mut bin = ((FFT_SIZE as f32 + 1.0) * hz / SAMPLE_RATE_16K).floor() as usize;
|
|
bin = bin.min(N_BINS - 1);
|
|
if idx > 0 && bin == bins[idx - 1] {
|
|
bin = (bin + 1).min(N_BINS - 1);
|
|
}
|
|
bins[idx] = bin;
|
|
}
|
|
|
|
let mut filters = vec![[0.0_f32; N_BINS]; MEL_BANDS];
|
|
for band in 0..MEL_BANDS {
|
|
let left = bins[band];
|
|
let center = bins[band + 1].max(left + 1);
|
|
let right = bins[band + 2].max(center + 1).min(N_BINS - 1);
|
|
for (i, weight) in filters[band]
|
|
.iter_mut()
|
|
.enumerate()
|
|
.take(center.min(N_BINS))
|
|
.skip(left)
|
|
{
|
|
*weight = (i - left) as f32 / (center - left) as f32;
|
|
}
|
|
for (i, weight) in filters[band]
|
|
.iter_mut()
|
|
.enumerate()
|
|
.take(right + 1)
|
|
.skip(center)
|
|
{
|
|
*weight = (right - i) as f32 / (right - center).max(1) as f32;
|
|
}
|
|
}
|
|
filters
|
|
}
|
|
|
|
fn power_spectrum(
|
|
fft: &dyn rustfft::Fft<f32>,
|
|
fft_buffer: &mut [Complex32],
|
|
frame: &[f32],
|
|
) -> [f32; N_BINS] {
|
|
debug_assert_eq!(fft_buffer.len(), FFT_SIZE);
|
|
fft_buffer.fill(Complex32::ZERO);
|
|
for (idx, sample) in frame.iter().take(WINDOW_16K).enumerate() {
|
|
let hann = 0.5 - 0.5 * (2.0 * std::f32::consts::PI * idx as f32 / WINDOW_16K as f32).cos();
|
|
fft_buffer[idx].re = f32_to_i16(*sample) as f32 * hann;
|
|
}
|
|
|
|
fft.process(fft_buffer);
|
|
|
|
let mut out = [0.0_f32; N_BINS];
|
|
for (dst, bin) in out.iter_mut().zip(fft_buffer.iter()) {
|
|
*dst = bin.norm_sqr();
|
|
}
|
|
out
|
|
}
|
|
|
|
fn estimate_pitch_hz(frame: &[f32]) -> f32 {
|
|
let min_lag = (SAMPLE_RATE_16K / 400.0) as usize;
|
|
let max_lag = (SAMPLE_RATE_16K / 60.0) as usize;
|
|
let mut best_lag = 0_usize;
|
|
let mut best_corr = 0.0_f32;
|
|
for lag in min_lag..=max_lag.min(frame.len().saturating_sub(1)) {
|
|
let mut corr = 0.0_f32;
|
|
let mut energy = 0.0_f32;
|
|
for i in lag..frame.len() {
|
|
corr += frame[i] * frame[i - lag];
|
|
energy += frame[i - lag] * frame[i - lag];
|
|
}
|
|
let norm = if energy > 1.0e-8 {
|
|
corr / energy.sqrt()
|
|
} else {
|
|
0.0
|
|
};
|
|
if norm > best_corr {
|
|
best_corr = norm;
|
|
best_lag = lag;
|
|
}
|
|
}
|
|
if best_lag == 0 || best_corr < 0.01 {
|
|
0.0
|
|
} else {
|
|
SAMPLE_RATE_16K / best_lag as f32
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Background worker — same pattern as SileroOnnxVadWorker so the realtime
|
|
// callback never blocks on STFT / pitch / ONNX inference.
|
|
// ---------------------------------------------------------------------------
|
|
|
|
use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64};
|
|
use std::sync::Arc;
|
|
use std::thread::JoinHandle;
|
|
|
|
/// Maximum number of 10 ms frames the worker may lag before the callback
|
|
/// treats its output as stale and uses WebRTC fallback instead.
|
|
const TEN_MAX_STALE_FRAMES: u64 = 8;
|
|
|
|
struct TenFrameMessage {
|
|
seq: u64,
|
|
frame: [f32; INPUT_FRAME_10MS],
|
|
}
|
|
|
|
/// Background TEN VAD worker. The realtime callback only enqueues 10 ms
|
|
/// frames and reads the latest probability atomically.
|
|
pub struct TenOnnxVadWorker {
|
|
tx: Option<std::sync::mpsc::SyncSender<TenFrameMessage>>,
|
|
latest_probability: Arc<AtomicU32>,
|
|
latest_processed_seq: Arc<AtomicU64>,
|
|
alive: Arc<AtomicBool>,
|
|
handle: Option<JoinHandle<()>>,
|
|
}
|
|
|
|
impl TenOnnxVadWorker {
|
|
/// Start a background TEN worker if the model loads.
|
|
pub fn try_new(model_path: &str) -> Option<Self> {
|
|
let vad = TenOnnxVad::try_new(model_path)?;
|
|
let latest_probability = Arc::new(AtomicU32::new(0.0_f32.to_bits()));
|
|
let latest_processed_seq = Arc::new(AtomicU64::new(u64::MAX));
|
|
let alive = Arc::new(AtomicBool::new(true));
|
|
let (tx, rx) = std::sync::mpsc::sync_channel::<TenFrameMessage>(128);
|
|
let prob_arc = latest_probability.clone();
|
|
let seq_arc = latest_processed_seq.clone();
|
|
let alive_arc = alive.clone();
|
|
|
|
let handle = std::thread::Builder::new()
|
|
.name("chanora-ten-vad".to_string())
|
|
.spawn(move || {
|
|
let mut vad = vad;
|
|
while alive_arc.load(std::sync::atomic::Ordering::Relaxed) {
|
|
let msg = match rx.recv() {
|
|
Ok(m) => m,
|
|
Err(_) => break,
|
|
};
|
|
let mut frame_f32 = [0.0_f32; INPUT_FRAME_10MS];
|
|
frame_f32.copy_from_slice(&msg.frame);
|
|
let out = VoiceActivityDetector::process_10ms(&mut vad, &frame_f32);
|
|
prob_arc.store(
|
|
out.probability.clamp(0.0, 1.0).to_bits(),
|
|
std::sync::atomic::Ordering::Relaxed,
|
|
);
|
|
seq_arc.store(msg.seq, std::sync::atomic::Ordering::Relaxed);
|
|
}
|
|
})
|
|
.ok()?;
|
|
|
|
Some(Self {
|
|
tx: Some(tx),
|
|
latest_probability,
|
|
latest_processed_seq,
|
|
alive,
|
|
handle: Some(handle),
|
|
})
|
|
}
|
|
|
|
/// Best-effort enqueue of a 10 ms frame for background inference.
|
|
pub fn try_send(&self, seq: u64, frame: &[f32; INPUT_FRAME_10MS]) -> bool {
|
|
let Some(tx) = &self.tx else {
|
|
return false;
|
|
};
|
|
tx.try_send(TenFrameMessage { seq, frame: *frame }).is_ok()
|
|
}
|
|
|
|
/// Latest probability published by the background worker.
|
|
pub fn latest_probability(&self) -> f32 {
|
|
f32::from_bits(
|
|
self.latest_probability
|
|
.load(std::sync::atomic::Ordering::Relaxed),
|
|
)
|
|
}
|
|
|
|
/// True when the worker is too far behind to trust its output.
|
|
pub fn is_stale(&self, capture_seq: u64) -> bool {
|
|
let latest = self
|
|
.latest_processed_seq
|
|
.load(std::sync::atomic::Ordering::Relaxed);
|
|
latest == u64::MAX || capture_seq.saturating_sub(latest) > TEN_MAX_STALE_FRAMES
|
|
}
|
|
}
|
|
|
|
impl Drop for TenOnnxVadWorker {
|
|
fn drop(&mut self) {
|
|
self.alive
|
|
.store(false, std::sync::atomic::Ordering::Relaxed);
|
|
drop(self.tx.take());
|
|
if let Some(h) = self.handle.take() {
|
|
let _ = h.join();
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn mel_filter_bank_has_expected_shape() {
|
|
let filters = build_mel_filters();
|
|
assert_eq!(filters.len(), MEL_BANDS);
|
|
assert!(filters.iter().all(|f| f.iter().any(|&v| v > 0.0)));
|
|
}
|
|
|
|
#[test]
|
|
fn preprocessing_produces_finite_features() {
|
|
let filters = build_mel_filters();
|
|
let mut planner = FftPlanner::<f32>::new();
|
|
let fft = planner.plan_fft_forward(FFT_SIZE);
|
|
let mut fft_buffer = vec![Complex32::ZERO; FFT_SIZE];
|
|
let frame = vec![0.0_f32; WINDOW_16K];
|
|
let feature = compute_feature(&filters, fft.as_ref(), &mut fft_buffer, &frame);
|
|
assert!(feature.iter().all(|v| v.is_finite()));
|
|
}
|
|
}
|