re-teamspeak/src/iced-app/src/noise_cancel.rs

use serde::{Deserialize, Serialize};
use sonora::config::{
    AdaptiveDigital, GainController2, HighPassFilter, NoiseSuppression,
    NoiseSuppressionLevel,
};
use sonora::{AudioProcessing, Config, StreamConfig};

const SONORA_FRAME_SIZE: usize = 480;

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum NoiseCancelMethod {
    None,
    Nnnoiseless,
    Sonora,
}

impl NoiseCancelMethod {
    pub fn label(&self) -> &str {
        match self {
            Self::None => "Off",
            Self::Nnnoiseless => "RNNoise (nnnoiseless)",
            Self::Sonora => "WebRTC (Sonora)",
        }
    }

    pub fn all() -> &'static [NoiseCancelMethod] {
        &[Self::None, Self::Nnnoiseless, Self::Sonora]
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub struct SonoraSettings {
    pub noise_suppression_level: NoiseSuppressionLevelSetting,
    pub agc_enabled: bool,
    pub high_pass_enabled: bool,
}

impl Default for SonoraSettings {
    fn default() -> Self {
        Self {
            noise_suppression_level: NoiseSuppressionLevelSetting::Moderate,
            agc_enabled: true,
            high_pass_enabled: true,
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum NoiseSuppressionLevelSetting {
    Low,
    Moderate,
    High,
    VeryHigh,
}

impl NoiseSuppressionLevelSetting {
    pub fn label(self) -> &'static str {
        match self {
            Self::Low => "Low",
            Self::Moderate => "Moderate",
            Self::High => "High",
            Self::VeryHigh => "Very High",
        }
    }

    pub fn all() -> &'static [NoiseSuppressionLevelSetting] {
        &[Self::Low, Self::Moderate, Self::High, Self::VeryHigh]
    }

    fn to_sonora(self) -> NoiseSuppressionLevel {
        match self {
            Self::Low => NoiseSuppressionLevel::Low,
            Self::Moderate => NoiseSuppressionLevel::Moderate,
            Self::High => NoiseSuppressionLevel::High,
            Self::VeryHigh => NoiseSuppressionLevel::VeryHigh,
        }
    }
}

pub struct NoiseReducer {
    method: NoiseCancelMethod,
    sonora_settings: SonoraSettings,
    nnnoiseless: Option<Box<nnnoiseless::DenoiseState<'static>>>,
    sonora: Option<SonoraNoiseReducer>,
    residual_buf: Vec<f32>,
}

struct SonoraNoiseReducer {
    processor: AudioProcessing,
    input: Vec<f32>,
    output: Vec<f32>,
}

impl SonoraNoiseReducer {
    fn new(settings: SonoraSettings) -> Self {
        let stream = StreamConfig::new(48_000, 1);
        let config = Config {
            high_pass_filter: settings.high_pass_enabled.then(HighPassFilter::default),
            noise_suppression: Some(NoiseSuppression {
                level: settings.noise_suppression_level.to_sonora(),
                ..NoiseSuppression::default()
            }),
            gain_controller2: settings.agc_enabled.then(|| GainController2 {
                adaptive_digital: Some(AdaptiveDigital::default()),
                ..GainController2::default()
            }),
            ..Config::default()
        };

        Self {
            processor: AudioProcessing::builder()
                .config(config)
                .capture_config(stream)
                .render_config(stream)
                .build(),
            input: vec![0.0; SONORA_FRAME_SIZE],
            output: vec![0.0; SONORA_FRAME_SIZE],
        }
    }

    fn process(&mut self, frame: &mut [f32]) {
        self.input.copy_from_slice(frame);
        let src = [&self.input[..]];
        let mut dest = [&mut self.output[..]];
        if let Err(error) = self.processor.process_capture_f32(&src, &mut dest) {
            tracing::debug!("Sonora noise suppression error: {error}");
            return;
        }
        frame.copy_from_slice(&self.output);
    }
}

impl NoiseReducer {
    pub fn new(method: NoiseCancelMethod, sonora_settings: SonoraSettings) -> Self {
        let mut this = Self {
            method,
            sonora_settings,
            nnnoiseless: None,
            sonora: None,
            residual_buf: Vec::new(),
        };
        this.init_method();
        this
    }

    fn init_method(&mut self) {
        self.nnnoiseless = match self.method {
            NoiseCancelMethod::Nnnoiseless => Some(nnnoiseless::DenoiseState::new()),
            _ => None,
        };
        self.sonora = match self.method {
            NoiseCancelMethod::Sonora => Some(SonoraNoiseReducer::new(self.sonora_settings)),
            _ => None,
        };
        self.residual_buf.clear();
    }

    pub fn set_method(&mut self, method: NoiseCancelMethod) {
        self.method = method;
        self.init_method();
    }

    pub fn set_sonora_settings(&mut self, settings: SonoraSettings) {
        self.sonora_settings = settings;
        if self.method == NoiseCancelMethod::Sonora {
            self.init_method();
        }
    }

    pub fn process(&mut self, _samples: &mut [f32]) {
        match self.method {
            NoiseCancelMethod::None => {}
            NoiseCancelMethod::Nnnoiseless => self.process_nnnoiseless(_samples),
            NoiseCancelMethod::Sonora => self.process_sonora(_samples),
        }
    }

    fn process_nnnoiseless(&mut self, samples: &mut [f32]) {
        let denoise = match &mut self.nnnoiseless {
            Some(d) => d,
            None => return,
        };

        self.residual_buf.extend_from_slice(samples);
        let frame_size = nnnoiseless::DenoiseState::FRAME_SIZE; // 480
        let mut output = vec![0.0f32; frame_size];
        let mut write_pos = 0;

        while self.residual_buf.len() >= frame_size {
            let frame: Vec<f32> = self.residual_buf.drain(..frame_size).collect();
            let mut input = [0.0f32; 480];
            for (i, &s) in frame.iter().enumerate().take(frame_size) {
                input[i] = s * 32768.0;
            }
            denoise.process_frame(&mut output, &input);
            for i in 0..frame_size {
                if write_pos + i < samples.len() {
                    samples[write_pos + i] = output[i] / 32768.0;
                }
            }
            write_pos += frame_size;
        }
    }

    fn process_sonora(&mut self, samples: &mut [f32]) {
        let reducer = match &mut self.sonora {
            Some(s) => s,
            None => return,
        };

        // Sonora processes 10ms frames (480 samples at 48kHz)
        let frame_size = SONORA_FRAME_SIZE;
        let mut offset = 0;

        while offset + frame_size <= samples.len() {
            let frame = &mut samples[offset..offset + frame_size];
            reducer.process(frame);
            offset += frame_size;
        }
    }
}