Skip to main content

active_call/media/
agc.rs

1use crate::media::{AudioFrame, Sample, Samples, processor::Processor};
2use anyhow::Result;
3use serde::{Deserialize, Serialize};
4use serde_with::skip_serializing_none;
5use sonora_agc2::adaptive_digital_gain_controller::{AdaptiveDigitalGainController, FrameInfo};
6use sonora_agc2::common::{
7    ADJACENT_SPEECH_FRAMES_THRESHOLD, FRAME_DURATION_MS, MIN_LEVEL_DBFS,
8    SATURATION_PROTECTOR_INITIAL_HEADROOM_DB, float_s16_to_dbfs,
9};
10use sonora_agc2::limiter::Limiter;
11use sonora_agc2::noise_level_estimator::NoiseLevelEstimator;
12use sonora_agc2::saturation_protector::SaturationProtector;
13use sonora_agc2::speech_level_estimator::{AdaptiveDigitalConfig, SpeechLevelEstimator};
14
15// AGC2 processes audio in fixed 10ms sub-frames.
16const SUB_FRAME_MS: u32 = FRAME_DURATION_MS as u32;
17
18#[skip_serializing_none]
19#[derive(Clone, Debug, Deserialize, Serialize)]
20#[serde(rename_all = "camelCase")]
21#[serde(default)]
22pub struct AGCOption {
23    /// Target headroom below 0 dBFS in dB. AGC2 default: 5.0.
24    pub headroom_db: Option<f32>,
25    /// Maximum gain in dB. AGC2 default: 50.0.
26    pub max_gain_db: Option<f32>,
27    /// Initial gain applied before the speech level estimator becomes confident. AGC2 default: 15.0.
28    pub initial_gain_db: Option<f32>,
29    /// Max gain change in dB per second (attack/release rate). AGC2 default: 6.0.
30    pub max_gain_change_db_per_second: Option<f32>,
31    /// Noise floor cap above which AGC2 will not amplify. AGC2 default: -50.0.
32    pub max_output_noise_level_dbfs: Option<f32>,
33    /// Number of consecutive speech sub-frames required before allowing gain increase. AGC2 default: 12 (≈120 ms).
34    pub adjacent_speech_frames_threshold: Option<i32>,
35    /// Run the limiter after the adaptive gain. Default: true.
36    pub enable_limiter: Option<bool>,
37}
38
39impl Default for AGCOption {
40    fn default() -> Self {
41        Self {
42            headroom_db: None,
43            max_gain_db: None,
44            initial_gain_db: None,
45            max_gain_change_db_per_second: None,
46            max_output_noise_level_dbfs: None,
47            adjacent_speech_frames_threshold: None,
48            enable_limiter: None,
49        }
50    }
51}
52
53pub struct AutomaticGainControl {
54    sample_rate: u32,
55    sub_frame_samples: usize,
56    level_estimator: SpeechLevelEstimator,
57    saturation_protector: SaturationProtector,
58    noise_estimator: NoiseLevelEstimator,
59    controller: AdaptiveDigitalGainController,
60    limiter: Option<Limiter>,
61    f32_buf: Vec<f32>,
62    leftover: Vec<i16>,
63    last_speech_probability: f32,
64    last_applied_gain_linear: f32,
65}
66
67impl AutomaticGainControl {
68    pub fn new(sample_rate: u32, option: AGCOption) -> Result<Self> {
69        let config = AdaptiveDigitalConfig {
70            headroom_db: option.headroom_db.unwrap_or(5.0),
71            max_gain_db: option.max_gain_db.unwrap_or(50.0),
72            initial_gain_db: option.initial_gain_db.unwrap_or(15.0),
73            max_gain_change_db_per_second: option.max_gain_change_db_per_second.unwrap_or(6.0),
74            max_output_noise_level_dbfs: option.max_output_noise_level_dbfs.unwrap_or(-50.0),
75        };
76        let adjacent_threshold = option
77            .adjacent_speech_frames_threshold
78            .unwrap_or(ADJACENT_SPEECH_FRAMES_THRESHOLD);
79
80        let sub_frame_samples = (sample_rate as usize * SUB_FRAME_MS as usize) / 1000;
81        if sub_frame_samples == 0 || sub_frame_samples % 20 != 0 {
82            anyhow::bail!(
83                "sample rate {sample_rate} produces {sub_frame_samples} samples per 10ms, must be > 0 and divisible by 20"
84            );
85        }
86
87        let level_estimator = SpeechLevelEstimator::new(&config, adjacent_threshold);
88        let saturation_protector =
89            SaturationProtector::new(SATURATION_PROTECTOR_INITIAL_HEADROOM_DB, adjacent_threshold);
90        let noise_estimator = NoiseLevelEstimator::default();
91        let controller = AdaptiveDigitalGainController::new(config, adjacent_threshold);
92        let limiter = match option.enable_limiter.unwrap_or(true) {
93            true => Some(Limiter::new(sub_frame_samples)),
94            false => None,
95        };
96
97        Ok(Self {
98            sample_rate,
99            sub_frame_samples,
100            level_estimator,
101            saturation_protector,
102            noise_estimator,
103            controller,
104            limiter,
105            f32_buf: vec![0.0; sub_frame_samples],
106            leftover: Vec::with_capacity(sub_frame_samples),
107            last_speech_probability: 0.0,
108            last_applied_gain_linear: 1.0,
109        })
110    }
111
112    #[cfg(test)]
113    pub(crate) fn current_gain_for_test(&self) -> f32 {
114        self.last_applied_gain_linear
115    }
116
117    fn process_sub_frame(&mut self, sub_frame_i16: &mut [i16]) {
118        debug_assert_eq!(sub_frame_i16.len(), self.sub_frame_samples);
119
120        // Convert i16 -> S16-float (range [-32768, 32767]).
121        for (dst, src) in self.f32_buf.iter_mut().zip(sub_frame_i16.iter()) {
122            *dst = *src as f32;
123        }
124
125        let speech_probability = self.last_speech_probability;
126
127        let mut peak_abs = 0.0_f32;
128        let mut sum_sq = 0.0_f32;
129        for &s in self.f32_buf.iter() {
130            peak_abs = peak_abs.max(s.abs());
131            sum_sq += s * s;
132        }
133        let rms_lin = (sum_sq / self.f32_buf.len() as f32).sqrt();
134        // Clamp to satisfy SpeechLevelEstimator debug asserts on silent frames.
135        let rms_dbfs = float_s16_to_dbfs(rms_lin).clamp(MIN_LEVEL_DBFS, 30.0);
136        let peak_dbfs = float_s16_to_dbfs(peak_abs).clamp(MIN_LEVEL_DBFS, 30.0);
137
138        self.level_estimator.update(rms_dbfs, speech_probability);
139        self.saturation_protector.analyze(
140            speech_probability,
141            peak_dbfs,
142            self.level_estimator.level_dbfs(),
143        );
144
145        let mono_view: [&[f32]; 1] = [&self.f32_buf];
146        let noise_rms_dbfs = self.noise_estimator.analyze(&mono_view);
147
148        let limiter_envelope_dbfs = self
149            .limiter
150            .as_ref()
151            .map(|l| l.last_audio_level())
152            .unwrap_or(MIN_LEVEL_DBFS);
153
154        let info = FrameInfo {
155            speech_probability,
156            speech_level_dbfs: self.level_estimator.level_dbfs(),
157            speech_level_reliable: self.level_estimator.is_confident(),
158            noise_rms_dbfs,
159            headroom_db: self.saturation_protector.headroom_db(),
160            limiter_envelope_dbfs,
161        };
162
163        // Sample a non-zero input to estimate the applied linear gain after the
164        // adaptive controller. Stored so tests can read back the gain decision
165        // independently of the input level.
166        let probe_idx = self
167            .f32_buf
168            .iter()
169            .position(|&s| s.abs() > 1.0)
170            .unwrap_or(0);
171        let probe_in = self.f32_buf[probe_idx];
172
173        let mut channels: [&mut [f32]; 1] = [&mut self.f32_buf];
174        self.controller.process(&info, &mut channels);
175
176        if probe_in.abs() > 1.0 {
177            self.last_applied_gain_linear = self.f32_buf[probe_idx] / probe_in;
178        }
179
180        if let Some(limiter) = self.limiter.as_mut() {
181            let mut channels: [&mut [f32]; 1] = [&mut self.f32_buf];
182            limiter.process(&mut channels);
183        }
184
185        // Convert back to i16 with saturating cast.
186        for (dst, src) in sub_frame_i16.iter_mut().zip(self.f32_buf.iter()) {
187            *dst = src.clamp(i16::MIN as f32, i16::MAX as f32) as Sample;
188        }
189    }
190}
191
192impl Processor for AutomaticGainControl {
193    fn process_frame(&mut self, frame: &mut AudioFrame) -> Result<()> {
194        let samples = match &mut frame.samples {
195            Samples::PCM { samples } if !samples.is_empty() => samples,
196            _ => return Ok(()),
197        };
198
199        if frame.sample_rate != self.sample_rate {
200            // Frame from an unexpected rate; pass-through.
201            return Ok(());
202        }
203
204        if let Some(p) = frame.speech_probability {
205            self.last_speech_probability = p.clamp(0.0, 1.0);
206        }
207
208        // Accumulate input into sub-frame-sized chunks. The leftover buffer
209        // carries samples across frames so a 20ms input becomes two 10ms
210        // AGC2 sub-frames cleanly.
211        let mut work = std::mem::take(samples);
212        let mut produced: Vec<i16> = Vec::with_capacity(work.len() + self.leftover.len());
213
214        if !self.leftover.is_empty() {
215            let mut combined = std::mem::take(&mut self.leftover);
216            combined.append(&mut work);
217            work = combined;
218        }
219
220        let mut idx = 0;
221        while idx + self.sub_frame_samples <= work.len() {
222            let end = idx + self.sub_frame_samples;
223            self.process_sub_frame(&mut work[idx..end]);
224            produced.extend_from_slice(&work[idx..end]);
225            idx = end;
226        }
227
228        if idx < work.len() {
229            self.leftover.extend_from_slice(&work[idx..]);
230        }
231
232        *samples = produced;
233        Ok(())
234    }
235}