279 lines
8.7 KiB
Rust
279 lines
8.7 KiB
Rust
use std::{
|
|
io::Error,
|
|
sync::{mpsc, Arc, Mutex},
|
|
time::Duration,
|
|
};
|
|
|
|
use cpal::{
|
|
traits::{DeviceTrait, HostTrait, StreamTrait},
|
|
Device, Sample, SizedSample,
|
|
};
|
|
|
|
use crate::audio_toolkit::{
|
|
audio::FrameResampler,
|
|
constants,
|
|
vad::{self, VadFrame},
|
|
VoiceActivityDetector,
|
|
};
|
|
|
|
enum Cmd {
|
|
Start,
|
|
Stop(mpsc::Sender<Vec<f32>>),
|
|
Shutdown,
|
|
}
|
|
|
|
pub struct AudioRecorder {
|
|
device: Option<Device>,
|
|
cmd_tx: Option<mpsc::Sender<Cmd>>,
|
|
worker_handle: Option<std::thread::JoinHandle<()>>,
|
|
vad: Option<Arc<Mutex<Box<dyn vad::VoiceActivityDetector>>>>,
|
|
}
|
|
|
|
impl AudioRecorder {
|
|
pub fn new() -> Result<Self, Box<dyn std::error::Error>> {
|
|
Ok(AudioRecorder {
|
|
device: None,
|
|
cmd_tx: None,
|
|
worker_handle: None,
|
|
vad: None,
|
|
})
|
|
}
|
|
|
|
pub fn with_vad(mut self, vad: Box<dyn VoiceActivityDetector>) -> Self {
|
|
self.vad = Some(Arc::new(Mutex::new(vad)));
|
|
self
|
|
}
|
|
|
|
pub fn open(&mut self, device: Option<Device>) -> Result<(), Box<dyn std::error::Error>> {
|
|
if self.worker_handle.is_some() {
|
|
return Ok(()); // already open
|
|
}
|
|
|
|
let (sample_tx, sample_rx) = mpsc::channel::<Vec<f32>>();
|
|
let (cmd_tx, cmd_rx) = mpsc::channel::<Cmd>();
|
|
|
|
let host = cpal::default_host();
|
|
let device = match device {
|
|
Some(dev) => dev,
|
|
None => host
|
|
.default_input_device()
|
|
.ok_or_else(|| Error::new(std::io::ErrorKind::NotFound, "No input device found"))?,
|
|
};
|
|
|
|
let thread_device = device.clone();
|
|
let vad = self.vad.clone();
|
|
|
|
let worker = std::thread::spawn(move || {
|
|
let config = AudioRecorder::get_preferred_config(&thread_device)
|
|
.expect("failed to fetch preferred config");
|
|
|
|
let sample_rate = config.sample_rate().0;
|
|
let channels = config.channels() as usize;
|
|
|
|
println!(
|
|
"Using device: {:?}\nSample rate: {}\nChannels: {}\nFormat: {:?}",
|
|
thread_device.name(),
|
|
sample_rate,
|
|
channels,
|
|
config.sample_format()
|
|
);
|
|
|
|
let stream = match config.sample_format() {
|
|
cpal::SampleFormat::I8 => {
|
|
AudioRecorder::build_stream::<i8>(&thread_device, &config, sample_tx, channels)
|
|
.unwrap()
|
|
}
|
|
cpal::SampleFormat::I16 => {
|
|
AudioRecorder::build_stream::<i16>(&thread_device, &config, sample_tx, channels)
|
|
.unwrap()
|
|
}
|
|
cpal::SampleFormat::I32 => {
|
|
AudioRecorder::build_stream::<i32>(&thread_device, &config, sample_tx, channels)
|
|
.unwrap()
|
|
}
|
|
cpal::SampleFormat::F32 => {
|
|
AudioRecorder::build_stream::<f32>(&thread_device, &config, sample_tx, channels)
|
|
.unwrap()
|
|
}
|
|
_ => panic!("unsupported sample format"),
|
|
};
|
|
|
|
stream.play().expect("failed to start stream");
|
|
|
|
// keep the stream alive while we process samples
|
|
run_consumer(sample_rate, vad, sample_rx, cmd_rx);
|
|
// stream is dropped here, after run_consumer returns
|
|
});
|
|
|
|
self.device = Some(device);
|
|
self.cmd_tx = Some(cmd_tx);
|
|
self.worker_handle = Some(worker);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
pub fn start(&self) -> Result<(), Box<dyn std::error::Error>> {
|
|
if let Some(tx) = &self.cmd_tx {
|
|
tx.send(Cmd::Start)?;
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
pub fn stop(&self) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
|
|
let (resp_tx, resp_rx) = mpsc::channel();
|
|
if let Some(tx) = &self.cmd_tx {
|
|
tx.send(Cmd::Stop(resp_tx))?;
|
|
}
|
|
Ok(resp_rx.recv()?) // wait for the samples
|
|
}
|
|
|
|
pub fn close(&mut self) -> Result<(), Box<dyn std::error::Error>> {
|
|
if let Some(tx) = self.cmd_tx.take() {
|
|
let _ = tx.send(Cmd::Shutdown);
|
|
}
|
|
if let Some(h) = self.worker_handle.take() {
|
|
let _ = h.join();
|
|
}
|
|
self.device = None;
|
|
Ok(())
|
|
}
|
|
|
|
fn build_stream<T>(
|
|
device: &cpal::Device,
|
|
config: &cpal::SupportedStreamConfig,
|
|
sample_tx: mpsc::Sender<Vec<f32>>,
|
|
channels: usize,
|
|
) -> Result<cpal::Stream, cpal::BuildStreamError>
|
|
where
|
|
T: Sample + SizedSample + Send + 'static,
|
|
f32: cpal::FromSample<T>,
|
|
{
|
|
let mut output_buffer = Vec::new();
|
|
|
|
let stream_cb = move |data: &[T], _: &cpal::InputCallbackInfo| {
|
|
output_buffer.clear();
|
|
|
|
if channels == 1 {
|
|
// Direct conversion without intermediate Vec
|
|
output_buffer.extend(data.iter().map(|&sample| sample.to_sample::<f32>()));
|
|
} else {
|
|
// Convert to mono directly
|
|
let frame_count = data.len() / channels;
|
|
output_buffer.reserve(frame_count);
|
|
|
|
for frame in data.chunks_exact(channels) {
|
|
let mono_sample = frame
|
|
.iter()
|
|
.map(|&sample| sample.to_sample::<f32>())
|
|
.sum::<f32>()
|
|
/ channels as f32;
|
|
output_buffer.push(mono_sample);
|
|
}
|
|
}
|
|
|
|
if sample_tx.send(output_buffer.clone()).is_err() {
|
|
eprintln!("Failed to send samples");
|
|
}
|
|
};
|
|
|
|
device.build_input_stream(
|
|
&config.clone().into(),
|
|
stream_cb,
|
|
|err| eprintln!("Stream error: {}", err),
|
|
None,
|
|
)
|
|
}
|
|
|
|
fn get_preferred_config(
|
|
device: &cpal::Device,
|
|
) -> Result<cpal::SupportedStreamConfig, Box<dyn std::error::Error>> {
|
|
let supported_configs = device.supported_input_configs()?;
|
|
|
|
// Try to find a config that supports 16kHz
|
|
for config_range in supported_configs {
|
|
if config_range.min_sample_rate().0 <= constants::WHISPER_SAMPLE_RATE
|
|
&& config_range.max_sample_rate().0 >= constants::WHISPER_SAMPLE_RATE
|
|
{
|
|
// Found a config that supports 16kHz, use it
|
|
return Ok(
|
|
config_range.with_sample_rate(cpal::SampleRate(constants::WHISPER_SAMPLE_RATE))
|
|
);
|
|
}
|
|
}
|
|
|
|
// If no config supports 16kHz, fall back to default
|
|
Ok(device.default_input_config()?)
|
|
}
|
|
}
|
|
|
|
fn run_consumer(
|
|
in_sample_rate: u32,
|
|
vad: Option<Arc<Mutex<Box<dyn vad::VoiceActivityDetector>>>>,
|
|
sample_rx: mpsc::Receiver<Vec<f32>>,
|
|
cmd_rx: mpsc::Receiver<Cmd>,
|
|
) {
|
|
let mut frame_resampler = FrameResampler::new(
|
|
in_sample_rate as usize,
|
|
constants::WHISPER_SAMPLE_RATE as usize,
|
|
Duration::from_millis(30),
|
|
);
|
|
|
|
let mut processed_samples = Vec::<f32>::new();
|
|
let mut recording = false;
|
|
|
|
fn handle_frame(
|
|
samples: &[f32],
|
|
recording: bool,
|
|
vad: &Option<Arc<Mutex<Box<dyn vad::VoiceActivityDetector>>>>,
|
|
out_buf: &mut Vec<f32>,
|
|
) {
|
|
if !recording {
|
|
return;
|
|
}
|
|
|
|
if let Some(vad_arc) = vad {
|
|
let mut det = vad_arc.lock().unwrap();
|
|
match det.push_frame(samples).unwrap_or(VadFrame::Speech(samples)) {
|
|
VadFrame::Speech(buf) => out_buf.extend_from_slice(buf),
|
|
VadFrame::Noise => {}
|
|
}
|
|
} else {
|
|
out_buf.extend_from_slice(samples);
|
|
}
|
|
}
|
|
|
|
loop {
|
|
let raw = match sample_rx.recv() {
|
|
Ok(s) => s,
|
|
Err(_) => break, // stream closed
|
|
};
|
|
|
|
frame_resampler.push(&raw, &mut |frame: &[f32]| {
|
|
handle_frame(frame, recording, &vad, &mut processed_samples)
|
|
});
|
|
|
|
// non-blocking check for a command
|
|
while let Ok(cmd) = cmd_rx.try_recv() {
|
|
match cmd {
|
|
Cmd::Start => {
|
|
processed_samples.clear();
|
|
recording = true;
|
|
if let Some(v) = &vad {
|
|
v.lock().unwrap().reset();
|
|
}
|
|
}
|
|
Cmd::Stop(reply_tx) => {
|
|
recording = false;
|
|
|
|
frame_resampler.finish(&mut |frame: &[f32]| {
|
|
// we still want to process the last few frames
|
|
handle_frame(frame, true, &vad, &mut processed_samples)
|
|
});
|
|
|
|
let _ = reply_tx.send(std::mem::take(&mut processed_samples));
|
|
}
|
|
Cmd::Shutdown => return,
|
|
}
|
|
}
|
|
}
|
|
}
|