use crate::mel::{interleave_frames, log_mel_spectrogram, mel, norm_mel};
use crate::quant::quantize;
use crate::stft::Spectrogram;
use crate::vad::{duration_ms_for_n_frames, DetectionSettings, VoiceActivityDetector};
use js_sys::{Object, Reflect, Uint8Array, Uint8ClampedArray};
use ndarray::Array2;
use wasm_bindgen::prelude::*;
use web_sys::Worker;
#[wasm_bindgen]
pub struct SpeechToMel {
mel: Array2<f64>,
mel_vad: Array2<f64>,
fft: Spectrogram,
vad: VoiceActivityDetector,
hop_size: usize,
sampling_rate: f64,
accumulated_samples: Vec<f32>,
idx: usize,
}
#[wasm_bindgen]
impl SpeechToMel {
#[wasm_bindgen]
pub fn new(fft_size: usize, hop_size: usize, sampling_rate: f64, n_mels: usize) -> Self {
Self::new_with_settings(
fft_size,
hop_size,
sampling_rate,
n_mels,
DetectionSettings {
min_energy: 1.0,
min_y: 3,
min_x: 3,
min_mel: 0,
},
)
}
#[wasm_bindgen(js_name = newWithVadSettings)]
pub fn new_with_vad_settings(
fft_size: usize,
hop_size: usize,
sampling_rate: f64,
n_mels: usize,
min_energy: f64,
min_y: usize,
min_x: usize,
min_mel: usize,
) -> Self {
Self::new_with_settings(
fft_size,
hop_size,
sampling_rate,
n_mels,
DetectionSettings {
min_energy,
min_y,
min_x,
min_mel,
},
)
}
fn new_with_settings(
fft_size: usize,
hop_size: usize,
sampling_rate: f64,
n_mels: usize,
settings: DetectionSettings,
) -> Self {
let filters = mel(sampling_rate, fft_size, n_mels, None, None, false, true);
let filters2 = mel(sampling_rate, fft_size, n_mels / 4, None, None, false, true);
let stft = Spectrogram::new(fft_size, hop_size);
let vad = VoiceActivityDetector::new(&settings);
Self {
accumulated_samples: Vec::new(),
mel: filters,
mel_vad: filters2,
fft: stft,
vad,
sampling_rate,
hop_size,
idx: 0,
}
}
#[wasm_bindgen]
pub fn get(&mut self) -> JsValue {
let empty = vec![0.0; 0];
self.add(empty, false)
}
#[wasm_bindgen]
pub fn add(&mut self, data: Vec<f32>, vad: bool) -> JsValue {
let result = Object::new();
Reflect::set(&result, &JsValue::from_str("ok"), &JsValue::from(false)).unwrap();
self.accumulated_samples.extend_from_slice(&data);
if self.accumulated_samples.len() >= self.hop_size {
let (samples, rest) = self.accumulated_samples.split_at(self.hop_size);
Reflect::set(
&result,
&JsValue::from_str("len"),
&JsValue::from(samples.len()),
)
.unwrap();
if let Some(fft) = self.fft.add(&samples.to_vec()) {
let frame = &log_mel_spectrogram(&fft, &self.mel);
let frame2 = norm_mel(&log_mel_spectrogram(&fft, &self.mel_vad));
let (quant_frame, range) = quantize(&interleave_frames(&[frame.clone()], false, 0));
let frame_array = Uint8Array::from(&quant_frame[..]);
let frame_clamped_array = Uint8ClampedArray::new(&frame_array.buffer());
Reflect::set(&result, &JsValue::from_str("frame"), &frame_clamped_array).unwrap();
Reflect::set(&result, &JsValue::from_str("ok"), &JsValue::from(true)).unwrap();
Reflect::set(
&result,
&JsValue::from_str("min"),
&JsValue::from(range.min),
)
.unwrap();
Reflect::set(
&result,
&JsValue::from_str("max"),
&JsValue::from(range.max),
)
.unwrap();
let ms = duration_ms_for_n_frames(self.hop_size, self.sampling_rate, self.idx);
Reflect::set(&result, &JsValue::from_str("idx"), &JsValue::from(self.idx)).unwrap();
Reflect::set(&result, &JsValue::from_str("ms"), &JsValue::from(ms)).unwrap();
if vad {
if let Some(gap) = self.vad.add(&frame2) {
Reflect::set(&result, &JsValue::from_str("va"), &JsValue::from(gap))
.unwrap();
}
}
}
self.accumulated_samples = rest.to_vec();
self.idx = self.idx.wrapping_add(1);
}
JsValue::from(result)
}
}
#[wasm_bindgen]
pub fn startup(path: String) -> Worker {
let worker_handle = Worker::new(&path).unwrap();
worker_handle
}