use mel_spec::quant::{load_tga_8bit, to_array2};
use mel_spec::vad::{
vad_boundaries, DetectionSettings, VadFrameTiming, VoiceActivityDetector,
VoiceActivityTimestamps,
};
use ndarray::{concatenate, s, Array, Array2, Axis};
use std::collections::HashSet;
struct LegacyVoiceActivityDetector {
mel_buffer: Vec<Array2<f64>>,
settings: DetectionSettings,
idx: usize,
}
impl LegacyVoiceActivityDetector {
fn new(settings: &DetectionSettings) -> Self {
Self {
mel_buffer: Vec::new(),
settings: *settings,
idx: 0,
}
}
fn add(&mut self, frame: &Array2<f64>) -> Option<bool> {
let min_x = self.settings.min_x;
if self.idx == 128 {
self.mel_buffer = self.mel_buffer[(self.mel_buffer.len() - min_x)..].to_vec();
self.idx = min_x;
}
self.mel_buffer.push(frame.to_owned());
self.idx += 1;
if self.idx < min_x {
return None;
}
let window = &self.mel_buffer[self.idx - min_x..];
let edge_info = legacy_vad_boundaries(window, &self.settings);
let intersected = edge_info.intersected();
if intersected.is_empty() {
Some(false)
} else {
Some(intersected[0] == 0)
}
}
}
#[derive(Debug)]
struct LegacyEdgeInfo {
non_intersected_columns: Vec<usize>,
intersected_columns: Vec<usize>,
}
impl LegacyEdgeInfo {
fn new(non_intersected_columns: Vec<usize>, intersected_columns: Vec<usize>) -> Self {
Self {
non_intersected_columns,
intersected_columns,
}
}
fn non_intersected(&self) -> Vec<usize> {
self.non_intersected_columns.clone()
}
fn intersected(&self) -> Vec<usize> {
self.intersected_columns.clone()
}
}
fn legacy_vad_boundaries(frames: &[Array2<f64>], settings: &DetectionSettings) -> LegacyEdgeInfo {
let array_views: Vec<_> = frames.iter().map(|a| a.view()).collect();
let min_energy = settings.min_energy;
let min_y = settings.min_y;
let min_mel = settings.min_mel;
let merged_frames = concatenate(Axis(1), &array_views).unwrap();
let shape = merged_frames.raw_dim();
let width = shape[1];
let height = shape[0];
let sobel_x =
Array::from_shape_vec((3, 3), vec![-1.0, 0.0, 1.0, -2.0, 0.0, 2.0, -1.0, 0.0, 1.0])
.unwrap();
let sobel_y =
Array::from_shape_vec((3, 3), vec![-1.0, -2.0, -1.0, 0.0, 0.0, 0.0, 1.0, 2.0, 1.0])
.unwrap();
let gradient_mag = Array::from_shape_fn((height - 2, width - 2), |(y, x)| {
let view = merged_frames.slice(s![y..y + 3, x..x + 3]);
let mut gradient_x = 0.0;
let mut gradient_y = 0.0;
for j in 0..3 {
for i in 0..3 {
gradient_x += view[[j, i]] * sobel_x[[j, i]];
gradient_y += view[[j, i]] * sobel_y[[j, i]];
}
}
(gradient_x * gradient_x + gradient_y * gradient_y).sqrt()
});
let mut raw_classification = Vec::with_capacity(width - 2);
for x in 0..(width - 2) {
let mut count = 0;
for y in 0..(height - 2) {
let grad = gradient_mag[(y, x)];
if y >= min_mel && grad >= min_energy {
count += 1;
}
}
raw_classification.push(count >= min_y);
}
let smoothed_classification = legacy_smooth_mask(&raw_classification, 4);
let mut intersected_columns = Vec::new();
let mut non_intersected_columns = Vec::new();
for (x, &active) in smoothed_classification.iter().enumerate() {
if active {
intersected_columns.push(x);
} else {
non_intersected_columns.push(x);
}
}
LegacyEdgeInfo::new(non_intersected_columns, intersected_columns)
}
fn legacy_smooth_mask(mask: &[bool], window: usize) -> Vec<bool> {
let n = mask.len();
let mut smoothed = vec![false; n];
for i in 0..n {
let start = if i < window { 0 } else { i - window };
let end = if i + window + 1 > n {
n
} else {
i + window + 1
};
let count_true = mask[start..end].iter().filter(|&&val| val).count();
if count_true * 2 >= (end - start) {
smoothed[i] = true;
}
}
smoothed
}
#[test]
fn vad_boundaries_matches_legacy_on_reference_fixtures() {
let n_mels = 80;
let cases = [
(
DetectionSettings {
min_energy: 1.0,
min_y: 3,
min_x: 6,
min_mel: 0,
},
"./testdata/quantized_mel_golden.tga",
),
(
DetectionSettings {
min_energy: 1.0,
min_y: 10,
min_x: 10,
min_mel: 0,
},
"./testdata/blank/frame_23760.tga",
),
(
DetectionSettings {
min_energy: 1.0,
min_y: 10,
min_x: 10,
min_mel: 0,
},
"./testdata/speech/frame_27125.tga",
),
(
DetectionSettings {
min_energy: 1.0,
min_y: 6,
min_x: 1,
min_mel: 0,
},
"./testdata/jfk_full_speech_chunk0_golden.tga",
),
];
for (settings, path) in cases {
let dequantized_mel = load_tga_8bit(path).unwrap();
let frames = to_array2(&dequantized_mel, n_mels);
let current = vad_boundaries(&[frames.clone()], &settings);
let legacy = legacy_vad_boundaries(&[frames], &settings);
assert_eq!(current.intersected(), legacy.intersected(), "{path}");
assert_eq!(
current.non_intersected(),
legacy.non_intersected(),
"{path}"
);
assert_eq!(current.gradient_positions(), HashSet::new(), "{path}");
}
}
#[test]
fn streaming_vad_matches_legacy_on_quantized_fixture() {
let n_mels = 80;
let settings = DetectionSettings {
min_energy: 1.0,
min_y: 3,
min_x: 3,
min_mel: 0,
};
let dequantized_mel = load_tga_8bit("./testdata/quantized_mel_golden.tga").unwrap();
let frames = to_array2(&dequantized_mel, n_mels);
let chunks: Vec<Array2<f64>> = frames
.axis_chunks_iter(Axis(1), 1)
.map(|chunk| chunk.to_owned())
.collect();
let mut current = VoiceActivityDetector::new(&settings);
let mut legacy = LegacyVoiceActivityDetector::new(&settings);
let current_outputs: Vec<Option<bool>> =
chunks.iter().map(|chunk| current.add(chunk)).collect();
let legacy_outputs: Vec<Option<bool>> = chunks.iter().map(|chunk| legacy.add(chunk)).collect();
assert_eq!(current_outputs, legacy_outputs);
}
#[test]
fn streaming_vad_can_return_stft_timestamps() {
let n_mels = 80;
let settings = DetectionSettings {
min_energy: 1.0,
min_y: 3,
min_x: 3,
min_mel: 0,
};
let dequantized_mel = load_tga_8bit("./testdata/quantized_mel_golden.tga").unwrap();
let frames = to_array2(&dequantized_mel, n_mels);
let chunks: Vec<Array2<f64>> = frames
.axis_chunks_iter(Axis(1), 1)
.map(|chunk| chunk.to_owned())
.collect();
let timing = VadFrameTiming::new(400, 160, 16_000.0);
let mut vad = VoiceActivityDetector::new_with_timing(&settings, timing);
let first_activity = chunks
.iter()
.filter_map(|chunk| vad.add_activity(chunk))
.next()
.unwrap();
assert_eq!(first_activity.frame_index, 2);
assert_eq!(
first_activity.timestamps,
Some(VoiceActivityTimestamps {
start_ms: 20,
center_ms: 33,
end_ms: 45,
})
);
}