Merge pull request #95 from teamclouday/speexdsp

Speexdsp improvements

Author: teamclouday

RustApp/src/audio/denoise.rs RustApp/src/audio/denoise_rnnoise.rsRustApp/src/audio/denoise.rs renamed to RustApp/src/audio/denoise_rnnoise.rs

@@ -2,6 +2,8 @@ use std::sync::{LazyLock, Mutex};
 
 use nnnoiseless::DenoiseState;
 
+pub const DENOISE_RNNOISE_SAMPLE_RATE: u32 = 48000;
+
 struct DenoiseCache {
     sample_buffer: Vec<Vec<f32>>,
     denoisers: Vec<Box<DenoiseState<'static>>>,

RustApp/src/audio/denoise_speex.rs

@@ -1,54 +1,155 @@
-use speexdsp::preprocess::*;
+use std::sync::{LazyLock, Mutex};
+
+use speexdsp::preprocess::{SpeexPreprocess, SpeexPreprocessConst};
+
+use crate::audio::AudioProcessParams;
 
 // xxx: do we really need to change the sample rate ?
 // apparently, speexdsp is optimized for low sample rate (8000, 16000), according to chatgpt,
 // but 16000 just doesn't work on my end
 pub const DENOISE_SPEEX_SAMPLE_RATE: u32 = 48000;
+const FRAME_SIZE: usize = (DENOISE_SPEEX_SAMPLE_RATE as f32 * 0.02) as usize; // 20 ms frame
 
-pub struct DenoiseSpeexCache {
+struct DenoiseSpeexCache {
+    sample_buffer: Vec<Vec<i16>>,
     denoisers: Vec<SpeexPreprocess>,
+    config_noise_suppress: i32,
+    config_vad_enabled: bool,
+    config_vad_threshold: u32,
+    config_agc_enabled: bool,
+    config_agc_target: u32,
+    config_dereverb_enabled: bool,
+    config_dereverb_level: f32,
+}
+
+impl DenoiseSpeexCache {
+    fn is_config_changed(&self, config: &AudioProcessParams) -> bool {
+        self.config_noise_suppress != config.speex_noise_suppress
+            || self.config_vad_enabled != config.speex_vad_enabled
+            || self.config_vad_threshold != config.speex_vad_threshold
+            || self.config_agc_enabled != config.speex_agc_enabled
+            || self.config_agc_target != config.speex_agc_target
+            || self.config_dereverb_enabled != config.speex_dereverb_enabled
+            || self.config_dereverb_level != config.speex_dereverb_level
+    }
 }
 
 // safe because packets are processed in order, and not concurrently
 unsafe impl Send for DenoiseSpeexCache {}
 
+// safe because packets are processed in order, and not concurrently
+static DENOISE_CACHE: LazyLock<Mutex<Option<DenoiseSpeexCache>>> =
+    LazyLock::new(|| Mutex::new(None));
+
 pub fn denoise_speex_f32_stream(
-    data: &mut [Vec<i16>],
-    cache: &mut Option<DenoiseSpeexCache>,
-    noise_suppress: i32,
-) -> anyhow::Result<()> {
-    const FRAME_SIZE: usize = (DENOISE_SPEEX_SAMPLE_RATE as f32 * 0.02) as usize; // 20 ms frame
-
-    if cache.is_none() {
-        *cache = Some(DenoiseSpeexCache {
+    data: &[Vec<f32>],
+    config: &AudioProcessParams,
+) -> anyhow::Result<Vec<Vec<f32>>> {
+    let mut denoise_cache = DENOISE_CACHE.lock().unwrap();
+
+    if denoise_cache.is_none()
+        || data.len() != denoise_cache.as_ref().unwrap().denoisers.len()
+        || denoise_cache.as_ref().unwrap().is_config_changed(config)
+    {
+        *denoise_cache = Some(DenoiseSpeexCache {
+            sample_buffer: vec![Vec::with_capacity(FRAME_SIZE); data.len()],
             denoisers: data
                 .iter()
                 .map(|_| {
                     let mut st =
                         SpeexPreprocess::new(FRAME_SIZE, DENOISE_SPEEX_SAMPLE_RATE as usize)
                             .unwrap();
-                    st.set_denoise(true);
-                    st.set_noise_suppress(noise_suppress);
+                    st.preprocess_ctl(SpeexPreprocessConst::SPEEX_PREPROCESS_SET_DENOISE, 1)
+                        .unwrap();
+                    st.set_noise_suppress(config.speex_noise_suppress);
+                    st.preprocess_ctl(
+                        SpeexPreprocessConst::SPEEX_PREPROCESS_SET_VAD,
+                        if config.speex_vad_enabled { 1 } else { 0 },
+                    )
+                    .unwrap();
+                    st.preprocess_ctl(
+                        SpeexPreprocessConst::SPEEX_PREPROCESS_SET_PROB_START,
+                        config.speex_vad_threshold,
+                    )
+                    .unwrap();
+                    st.preprocess_ctl(
+                        SpeexPreprocessConst::SPEEX_PREPROCESS_SET_AGC,
+                        if config.speex_agc_enabled { 1 } else { 0 },
+                    )
+                    .unwrap();
+                    st.preprocess_ctl(
+                        SpeexPreprocessConst::SPEEX_PREPROCESS_SET_AGC_TARGET,
+                        config.speex_agc_target,
+                    )
+                    .unwrap();
+                    st.preprocess_ctl(
+                        SpeexPreprocessConst::SPEEX_PREPROCESS_SET_DEREVERB,
+                        if config.speex_dereverb_enabled { 1 } else { 0 },
+                    )
+                    .unwrap();
+                    st.preprocess_ctl(
+                        SpeexPreprocessConst::SPEEX_PREPROCESS_SET_DEREVERB_LEVEL,
+                        config.speex_dereverb_level,
+                    )
+                    .unwrap();
                     st
                 })
                 .collect(),
+            config_noise_suppress: config.speex_noise_suppress,
+            config_vad_enabled: config.speex_vad_enabled,
+            config_vad_threshold: config.speex_vad_threshold,
+            config_agc_enabled: config.speex_agc_enabled,
+            config_agc_target: config.speex_agc_target,
+            config_dereverb_enabled: config.speex_dereverb_enabled,
+            config_dereverb_level: config.speex_dereverb_level,
         });
     }
 
-    for (channel, st) in data
-        .iter_mut()
-        .zip(cache.as_mut().unwrap().denoisers.iter_mut())
-    {
-        for frame in channel.chunks_exact_mut(FRAME_SIZE) {
-            match st.preprocess_run(frame) {
+    let cache = denoise_cache.as_mut().unwrap();
+    let mut output: Vec<Vec<f32>> = vec![Vec::new(); data.len()];
+
+    // Convert f32 to i16
+    let data_i16: Vec<Vec<i16>> = data
+        .iter()
+        .map(|channel| {
+            channel
+                .iter()
+                .map(|&x| (x * i16::MAX as f32).clamp(i16::MIN as f32, i16::MAX as f32) as i16)
+                .collect()
+        })
+        .collect();
+
+    // Append new data into the cache
+    for channel_idx in 0..data_i16.len() {
+        cache.sample_buffer[channel_idx].extend_from_slice(&data_i16[channel_idx]);
+    }
+
+    while cache.sample_buffer[0].len() >= FRAME_SIZE {
+        for channel_idx in 0..data.len() {
+            match cache.denoisers[channel_idx]
+                .preprocess_run(&mut cache.sample_buffer[channel_idx][0..FRAME_SIZE])
+            {
                 0 => {
-                    frame.fill(0);
+                    cache.sample_buffer[channel_idx][0..FRAME_SIZE].fill(0);
                 }
                 1 => {}
                 _ => panic!(),
             }
+
+            // Scale back to -1.0 to 1.0 range
+            output[channel_idx].extend_from_slice(
+                &cache.sample_buffer[channel_idx][0..FRAME_SIZE]
+                    .iter()
+                    .map(|&x| x as f32 / i16::MAX as f32)
+                    .collect::<Vec<f32>>(),
+            );
+        }
+
+        // Clear the sample buffer for the next round
+        for channel in &mut cache.sample_buffer {
+            channel.drain(0..FRAME_SIZE);
         }
     }
 
-    Ok(())
+    Ok(output)
 }

RustApp/src/audio/mod.rs

@@ -3,16 +3,47 @@ use byteorder::{ByteOrder, NativeEndian, WriteBytesExt};
 use rtrb::Consumer;
 
 use crate::{
-    config::{AudioFormat, ChannelCount, SampleRate},
+    config::{AudioFormat, ChannelCount, Config, DenoiseKind, SampleRate},
     ui::app::{AppState, Stream},
 };
 
-mod denoise;
+mod denoise_rnnoise;
 mod denoise_speex;
 mod player;
 pub mod process;
 mod resampler;
-pub use denoise_speex::DenoiseSpeexCache;
+
+/// Audio processing parameters
+#[derive(Clone, Debug)]
+pub struct AudioProcessParams {
+    pub target_format: AudioPacketFormat,
+    pub denoise: Option<DenoiseKind>,
+    pub amplify: Option<f32>,
+    pub speex_noise_suppress: i32,
+    pub speex_vad_enabled: bool,
+    pub speex_vad_threshold: u32,
+    pub speex_agc_enabled: bool,
+    pub speex_agc_target: u32,
+    pub speex_dereverb_enabled: bool,
+    pub speex_dereverb_level: f32,
+}
+
+impl AudioProcessParams {
+    pub fn new(target_format: AudioPacketFormat, config: Config) -> Self {
+        Self {
+            target_format,
+            denoise: config.denoise.then_some(config.denoise_kind),
+            amplify: config.amplify.then_some(config.amplify_value),
+            speex_noise_suppress: config.speex_noise_suppress,
+            speex_vad_enabled: config.speex_vad_enabled,
+            speex_vad_threshold: config.speex_vad_threshold,
+            speex_agc_enabled: config.speex_agc_enabled,
+            speex_agc_target: config.speex_agc_target,
+            speex_dereverb_enabled: config.speex_dereverb_enabled,
+            speex_dereverb_level: config.speex_dereverb_level,
+        }
+    }
+}
 
 impl AppState {
     pub fn start_audio_stream(

RustApp/src/audio/player.rs

@@ -39,7 +39,9 @@ pub fn start_audio_stream(
     }
 
     if !supported {
-        bail!("unsupported output audio format or sample rate.");
+        bail!(
+            "Unsupported output audio format or sample rate. Please apply recommended format from settings page."
+        );
     }
 
     let config = cpal::StreamConfig {

RustApp/src/audio/process.rs

@@ -1,34 +1,15 @@
 use std::borrow::Cow;
 
 use crate::{
-    audio::denoise_speex::{DENOISE_SPEEX_SAMPLE_RATE, denoise_speex_f32_stream},
-    config::{AudioFormat, Config, DenoiseKind},
+    audio::{
+        denoise_rnnoise::DENOISE_RNNOISE_SAMPLE_RATE,
+        denoise_speex::{DENOISE_SPEEX_SAMPLE_RATE, denoise_speex_f32_stream},
+    },
+    config::{AudioFormat, DenoiseKind},
     streamer::{AudioPacketMessage, AudioStream},
 };
 
-use super::{
-    AudioBytes, AudioPacketFormat, denoise::denoise_f32_stream, resampler::resample_f32_stream,
-};
-
-/// Audio processing parameters
-#[derive(Clone, Debug)]
-pub struct AudioProcessParams {
-    pub target_format: AudioPacketFormat,
-    pub denoise: Option<DenoiseKind>,
-    pub amplify: Option<f32>,
-    pub speex_noise_suppress: i32,
-}
-
-impl AudioProcessParams {
-    pub fn new(target_format: AudioPacketFormat, config: Config) -> Self {
-        Self {
-            target_format,
-            denoise: config.denoise.then_some(config.denoise_kind),
-            amplify: config.amplify.then_some(config.amplify_value),
-            speex_noise_suppress: config.speex_noise_suppress,
-        }
-    }
-}
+use super::{AudioBytes, denoise_rnnoise::denoise_f32_stream, resampler::resample_f32_stream};
 
 impl AudioStream {
     /// This function converts an audio stream from packet into producer
@@ -65,14 +46,15 @@ impl AudioStream {
         if let Some(denoise) = &config.denoise {
             match denoise {
                 DenoiseKind::Rnnoise => {
-                    const DENOISE_SAMPLE_RATE: u32 = 48000;
-
-                    let prepared_buffer = if current_sample_rate == DENOISE_SAMPLE_RATE {
+                    let prepared_buffer = if current_sample_rate == DENOISE_RNNOISE_SAMPLE_RATE {
                         Cow::Borrowed(&buffer)
                     } else {
-                        let tmp =
-                            resample_f32_stream(&buffer, current_sample_rate, DENOISE_SAMPLE_RATE)?;
-                        current_sample_rate = DENOISE_SAMPLE_RATE;
+                        let tmp = resample_f32_stream(
+                            &buffer,
+                            current_sample_rate,
+                            DENOISE_RNNOISE_SAMPLE_RATE,
+                        )?;
+                        current_sample_rate = DENOISE_RNNOISE_SAMPLE_RATE;
                         Cow::Owned(tmp)
                     };
 
@@ -92,21 +74,7 @@ impl AudioStream {
                         Cow::Owned(tmp)
                     };
 
-                    let mut prepared_buffer: Vec<Vec<i16>> = prepared_buffer
-                        .iter()
-                        .map(|v| v.iter().map(|v| AudioBytes::from_f32(*v)).collect())
-                        .collect();
-
-                    denoise_speex_f32_stream(
-                        &mut prepared_buffer,
-                        &mut self.denoise_speex_cache,
-                        config.speex_noise_suppress,
-                    )?;
-
-                    buffer = prepared_buffer
-                        .into_iter()
-                        .map(|v| v.into_iter().map(|v| AudioBytes::to_f32(&v)).collect())
-                        .collect();
+                    buffer = denoise_speex_f32_stream(&prepared_buffer, &config)?;
                 }
             }
         }

RustApp/src/config.rs

@@ -23,6 +23,12 @@ pub struct Config {
     pub denoise: bool,
     pub denoise_kind: DenoiseKind,
     pub speex_noise_suppress: i32,
+    pub speex_vad_enabled: bool,
+    pub speex_vad_threshold: u32,
+    pub speex_agc_enabled: bool,
+    pub speex_agc_target: u32,
+    pub speex_dereverb_enabled: bool,
+    pub speex_dereverb_level: f32,
     pub theme: AppTheme,
     pub amplify: bool,
     pub amplify_value: f32,
@@ -62,30 +68,26 @@ impl Default for Config {
             theme: Default::default(),
             amplify: false,
             amplify_value: 2.0,
-            speex_noise_suppress: -30,
+            speex_noise_suppress: -30, // range: [-100, 0]
+            speex_vad_enabled: false,
+            speex_vad_threshold: 80, // range: [0, 100]
+            speex_agc_enabled: false,
+            speex_agc_target: 8000, // range: [8000, 65535]
+            speex_dereverb_enabled: false,
+            speex_dereverb_level: 0.5, // range: [0.0, 1.0]
         }
     }
 }
 
-pub struct ConfigCache {
-    pub amplify_value: String,
-    pub speex_noise_suppress: String,
-}
-
-impl ConfigCache {
-    pub fn new(config: &Config) -> Self {
-        Self {
-            amplify_value: config.amplify_value.to_string(),
-            speex_noise_suppress: config.speex_noise_suppress.to_string(),
-        }
-    }
-
-    pub fn parse_amplify_value(&self) -> Option<f32> {
-        self.amplify_value.replace(',', ".").parse().ok()
-    }
-
-    pub fn parse_speex_noise_suppress(&self) -> Option<i32> {
-        self.speex_noise_suppress.parse().ok()
+impl Config {
+    pub fn reset_denoise_settings(&mut self) {
+        self.speex_noise_suppress = -30;
+        self.speex_vad_enabled = false;
+        self.speex_vad_threshold = 80;
+        self.speex_agc_enabled = false;
+        self.speex_agc_target = 8000;
+        self.speex_dereverb_enabled = false;
+        self.speex_dereverb_level = 0.5;
     }
 }