zk-audio 0.1.0

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

pub struct HighBandCleanupProcessor {
    profile: AudioProfile,
    detect_split_coeffs: Option<Coefficients<f32>>,
    split_coeffs: Option<Coefficients<f32>>,
    cleanup_coeffs: Option<Coefficients<f32>>,
    detect_split_low_pass: Option<DirectForm1<f32>>,
    split_low_pass: Option<DirectForm1<f32>>,
    cleanup_low_pass: Option<DirectForm1<f32>>,
    reduction: f32,
    noise_floor: f32,
    speech_memory: f32,
    last_rms: f32,
}

impl HighBandCleanupProcessor {
    pub fn new(profile: AudioProfile) -> Self {
        Self {
            profile,
            detect_split_coeffs: None,
            split_coeffs: None,
            cleanup_coeffs: None,
            detect_split_low_pass: None,
            split_low_pass: None,
            cleanup_low_pass: None,
            reduction: 0.0,
            noise_floor: 0.01,
            speech_memory: 0.0,
            last_rms: 0.0,
        }
    }
}

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

    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 high-band cleanup"))?;
        let split = Coefficients::<f32>::from_params(
            Type::LowPass,
            sample_rate,
            tuning.high_cleanup_split_hz.hz(),
            Q_BUTTERWORTH_F32,
        )
        .map_err(|_| AudioError::new("Invalid split configuration for high-band cleanup"))?;
        let cleanup = Coefficients::<f32>::from_params(
            Type::LowPass,
            sample_rate,
            (tuning.high_cleanup_split_hz * 0.78).max(1_800.0).hz(),
            Q_BUTTERWORTH_F32,
        )
        .map_err(|_| AudioError::new("Invalid cleanup configuration for high-band cleanup"))?;

        self.detect_split_coeffs = Some(split);
        self.split_coeffs = Some(split);
        self.cleanup_coeffs = Some(cleanup);
        self.detect_split_low_pass = Some(DirectForm1::<f32>::new(split));
        self.split_low_pass = Some(DirectForm1::<f32>::new(split));
        self.cleanup_low_pass = Some(DirectForm1::<f32>::new(cleanup));
        self.reduction = 0.0;
        self.noise_floor = 0.01;
        self.speech_memory = 0.0;
        self.last_rms = 0.0;
        Ok(())
    }

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

        let detect_split = self
            .detect_split_low_pass
            .as_mut()
            .ok_or_else(|| AudioError::new("High-band cleanup detector split not prepared"))?;
        let split = self
            .split_low_pass
            .as_mut()
            .ok_or_else(|| AudioError::new("High-band cleanup split not prepared"))?;
        let cleanup = self
            .cleanup_low_pass
            .as_mut()
            .ok_or_else(|| AudioError::new("High-band cleanup filter not prepared"))?;

        let mut low_energy = 0.0;
        let mut high_energy = 0.0;
        for sample in &frame.samples {
            let low = detect_split.run(*sample);
            let high = *sample - low;
            low_energy += low * low;
            high_energy += high * high;
        }
        low_energy = (low_energy / frame.samples.len() as f32).sqrt();
        high_energy = (high_energy / frame.samples.len() as f32).sqrt();

        let features = analyze_frame(&frame.samples, self.noise_floor);
        let frame_class = classify_frame(self.profile, features, self.noise_floor);
        let rms = features.rms;
        let peak = features.peak;
        if peak < self.noise_floor * 3.0 && rms < self.noise_floor * 2.0 {
            self.noise_floor = self.noise_floor * 0.94 + rms.max(0.0005) * 0.06;
        } else {
            self.noise_floor = self.noise_floor * 0.997 + rms.max(0.0005) * 0.003;
        }

        let tuning = profile_tuning(self.profile);
        let high_ratio = high_energy / (low_energy + 1e-4);
        let falling_energy = if self.last_rms > 1e-4 {
            ((self.last_rms - rms) / self.last_rms.max(rms)).clamp(0.0, 1.0)
        } else {
            0.0
        };
        let target = match frame_class {
            FrameClass::NoiseOnly => tuning.high_cleanup_strength * 0.95,
            FrameClass::Transitional => {
                let tail_focus =
                    (self.speech_memory * 0.65 + high_ratio * 0.22 + falling_energy * 0.30)
                        .clamp(0.0, 1.2);
                tuning.high_cleanup_strength * tail_focus
            }
            FrameClass::SpeechLike => {
                let harshness = ((high_ratio - 0.34) / 1.00).clamp(0.0, 1.0);
                tuning.high_cleanup_strength * harshness * 0.72
            }
        }
        .clamp(0.0, tuning.high_cleanup_strength.max(0.0));

        let smoothing = if target > self.reduction {
            tuning.high_cleanup_attack
        } else {
            tuning.high_cleanup_release
        };
        let start_reduction = self.reduction;
        self.reduction = start_reduction * (1.0 - smoothing) + target * 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 reduction = start_reduction + (self.reduction - start_reduction) * t;
            let low = split.run(*sample);
            let high = *sample - low;
            let smoothed_high = cleanup.run(high);
            let cleaned_high = high * (1.0 - reduction) + smoothed_high * reduction;
            *sample = low + cleaned_high;
        }

        self.speech_memory = match frame_class {
            FrameClass::SpeechLike => (self.speech_memory * 0.84 + 0.22).clamp(0.0, 1.0),
            FrameClass::Transitional => self.speech_memory * 0.88,
            FrameClass::NoiseOnly => self.speech_memory * 0.72,
        };
        self.last_rms = rms;
        Ok(())
    }

    fn reset(&mut self) {
        self.detect_split_low_pass = self.detect_split_coeffs.map(DirectForm1::<f32>::new);
        self.split_low_pass = self.split_coeffs.map(DirectForm1::<f32>::new);
        self.cleanup_low_pass = self.cleanup_coeffs.map(DirectForm1::<f32>::new);
        self.reduction = 0.0;
        self.noise_floor = 0.01;
        self.speech_memory = 0.0;
        self.last_rms = 0.0;
    }
}

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

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

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

        let mut frame = AudioFrame {
            samples: (0..512)
                .map(|i| if i % 2 == 0 { 0.06 } else { -0.06 })
                .collect(),
            spec,
        };
        let before = frame
            .samples
            .windows(2)
            .map(|pair| (pair[1] - pair[0]).abs())
            .sum::<f32>()
            / (frame.samples.len() - 1) as f32;
        for _ in 0..4 {
            processor.process(&mut frame).unwrap();
        }
        let after = frame
            .samples
            .windows(2)
            .map(|pair| (pair[1] - pair[0]).abs())
            .sum::<f32>()
            / (frame.samples.len() - 1) as f32;

        assert!(after < before);
    }
}