1mod analysis;
2mod aq;
3mod lf;
4mod q;
5mod threads;
6
7#[cfg(test)]
8mod tests;
9
10pub use analysis::{Analyser, FrameAnalysis};
11pub use aq::{SEGMENTS, Segmentation, tree_probs};
12pub use lf::{MAX_LEVEL, guess_filter_level, pick_filter_level};
13pub use q::{crf_to_q, lambda, q_index_for, qstep};
14pub use threads::{
15 LAG, Row, Sbs, TileLayout, Wavefront, max_log2_tile_cols, min_log2_tile_cols, tile_offset,
16};
17
18#[derive(Clone, Debug, PartialEq)]
19pub enum RcMode {
20 ConstantQ {
21 q_index: u8,
22 },
23 Crf {
24 crf: u8,
25 },
26 Bitrate {
27 kbps: u32,
28 buffer_ms: u32,
29 },
30 TwoPass {
31 kbps: u32,
32 stats: Vec<FrameAnalysis>,
33 },
34}
35
36#[derive(Clone, Copy, Debug, PartialEq, Eq)]
37pub enum Aq {
38 Off,
39 Variance,
40}
41
42#[derive(Clone, Debug, PartialEq)]
43pub struct RcConfig {
44 pub mode: RcMode,
45 pub width: usize,
46 pub height: usize,
47 pub fps_num: u32,
48 pub fps_den: u32,
49 pub keyint: u32,
50 pub min_q: u8,
51 pub max_q: u8,
52 pub aq: Aq,
53}
54
55impl RcConfig {
56 pub fn new(mode: RcMode, width: usize, height: usize) -> RcConfig {
57 RcConfig {
58 mode,
59 width,
60 height,
61 fps_num: 30,
62 fps_den: 1,
63 keyint: 240,
64 min_q: 0,
65 max_q: 255,
66 aq: Aq::Off,
67 }
68 }
69}
70
71#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
72pub enum FrameKind {
73 Key,
74 Inter,
75 Golden,
76 AltRef,
77}
78
79impl FrameKind {
80 fn slot(self) -> usize {
81 match self {
82 FrameKind::Key => 0,
83 FrameKind::Inter => 1,
84 FrameKind::Golden | FrameKind::AltRef => 2,
85 }
86 }
87}
88
89#[derive(Clone, Copy, Debug, PartialEq, Eq)]
90pub struct FramePlan {
91 pub base_q_idx: u8,
92 pub filter_level: u8,
93 pub sharpness: u8,
94 pub lambda: u32,
95 pub target_bits: u64,
96}
97
98const QCOMP: f64 = 0.6;
99const EXP: [f64; 3] = [0.85, 1.5, 1.5];
100const K0: [f64; 3] = [7.2, 42.0, 42.0];
101
102fn solve(terms: [(f64, f64); 2], bits: f64) -> f64 {
103 let f = |s: f64| terms.iter().map(|&(a, e)| a * s.powf(-e)).sum::<f64>();
104 let (mut lo, mut hi) = (0.5f64.ln(), 8000f64.ln());
105 for _ in 0..40 {
106 let mid = 0.5 * (lo + hi);
107 if f(mid.exp()) > bits {
108 lo = mid;
109 } else {
110 hi = mid;
111 }
112 }
113 (0.5 * (lo + hi)).exp()
114}
115
116pub struct RateControl {
117 cfg: RcConfig,
118 analyser: Analyser,
119 seg: Segmentation,
120 k: [f64; 3],
121 seen: [bool; 3],
122 frame_bits: f64,
123 buffer: f64,
124 buffer_size: f64,
125 log_cost: f64,
126 log_cost_n: f64,
127 ratio: f64,
128 a_inter: f64,
129 a_key: f64,
130 remaining: f64,
131 frame: usize,
132 suffix: Vec<(f64, f64)>,
133 geo: f64,
134 pending: Option<(FrameKind, u64)>,
135}
136
137impl RateControl {
138 pub fn new(cfg: RcConfig) -> RateControl {
139 let fps = cfg.fps_num.max(1) as f64 / cfg.fps_den.max(1) as f64;
140 let geo = match &cfg.mode {
141 RcMode::TwoPass { stats, .. } => two_pass_geo(stats, cfg.keyint),
142 _ => 1.0,
143 };
144 let (frame_bits, buffer_size, remaining, suffix) = match &cfg.mode {
145 RcMode::Bitrate { kbps, buffer_ms } => {
146 let b = *kbps as f64 * 1000.0;
147 (
148 b / fps,
149 b * (*buffer_ms).max(100) as f64 / 1000.0,
150 0.0,
151 Vec::new(),
152 )
153 }
154 RcMode::TwoPass { kbps, stats } => {
155 let b = *kbps as f64 * 1000.0;
156 let total = b * stats.len() as f64 / fps;
157 (
158 b / fps,
159 b * 2.0,
160 total,
161 two_pass_suffix(stats, cfg.keyint, geo),
162 )
163 }
164 _ => (0.0, 0.0, 0.0, Vec::new()),
165 };
166 RateControl {
167 analyser: Analyser::new(cfg.width, cfg.height),
168 seg: Segmentation::default(),
169 k: K0,
170 seen: [false; 3],
171 frame_bits,
172 buffer: 0.0,
173 buffer_size,
174 log_cost: 0.0,
175 log_cost_n: 0.0,
176 ratio: 0.5,
177 a_inter: 0.0,
178 a_key: 0.0,
179 remaining,
180 frame: 0,
181 geo,
182 suffix,
183 pending: None,
184 cfg,
185 }
186 }
187
188 pub fn analyse(&mut self, y: &[u8], stride: usize) -> FrameAnalysis {
189 self.analyser.run(y, stride)
190 }
191
192 pub fn segmentation(&self) -> &Segmentation {
193 &self.seg
194 }
195
196 pub fn plan(&mut self, kind: FrameKind, a: &FrameAnalysis) -> FramePlan {
197 let cost = a.cost(kind) as f64;
198 if a.has_prev && kind != FrameKind::Key {
199 self.ratio = 0.8 * self.ratio + 0.2 * a.inter_ratio().min(1.0);
200 self.log_cost += cost.ln();
201 self.log_cost_n += 1.0;
202 }
203 let g = if self.log_cost_n > 0.0 {
204 (self.log_cost / self.log_cost_n).exp()
205 } else {
206 cost
207 };
208 let boost = self.boost(kind);
209 let modulation = ((cost / g).powf(1.0 - QCOMP)).clamp(0.8, 1.25);
210 let step_for = |s: f64| match kind {
211 FrameKind::Key | FrameKind::Golden | FrameKind::AltRef => s / boost,
212 FrameKind::Inter => s * modulation,
213 };
214 let (step, target) = match &self.cfg.mode {
215 RcMode::ConstantQ { q_index } => (qstep(*q_index), 0.0),
216 RcMode::Crf { crf } => (step_for(qstep(crf_to_q(*crf))), 0.0),
217 RcMode::Bitrate { .. } => {
218 let (ek, ei) = (EXP[0], EXP[1]);
219 match kind {
220 FrameKind::Key => self.a_key = cost * boost.powf(ek),
221 _ => {
222 let a_now = cost * modulation.powf(-ei);
223 self.a_inter = if self.a_inter == 0.0 {
224 a_now
225 } else {
226 0.9 * self.a_inter + 0.1 * a_now
227 };
228 }
229 }
230 let a_inter = if self.a_inter > 0.0 {
231 self.a_inter
232 } else {
233 cost * self.ratio.clamp(0.05, 1.0)
234 };
235 let key_share = match self.cfg.keyint {
236 0 => 0.0,
237 n => 1.0 / n as f64,
238 };
239 let fill = (self.buffer / self.buffer_size).clamp(-0.6, 0.6);
240 let want = (self.frame_bits * (1.0 + fill)).max(1.0);
241 let s = solve(
242 [
243 (key_share * self.k[0] * self.a_key, ek),
244 ((1.0 - key_share) * self.k[1] * a_inter, ei),
245 ],
246 want,
247 );
248 let st = step_for(s);
249 (st, self.model(kind, cost, st))
250 }
251 RcMode::TwoPass { stats, .. } => {
252 let i = self.frame.min(stats.len().saturating_sub(1));
253 let (sk, si) = self.suffix.get(i).copied().unwrap_or((0.0, 0.0));
254 let s = solve(
255 [(self.k[0] * sk, EXP[0]), (self.k[1] * si, EXP[1])],
256 self.remaining.max(self.frame_bits * 0.25),
257 );
258 let st = match kind {
259 FrameKind::Inter => s * two_pass_modulation(cost, self.geo),
260 _ => s / boost,
261 };
262 (st, self.model(kind, cost, st))
263 }
264 };
265 let mut q = q_index_for(step);
266 if !matches!(self.cfg.mode, RcMode::ConstantQ { .. }) {
267 q = q.clamp(self.cfg.min_q, self.cfg.max_q.max(self.cfg.min_q));
268 }
269 if q > 0 && self.cfg.aq == Aq::Variance {
270 let (mi_cols, mi_rows) = (self.cfg.width.div_ceil(8), self.cfg.height.div_ceil(8));
271 self.seg.variance_aq(
272 q,
273 mi_cols,
274 mi_rows,
275 self.analyser.activity(),
276 self.analyser.blocks(),
277 );
278 } else {
279 self.seg.disable();
280 }
281 self.pending = Some((kind, a.cost(kind)));
282 FramePlan {
283 base_q_idx: q,
284 filter_level: guess_filter_level(q, kind),
285 sharpness: 0,
286 lambda: lambda(q),
287 target_bits: target.max(0.0) as u64,
288 }
289 }
290
291 pub fn update(&mut self, kind: FrameKind, bits: u64, base_q_idx: u8) {
292 let cost = match self.pending.take() {
293 Some((k, c)) if k == kind => c,
294 _ => return,
295 };
296 let slot = kind.slot();
297 let obs = (bits.max(8) as f64 * qstep(base_q_idx).powf(EXP[slot]) / cost as f64).max(1e-9);
298 if self.seen[slot] {
299 let r = (obs / self.k[slot]).clamp(0.25, 4.0);
300 self.k[slot] *= r.powf(0.4);
301 } else {
302 self.k[slot] = obs;
303 self.seen[slot] = true;
304 if slot == 2 && !self.seen[1] {
305 self.k[1] = obs;
306 }
307 if slot == 1 && !self.seen[2] {
308 self.k[2] = obs;
309 }
310 }
311 match self.cfg.mode {
312 RcMode::Bitrate { .. } => {
313 self.buffer = (self.buffer + self.frame_bits - bits as f64)
314 .clamp(-self.buffer_size, self.buffer_size);
315 }
316 RcMode::TwoPass { .. } => {
317 self.remaining -= bits as f64;
318 }
319 _ => {}
320 }
321 self.frame += 1;
322 }
323
324 fn boost(&self, kind: FrameKind) -> f64 {
325 let still = 1.0 - self.ratio.clamp(0.0, 1.0);
326 match kind {
327 FrameKind::Key => 1.95 + 0.9 * still,
328 FrameKind::Golden | FrameKind::AltRef => 1.12 + 0.3 * still,
329 FrameKind::Inter => 1.0,
330 }
331 }
332
333 fn model(&self, kind: FrameKind, cost: f64, step: f64) -> f64 {
334 let slot = kind.slot();
335 self.k[slot] * cost / step.powf(EXP[slot])
336 }
337}
338
339fn predicted_key(i: usize, keyint: u32) -> bool {
340 i == 0 || (keyint > 0 && i.is_multiple_of(keyint as usize))
341}
342
343fn two_pass_geo(stats: &[FrameAnalysis], keyint: u32) -> f64 {
344 let (mut s, mut n) = (0.0, 0.0);
345 for (i, a) in stats.iter().enumerate() {
346 if !predicted_key(i, keyint) {
347 s += (a.cost(FrameKind::Inter) as f64).ln();
348 n += 1.0;
349 }
350 }
351 if n > 0.0 { (s / n).exp() } else { 1.0 }
352}
353
354fn two_pass_modulation(cost: f64, geo: f64) -> f64 {
355 ((cost / geo).powf(1.0 - QCOMP)).clamp(0.5, 2.0)
356}
357
358fn two_pass_suffix(stats: &[FrameAnalysis], keyint: u32, geo: f64) -> Vec<(f64, f64)> {
359 let mut out = vec![(0.0, 0.0); stats.len()];
360 let (mut sk, mut si) = (0.0, 0.0);
361 let mut ratio = 0.5;
362 let ratios: Vec<f64> = stats
363 .iter()
364 .map(|a| {
365 if a.has_prev {
366 ratio = 0.8 * ratio + 0.2 * a.inter_ratio().min(1.0);
367 }
368 ratio
369 })
370 .collect();
371 for i in (0..stats.len()).rev() {
372 if predicted_key(i, keyint) {
373 let boost = 1.3 + 0.6 * (1.0 - ratios[i]);
374 sk += stats[i].cost(FrameKind::Key) as f64 * boost.powf(EXP[0]);
375 } else {
376 let c = stats[i].cost(FrameKind::Inter) as f64;
377 si += c * two_pass_modulation(c, geo).powf(-EXP[1]);
378 }
379 out[i] = (sk, si);
380 }
381 out
382}