use crate::pcm::{ms_to_samples, samples_to_ms};
const FRAME_MS: u64 = 20;
const LEAD_KEEP_FRAMES: usize = 10;
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct Config {
pub min_ms: u64,
pub max_ms: u64,
pub pause_ms: u64,
pub min_speech_ms: u64,
pub over_floor: f32,
pub min_rms: f32,
}
impl Default for Config {
fn default() -> Self {
Self {
min_ms: 6_000,
max_ms: 28_000,
pause_ms: 700,
min_speech_ms: 400,
over_floor: 3.0,
min_rms: 150.0,
}
}
}
#[derive(Debug, Clone, PartialEq)]
pub struct Chunk {
pub start_ms: u64,
pub samples: Vec<i16>,
pub speech_ms: u64,
}
impl Chunk {
pub fn duration_ms(&self) -> u64 {
samples_to_ms(self.samples.len())
}
}
pub struct Chunker {
cfg: Config,
buf_start_ms: u64,
buf: Vec<i16>,
frames: Vec<bool>,
partial: usize,
noise_floor: f32,
trailing_silence_ms: u64,
}
impl Chunker {
pub fn new(cfg: Config) -> Self {
Self {
cfg,
buf_start_ms: 0,
buf: Vec::new(),
frames: Vec::new(),
partial: 0,
noise_floor: 300.0,
trailing_silence_ms: 0,
}
}
pub fn position_ms(&self) -> u64 {
self.buf_start_ms + samples_to_ms(self.buf.len())
}
pub fn push(&mut self, end_ms: u64, samples: &[i16]) -> Vec<Chunk> {
let start_ms = end_ms.saturating_sub(samples_to_ms(samples.len()));
let pos = self.position_ms();
let mut out = Vec::new();
if start_ms > pos + 200 {
let gap = start_ms - pos;
if gap > self.cfg.pause_ms {
out.extend(self.cut_all());
self.buf_start_ms = start_ms;
self.trailing_silence_ms = 0;
} else {
out.extend(self.append(&vec![0; ms_to_samples(gap)]));
}
}
out.extend(self.append(samples));
out
}
pub fn flush(&mut self) -> Vec<Chunk> {
self.cut_all()
}
fn append(&mut self, samples: &[i16]) -> Vec<Chunk> {
let mut out = Vec::new();
let frame = ms_to_samples(FRAME_MS);
self.buf.extend_from_slice(samples);
self.partial += samples.len();
while self.partial >= frame {
let end = self.buf.len() - (self.partial - frame);
let is_speech = self.classify_at(end - frame, end);
self.frames.push(is_speech);
self.partial -= frame;
self.trailing_silence_ms = if is_speech {
0
} else {
self.trailing_silence_ms + FRAME_MS
};
if let Some(chunk) = self.maybe_cut() {
out.push(chunk);
}
}
out
}
fn classify_at(&mut self, from: usize, to: usize) -> bool {
let frame = &self.buf[from..to];
let rms = (frame.iter().map(|&s| (s as f32) * (s as f32)).sum::<f32>() / frame.len() as f32).sqrt();
let speech = rms > self.noise_floor * self.cfg.over_floor && rms > self.cfg.min_rms;
if rms < self.noise_floor {
self.noise_floor = self.noise_floor * 0.9 + rms * 0.1;
} else if !speech {
self.noise_floor = self.noise_floor * 0.995 + rms * 0.005;
}
self.noise_floor = self.noise_floor.max(20.0);
speech
}
fn maybe_cut(&mut self) -> Option<Chunk> {
let dur = self.frames.len() as u64 * FRAME_MS;
if dur >= self.cfg.min_ms && self.trailing_silence_ms >= self.cfg.pause_ms {
let keep_frames = (self.trailing_silence_ms / FRAME_MS / 2) as usize;
return self.cut_at(self.frames.len() - keep_frames);
}
if dur >= self.cfg.max_ms {
let look = (5_000 / FRAME_MS) as usize;
let from = self.frames.len().saturating_sub(look);
let at = (from..self.frames.len())
.rev()
.find(|&i| !self.frames[i])
.unwrap_or(self.frames.len());
return self.cut_at(at.max(1));
}
None
}
fn cut_at(&mut self, frame_index: usize) -> Option<Chunk> {
let frame = ms_to_samples(FRAME_MS);
let n = (frame_index * frame).min(self.buf.len());
let samples: Vec<i16> = self.buf.drain(..n).collect();
let flags: Vec<bool> = self.frames.drain(..frame_index.min(self.frames.len())).collect();
let start_ms = self.buf_start_ms;
self.buf_start_ms += samples_to_ms(samples.len());
self.trailing_silence_ms = self.frames.iter().rev().take_while(|s| !**s).count() as u64 * FRAME_MS;
let speech_ms = flags.iter().filter(|s| **s).count() as u64 * FRAME_MS;
let lead = flags
.iter()
.position(|s| *s)
.unwrap_or(flags.len())
.saturating_sub(LEAD_KEEP_FRAMES);
let skip = (lead * frame).min(samples.len());
(speech_ms >= self.cfg.min_speech_ms && samples.len() > skip).then(|| Chunk {
start_ms: start_ms + samples_to_ms(skip),
samples: samples[skip..].to_vec(),
speech_ms,
})
}
fn cut_all(&mut self) -> Vec<Chunk> {
let frame = ms_to_samples(FRAME_MS);
let frames = self.frames.len() + usize::from(self.partial > 0);
if self.partial > 0 {
self.frames.push(false);
self.buf.resize(frames * frame, 0);
self.partial = 0;
}
let chunk = self.cut_at(frames);
self.trailing_silence_ms = 0;
chunk.into_iter().collect()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::pcm::SAMPLE_RATE;
fn tone(ms: u64) -> Vec<i16> {
(0..ms_to_samples(ms))
.map(|i| {
((i as f32 / SAMPLE_RATE as f32 * 2.0 * std::f32::consts::PI * 300.0).sin() * 8000.0) as i16
})
.collect()
}
fn hiss(ms: u64) -> Vec<i16> {
(0..ms_to_samples(ms))
.map(|i| if i % 2 == 0 { 40 } else { -40 })
.collect()
}
fn feed(c: &mut Chunker, mut at_ms: u64, audio: &[i16]) -> (u64, Vec<Chunk>) {
let mut out = Vec::new();
for block in audio.chunks(ms_to_samples(10)) {
at_ms += samples_to_ms(block.len());
out.extend(c.push(at_ms, block));
}
(at_ms, out)
}
#[test]
fn cuts_at_the_first_pause_after_min_length() {
let mut c = Chunker::new(Config::default());
let mut audio = hiss(500);
audio.extend(tone(7_000));
audio.extend(hiss(1_000));
audio.extend(tone(3_000));
let (_, chunks) = feed(&mut c, 0, &audio);
assert_eq!(chunks.len(), 1, "режется в паузе после 7,5 с");
let first = &chunks[0];
assert_eq!(
first.start_ms, 300,
"тишина в начале срезана, 200 мс перед речью оставлены"
);
assert!(
(7_200..=7_900).contains(&first.duration_ms()),
"{}",
first.duration_ms()
);
assert!(first.speech_ms >= 6_800, "{}", first.speech_ms);
let rest = c.flush();
assert_eq!(rest.len(), 1);
assert_eq!(
rest[0].start_ms, 8_300,
"второй кусок — с начала речи минус 200 мс"
);
assert!(rest[0].speech_ms >= 2_800);
}
#[test]
fn short_pauses_inside_min_length_do_not_cut() {
let mut c = Chunker::new(Config::default());
let mut audio = Vec::new();
for _ in 0..2 {
audio.extend(tone(1_500));
audio.extend(hiss(900));
}
audio.extend(tone(1_500));
let (_, chunks) = feed(&mut c, 0, &audio);
assert!(chunks.is_empty(), "меньше min_ms — не режем");
assert_eq!(c.flush().len(), 1);
}
#[test]
fn long_speech_without_pause_is_cut_at_max() {
let mut c = Chunker::new(Config::default());
let (_, chunks) = feed(&mut c, 0, &tone(65_000));
assert_eq!(chunks.len(), 2);
for ch in &chunks {
assert!(ch.duration_ms() <= 28_000, "{}", ch.duration_ms());
}
assert_eq!(chunks[1].start_ms, chunks[0].duration_ms());
}
#[test]
fn silence_is_never_sent() {
let mut c = Chunker::new(Config::default());
let (_, chunks) = feed(&mut c, 0, &hiss(60_000));
assert!(chunks.is_empty());
assert!(c.flush().is_empty());
}
#[test]
fn steady_noise_becomes_the_floor_not_speech() {
let noise: Vec<i16> = (0..ms_to_samples(40_000))
.map(|i| if i % 2 == 0 { 600 } else { -600 })
.collect();
let mut c = Chunker::new(Config::default());
let (_, chunks) = feed(&mut c, 0, &noise);
let speech: u64 = chunks.iter().chain(c.flush().iter()).map(|c| c.speech_ms).sum();
assert!(speech < 1_000, "ровный шум распознан как речь: {speech} мс");
}
#[test]
fn gaps_in_the_stream_keep_time_right() {
let mut c = Chunker::new(Config::default());
let (end, _) = feed(&mut c, 0, &tone(3_000));
let (_, chunks) = feed(&mut c, end + 10_000, &tone(2_000));
assert_eq!(chunks.len(), 1, "длинная дыра закрывает кусок");
assert_eq!(chunks[0].start_ms, 0);
let rest = c.flush();
assert_eq!(rest[0].start_ms, 13_000, "время после дыры — настоящее");
}
#[test]
fn small_gaps_are_filled_with_silence() {
let mut c = Chunker::new(Config::default());
let (end, _) = feed(&mut c, 0, &tone(2_000));
feed(&mut c, end + 300, &tone(1_000));
assert_eq!(c.position_ms(), 3_300);
}
}