zk-audio 0.1.0

Audio processing library for voice recording and enhancement
Documentation
use super::support::{lerp, smooth_towards};
use super::{analyze_frame, classify_frame, FrameClass};
use crate::contracts::AudioProcessor;
use crate::core::{AudioError, AudioFrame, AudioProfile, AudioResult, AudioSpec};
use crate::profiles::profile_tuning;
use biquad::{Biquad, Coefficients, DirectForm1, Hertz, ToHertz, Type, Q_BUTTERWORTH_F32};

pub struct DeEsserProcessor {
    profile: AudioProfile,
    detect_high_pass: Option<DirectForm1<f32>>,
    detect_low_pass: Option<DirectForm1<f32>>,
    apply_split_low_pass: Option<DirectForm1<f32>>,
    band_noise_floor: f32,
    current_strength: f32,
}

impl DeEsserProcessor {
    pub fn new(profile: AudioProfile) -> Self {
        Self {
            profile,
            detect_high_pass: None,
            detect_low_pass: None,
            apply_split_low_pass: None,
            band_noise_floor: 0.0025,
            current_strength: 0.0,
        }
    }

    fn filtered_band(
        input: &[f32],
        high_pass: &mut DirectForm1<f32>,
        low_pass: &mut DirectForm1<f32>,
    ) -> Vec<f32> {
        let mut filtered = Vec::with_capacity(input.len());
        for sample in input {
            filtered.push(low_pass.run(high_pass.run(*sample)));
        }
        filtered
    }
}

impl AudioProcessor for DeEsserProcessor {
    fn name(&self) -> &'static str {
        "deesser"
    }

    fn prepare(&mut self, spec: AudioSpec) -> AudioResult<()> {
        let tuning = profile_tuning(self.profile);
        let sample_rate = Hertz::<f32>::from_hz(spec.sample_rate as f32)
            .map_err(|_| AudioError::new("Invalid sample rate for deesser"))?;
        let detect_high = Coefficients::<f32>::from_params(
            Type::HighPass,
            sample_rate,
            tuning.deesser_band_low_hz.hz(),
            Q_BUTTERWORTH_F32,
        )
        .map_err(|_| AudioError::new("Invalid deesser high-pass configuration"))?;
        let detect_low = Coefficients::<f32>::from_params(
            Type::LowPass,
            sample_rate,
            tuning.deesser_band_high_hz.hz(),
            Q_BUTTERWORTH_F32,
        )
        .map_err(|_| AudioError::new("Invalid deesser low-pass configuration"))?;

        self.detect_high_pass = Some(DirectForm1::<f32>::new(detect_high));
        self.detect_low_pass = Some(DirectForm1::<f32>::new(detect_low));
        let apply_split = Coefficients::<f32>::from_params(
            Type::LowPass,
            sample_rate,
            tuning.deesser_band_low_hz.hz(),
            Q_BUTTERWORTH_F32,
        )
        .map_err(|_| AudioError::new("Invalid deesser apply split configuration"))?;

        self.apply_split_low_pass = Some(DirectForm1::<f32>::new(apply_split));
        self.band_noise_floor = 0.0025;
        self.current_strength = 0.0;
        Ok(())
    }

    fn process(&mut self, frame: &mut AudioFrame) -> AudioResult<()> {
        if frame.samples.is_empty() {
            return Ok(());
        }

        let detect_high = self
            .detect_high_pass
            .as_mut()
            .ok_or_else(|| AudioError::new("Deesser detector high-pass not prepared"))?;
        let detect_low = self
            .detect_low_pass
            .as_mut()
            .ok_or_else(|| AudioError::new("Deesser detector low-pass not prepared"))?;
        let apply_split = self
            .apply_split_low_pass
            .as_mut()
            .ok_or_else(|| AudioError::new("Deesser apply split low-pass not prepared"))?;

        let full_features = analyze_frame(&frame.samples, self.band_noise_floor.max(0.0005));
        let frame_class = classify_frame(
            self.profile,
            full_features,
            self.band_noise_floor.max(0.0005),
        );
        let detected_band = Self::filtered_band(&frame.samples, detect_high, detect_low);
        let band_features = analyze_frame(&detected_band, self.band_noise_floor.max(0.0005));

        if frame_class == FrameClass::NoiseOnly {
            self.band_noise_floor =
                self.band_noise_floor * 0.92 + band_features.rms.max(0.0002) * 0.08;
        } else {
            self.band_noise_floor =
                self.band_noise_floor * 0.995 + band_features.rms.max(0.0002) * 0.005;
        }

        let tuning = profile_tuning(self.profile);
        let presence_ratio = band_features.rms / full_features.rms.max(0.0005);
        let band_speechiness = band_features.peak / (self.band_noise_floor + 1e-4);
        let ratio_activation = ((presence_ratio - tuning.deesser_presence_threshold)
            / tuning.deesser_presence_threshold.max(0.05))
        .clamp(0.0, 1.0);
        let speech_activation = ((band_speechiness - 1.6) / 2.8).clamp(0.0, 1.0);
        let base_strength = match frame_class {
            FrameClass::NoiseOnly => tuning.deesser_strength * 0.18,
            FrameClass::Transitional => tuning.deesser_strength * 0.55,
            FrameClass::SpeechLike => tuning.deesser_strength,
        };
        let target_strength =
            (base_strength * ratio_activation.max(speech_activation)).clamp(0.0, 0.95);
        let start_strength = self.current_strength;
        let smoothing = if target_strength > self.current_strength {
            tuning.deesser_attack
        } else {
            tuning.deesser_release
        };
        self.current_strength = smooth_towards(self.current_strength, target_strength, smoothing);

        let len = frame.samples.len().max(1) as f32;
        for (index, sample) in frame.samples.iter_mut().enumerate() {
            let t = index as f32 / len;
            let strength = lerp(start_strength, self.current_strength, t);
            let low = apply_split.run(*sample);
            let high = *sample - low;
            *sample = low + high * (1.0 - strength);
        }

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn test_spec() -> AudioSpec {
        AudioSpec {
            sample_rate: 44_100,
            channels: 1,
        }
    }

    #[test]
    fn deesser_reduces_presence_band_energy() {
        let spec = test_spec();
        let mut processor = DeEsserProcessor::new(AudioProfile::VoiceHvac);
        processor.prepare(spec).unwrap();

        let source = (0..512)
            .map(|i| {
                let t = i as f32 / spec.sample_rate as f32;
                let body = (std::f32::consts::TAU * 220.0 * t).sin() * 0.11;
                let sibilance = (std::f32::consts::TAU * 4_200.0 * t).sin() * 0.08;
                body + sibilance
            })
            .collect::<Vec<_>>();
        let mut processed = source.clone();
        for _ in 0..4 {
            let mut frame = AudioFrame {
                samples: source.clone(),
                spec,
            };
            processor.process(&mut frame).unwrap();
            processed = frame.samples;
        }
        let mean_delta = processed
            .iter()
            .zip(source.iter())
            .map(|(after, before)| (after - before).abs())
            .sum::<f32>()
            / source.len().max(1) as f32;

        assert!(processor.current_strength > 0.01);
        assert!(mean_delta > 0.001);
    }

    #[test]
    fn deesser_is_light_on_low_frequency_voice_body() {
        let spec = test_spec();
        let mut processor = DeEsserProcessor::new(AudioProfile::VoiceHvac);
        processor.prepare(spec).unwrap();

        let mut frame = AudioFrame {
            samples: (0..512)
                .map(|i| {
                    let t = i as f32 / spec.sample_rate as f32;
                    (std::f32::consts::TAU * 180.0 * t).sin() * 0.12
                })
                .collect(),
            spec,
        };

        let before = frame.samples.iter().map(|s| s.abs()).sum::<f32>() / 512.0;
        processor.process(&mut frame).unwrap();
        let after = frame.samples.iter().map(|s| s.abs()).sum::<f32>() / 512.0;

        assert!(after > before * 0.90);
    }
}