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
15const 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 pub headroom_db: Option<f32>,
25 pub max_gain_db: Option<f32>,
27 pub initial_gain_db: Option<f32>,
29 pub max_gain_change_db_per_second: Option<f32>,
31 pub max_output_noise_level_dbfs: Option<f32>,
33 pub adjacent_speech_frames_threshold: Option<i32>,
35 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 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 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 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 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 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 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}