voice_engine/media/
denoiser.rs

1use super::codecs::resample::LinearResampler;
2use crate::media::{AudioFrame, PcmBuf, Sample, Samples, processor::Processor};
3use anyhow::Result;
4use nnnoiseless::DenoiseState;
5use std::cell::RefCell;
6
7pub struct NoiseReducer {
8    resampler_target: RefCell<LinearResampler>,
9    resampler_source: RefCell<LinearResampler>,
10    denoiser: RefCell<Box<DenoiseState<'static>>>,
11}
12
13impl NoiseReducer {
14    pub fn new(input_sample_rate: usize) -> Result<Self> {
15        let resampler48k = LinearResampler::new(48000, input_sample_rate)?;
16        let resampler16k = LinearResampler::new(input_sample_rate, 48000 as usize)?;
17        let denoiser = DenoiseState::new();
18        Ok(Self {
19            resampler_target: RefCell::new(resampler48k),
20            resampler_source: RefCell::new(resampler16k),
21            denoiser: RefCell::new(denoiser),
22        })
23    }
24}
25unsafe impl Send for NoiseReducer {}
26unsafe impl Sync for NoiseReducer {}
27
28impl Processor for NoiseReducer {
29    fn process_frame(&self, frame: &mut AudioFrame) -> Result<()> {
30        // If empty frame, nothing to do
31        if frame.samples.is_empty() {
32            return Ok(());
33        }
34
35        let samples = match &frame.samples {
36            Samples::PCM { samples } => samples,
37            _ => return Ok(()),
38        };
39        let samples = self.resampler_source.borrow_mut().resample(samples);
40        let input_size = samples.len();
41
42        let output_padding_size = input_size + DenoiseState::FRAME_SIZE;
43        let mut output_buf = vec![0.0; output_padding_size];
44        let input_f32: Vec<f32> = samples.iter().map(|&s| s.into()).collect();
45
46        let mut offset = 0;
47        let mut buf;
48
49        while offset < input_size {
50            let remaining_size = input_size - offset;
51            let chunk_len = remaining_size.min(DenoiseState::FRAME_SIZE);
52            let end_offset = offset + chunk_len;
53
54            let input_chunk = if chunk_len < DenoiseState::FRAME_SIZE {
55                buf = vec![0.0; DenoiseState::FRAME_SIZE];
56                buf[..chunk_len].copy_from_slice(&input_f32[offset..end_offset]);
57                &buf
58            } else {
59                &input_f32[offset..end_offset]
60            };
61
62            // Process the current frame
63            self.denoiser.borrow_mut().process_frame(
64                &mut output_buf[offset..offset + DenoiseState::FRAME_SIZE],
65                &input_chunk,
66            );
67
68            offset += chunk_len;
69        }
70
71        let samples = output_buf[..input_size]
72            .iter()
73            .map(|&s| s as Sample)
74            .collect::<PcmBuf>();
75
76        frame.samples = Samples::PCM {
77            samples: self.resampler_target.borrow_mut().resample(&samples),
78        };
79
80        Ok(())
81    }
82}
voice_engine/media/denoiser.rs

voice_engine/media/
denoiser.rs
