collective

Browse / github.com/storytold/effectcraft / patch

av1enc: find the first frame's quantizer in fewer passes

av1enc: find the first frame's quantizer in fewer passes Part of #2390 (stacked on #2408: the base of this branch is the #2391 commit, so the diff here is the second commit only; merge #2408 first) ## Change (crates/av1enc/src/lib.rs, bitrate mode only) The first frame was coded four times with a fixed gain (30 quantizer steps per halving of the size). Now the first correction keeps that rule, and from the second correction on the next quantizer is placed in the log domain of the quantizer step, where the size is close to a line: false position between the nearest passes that enclose the target, otherwise a secant step with the measured slope. The

Status
open
Score
0.00
Adopters
0 other operators
Author
vangel76 (imported from pull request #2409, open)
Base
991904888a20b7da72b2a4880b21c7d912154736
Size
+168 −26 in 3 files
Review
pending
Shared
2026-10-28
Id
0191649dd77eb53c447449bf

How it went

No other operator has reported using it yet.

Diff

raw · crates/av1enc/README.md crates/av1enc/src/lib.rs crates/av1enc/tests/it/roundtrip.rs

diff --git a/crates/av1enc/README.md b/crates/av1enc/README.md--- a/crates/av1enc/README.md+++ b/crates/av1enc/README.md@@ -48,8 +48,10 @@ reconstruction, so prediction never drifts.   reconstruction loop; full coefficient coding (all_zero, tx type, eob_pt, eob_extra,   coeff_base_eob, coeff_base, coeff_br, dc_sign, Golomb) with the exact contexts. - **Rate control**: `ConstantQ(base_q_idx)`; `Bitrate { kbps }` adapts `base_q_idx` per frame-  towards the average target (key frames get a lower quantizer and a larger share; the first-  key frame is searched in a few passes).+  towards the average target (key frames get a lower quantizer and a larger share of the budget of+  their group of `keyint` frames, the inter frames the rest; the first key frame is searched in a+  few passes, a later key frame that misses its budget by more than 19 % is coded again; the+  deviation from the budget is paid back over 16 frames).  ## Limits diff --git a/crates/av1enc/src/lib.rs b/crates/av1enc/src/lib.rs--- a/crates/av1enc/src/lib.rs+++ b/crates/av1enc/src/lib.rs@@ -24,6 +24,7 @@ use bitw::{BitWriter, obu}; use decide::{Decider, FrameGeom, RefFrame}; use loopfilter::{LfInfo, loop_filter}; use pred::Plane;+use tables::AC_QLOOKUP; use writer::TileWriter;  /// Rate control.@@ -136,10 +137,118 @@ struct RateState {     frame_bits: f64,     /// Quantizer index for inter frames (continuous).     q: f64,+    /// Key frame quantizer index minus the inter frame one.+    key_off: f64,+    /// Quantizer index steps per halving of the frame size, measured from consecutive inter+    /// frames (the size/quantizer curve is much steeper than the default on noisy content).+    slope: f64,+    /// Quantizer index and size (in bits) of the previous inter frame.+    prev_inter: Option<(f64, f64)>,     /// Accumulated (spent - budget) bits.     debt: f64, } +/// Quantizer index steps per halving of the frame size before any measurement.+const RC_SLOPE: f64 = 28.0;+/// Fraction of the quantizer step the size error asks for that is applied per frame.+const RC_GAIN: f64 = 0.8;+/// Frames over which an accumulated deviation from the budget is paid back.+const RC_PAYBACK: f64 = 16.0;++impl RateState {+    /// The nominal bit budget and the starting quantizer index of a key or inter frame. A key+    /// frame gets a larger share, and the inter frames of the group get the rest, so that a whole+    /// group of `keyint` frames adds up to `keyint * frame_bits`.+    fn frame_budget(&self, keyint: u32, key: bool) -> (f64, f64) {+        let keyint = keyint.max(1) as f64;+        let boost = (1.0 + (keyint.clamp(1.0, 60.0) - 1.0).sqrt() * 0.75).min(0.5 * keyint + 0.5);+        if key {+            (self.frame_bits * boost, (self.q + self.key_off).clamp(1.0, 255.0))+        } else if keyint > 1.0 {+            (self.frame_bits * (keyint - boost) / (keyint - 1.0), self.q)+        } else {+            (self.frame_bits, self.q)+        }+    }++    /// Adapts the quantizers to the size of the frame just coded at `q_used` against `target`.+    fn update(&mut self, key: bool, first: bool, q_used: f64, bits: f64, target: f64) {+        let err = (bits / target).log2().clamp(-1.5, 1.5);+        if first && key {+            // four passes ran: the key frame quantizer is known, the inter one follows from the offset+            self.q = (q_used - self.key_off).clamp(1.0, 255.0);+        } else if key {+            self.key_off = q_used - self.q + RC_GAIN * self.slope * err;+        } else {+            if let Some((pq, pbits)) = self.prev_inter {+                let (dq, dr) = (q_used - pq, (pbits / bits).log2());+                if dq.abs() >= 3.0 && dr.abs() >= 0.15 && dq * dr > 0.0 {+                    self.slope = (0.6 * self.slope + 0.4 * (dq / dr)).clamp(3.0, 45.0);+                }+            }+            self.prev_inter = Some((q_used, bits));+            self.q = (q_used + RC_GAIN * self.slope * err).clamp(1.0, 255.0);+        }+    }+}++/// Most passes on the first frame, and how far from the target (octaves, 3 %) its size may be+/// before it is accepted.+const RC_FIRST_PASSES: usize = 4;+const RC_FIRST_TOLERANCE: f64 = 0.04;+/// Quantizer index steps per halving of the size, for the first correction of the first frame+/// (the starting quantizer comes from the bits per sample; the first pass shows how far off it+/// is, and from there on the passes are placed by `next_probe_q`).+const RC_FIRST_SLOPE: f64 = 30.0;+/// The same for a key frame after the first: coded again when its size is more than 0.25 octaves+/// (19 %) off the target.+const RC_KEY_PASSES: usize = 3;+const RC_KEY_TOLERANCE: f64 = 0.25;+/// Octaves of frame size per octave of quantizer step assumed for the first correction of a frame.+const RC_PROBE_SLOPE: f64 = 1.0;+/// Largest change of the quantizer step (octaves) between two passes on one frame.+const RC_PROBE_MAX_STEP: f64 = 2.0;++/// The quantizer index whose step size is closest to `2^x`.+fn q_for_step(bit_depth: u8, x: f64) -> u8 {+    let table = &AC_QLOOKUP[((bit_depth - 8) / 2) as usize];+    let mut best = 1;+    for q in 1..=255usize {+        if ((table[q] as f64).log2() - x).abs() < ((table[best] as f64).log2() - x).abs() {+            best = q;+        }+    }+    best as u8+}++/// The next quantizer to try for a frame, from the `(quantizer, log2(size / target))`+/// pairs coded so far. The size follows the quantizer step (not the index, which is linear in the+/// step only below about 60) in a near log-log line: false position between the closest pair of+/// passes that enclose the target, otherwise a step from the latest one with the slope of the+/// last two. `None` when no untried quantizer lies between the enclosing passes.+fn next_probe_q(bit_depth: u8, tried: &[(u8, f64)]) -> Option<u8> {+    let table = &AC_QLOOKUP[((bit_depth - 8) / 2) as usize];+    let x = |q: u8| (table[q as usize] as f64).log2();+    let over = tried.iter().filter(|t| t.1 > 0.0).max_by_key(|t| t.0);+    let under = tried.iter().filter(|t| t.1 <= 0.0).min_by_key(|t| t.0);+    if let (Some(&(q_lo, e_lo)), Some(&(q_hi, e_hi))) = (over, under) {+        if q_hi - q_lo < 2 {+            return None;+        }+        let next = x(q_lo) + (x(q_hi) - x(q_lo)) * e_lo / (e_lo - e_hi);+        return Some(q_for_step(bit_depth, next).clamp(q_lo + 1, q_hi - 1));+    }+    let &(q, err) = tried.last()?;+    let mut slope = RC_PROBE_SLOPE;+    if let [.., (q0, e0), (q1, e1)] = tried+        && (x(*q1) - x(*q0)).abs() > 0.05+        && (e0 - e1) / (x(*q1) - x(*q0)) > 0.0+    {+        slope = ((e0 - e1) / (x(*q1) - x(*q0))).clamp(0.25, 3.0);+    }+    Some(q_for_step(bit_depth, x(q) + (err / slope).clamp(-RC_PROBE_MAX_STEP, RC_PROBE_MAX_STEP)))+}+ /// AV1 encoder. Frames come out in input order (no hidden/reordered frames). pub struct Encoder {     cfg: EncoderConfig,@@ -185,7 +294,7 @@ impl Encoder {                 // initial guess from the bits per sample                 let bpp = frame_bits / (cfg.width as f64 * cfg.height as f64);                 let q = (110.0 - 40.0 * (bpp / 0.1).log2()).clamp(20.0, 230.0);-                Some(RateState { frame_bits, q, debt: 0.0 })+                Some(RateState { frame_bits, q, key_off: -24.0, slope: RC_SLOPE, prev_inter: None, debt: 0.0 })             }         };         let mut enc = Encoder { cfg, geom, level, seq_obu: Vec::new(), frame_num: 0, refr: None, rc };@@ -469,38 +578,52 @@ impl Encoder {      fn encode_rate_controlled(&mut self, src: &[Plane; 3], key: bool, rc: &mut RateState) -> (Vec<u8>, RefFrame) {         let first = self.frame_num == 0;-        let (frame_bits, debt, q_inter) = (rc.frame_bits, rc.debt, rc.q);-        // key frames get a lower quantizer and a larger share of the budget-        let key_boost = 1.0 + (self.cfg.keyint.clamp(1, 60) as f64 - 1.0).sqrt() * 0.75;-        let budget = if key { frame_bits * key_boost } else { frame_bits };-        // pay back a fraction of the accumulated deviation-        let target = (budget - debt * 0.25).max(budget * 0.25);-        let mut q = if key { q_inter - 24.0 } else { q_inter };-        let mut best: Option<(Vec<u8>, RefFrame, f64)> = None;-        let passes = if first { 4 } else { 1 };-        for pass in 0..passes {+        let (budget, q_start) = rc.frame_budget(self.cfg.keyint, key);+        // pay back the accumulated deviation over the next RC_PAYBACK frames+        let target = (budget - rc.debt / RC_PAYBACK).clamp(budget * 0.35, budget * 2.0);+        let mut q = q_start;+        let mut best: Option<(Vec<u8>, RefFrame, f64, f64)> = None;+        // The first frame is coded up to four times to find its quantizer (it fixes the+        // quantizers of everything after it), a later key frame up to three times (it is a large+        // part of the budget of its group). A frame is coded again only while its size is far from+        // the target and an untried quantizer is left to try; other frames are coded once.+        let (max_passes, tolerance) = if first {+            (RC_FIRST_PASSES, RC_FIRST_TOLERANCE)+        } else if key {+            (RC_KEY_PASSES, RC_KEY_TOLERANCE)+        } else {+            (1, 0.0)+        };+        let mut tried: Vec<(u8, f64)> = Vec::new();+        for _ in 0..max_passes {             let qi = q.round().clamp(1.0, 255.0) as u8;+            if tried.iter().any(|t| t.0 == qi) {+                break;+            }             let (obu, rf) = self.encode_frame(src, key, qi);             let bits = obu.len() as f64 * 8.0;             let err = (bits / target).log2();-            let better = best.as_ref().is_none_or(|b| err.abs() < (b.2 / target).log2().abs());-            if better {-                best = Some((obu, rf, bits));+            tried.push((qi, err));+            if best.as_ref().is_none_or(|b| err.abs() < (b.2 / target).log2().abs()) {+                best = Some((obu, rf, bits, qi as f64));+            }+            if err.abs() <= tolerance {+                break;             }-            if pass + 1 < passes {-                q = (q + 30.0 * err).clamp(1.0, 255.0);+            let next = if first && tried.len() == 1 {+                Some((q + RC_FIRST_SLOPE * err).clamp(1.0, 255.0))+            } else {+                next_probe_q(self.cfg.bit_depth, &tried).map(f64::from)+            };+            match next {+                Some(next) => q = next,+                None => break,             }         }         // At least one pass ran, so `best` is set.-        let Some((obu, rf, bits)) = best else { return self.encode_frame(src, key, q.round().clamp(1.0, 255.0) as u8) };+        let Some((obu, rf, bits, q_used)) = best else { return self.encode_frame(src, key, q.round().clamp(1.0, 255.0) as u8) };         rc.debt += bits - budget;-        // adapt the inter quantizer: about 30 index steps per doubling of the rate-        let err = (bits / target).log2().clamp(-1.0, 1.0);-        if first && key {-            rc.q = (q + 24.0).clamp(1.0, 255.0);-        } else {-            rc.q = (rc.q + 12.0 * err).clamp(1.0, 255.0);-        }+        rc.update(key, first, q_used, bits, target);         (obu, rf)     } diff --git a/crates/av1enc/tests/it/roundtrip.rs b/crates/av1enc/tests/it/roundtrip.rs--- a/crates/av1enc/tests/it/roundtrip.rs+++ b/crates/av1enc/tests/it/roundtrip.rs@@ -80,6 +80,23 @@ fn bitrate_mode_tracks_target() {     assert!(kbps > 400.0 / 2.0 && kbps < 400.0 * 2.0, "{kbps} kbit/s"); } +/// Over several key frame intervals the average rate stays within 5 % of the target, for a+/// range of targets (the key frame budget is taken out of the inter frames of its group).+#[test]+fn bitrate_mode_is_accurate_over_many_groups() {+    let n = 180;+    for (kbps, keyint, bd) in [(40, 30, 8), (100, 30, 8), (200, 60, 8), (150, 20, 10)] {+        let pics: Vec<Pic> = (0..n).map(|t| synth(192, 128, bd, t)).collect();+        let mut c = cfg(192, 128, bd as u8, 0, keyint);+        c.rate = RateControl::Bitrate { kbps };+        let mut enc = effectcraft_av1enc::Encoder::new(c).unwrap();+        let bits: usize = pics.iter().map(|p| enc.encode(&p.frame()).data.len() * 8).sum();+        let got = bits as f64 / (n as f64 / 30.0) / 1000.0;+        eprintln!("{bd} bit, keyint {keyint}: {got:.1} kbit/s for {kbps}");+        assert!((got / kbps as f64 - 1.0).abs() < 0.05, "{got:.1} kbit/s for a {kbps} kbit/s target ({bd} bit, keyint {keyint})");+    }+}+ #[test] fn headers() {     let enc = effectcraft_av1enc::Encoder::new(cfg(1920, 1080, 10, 100, 60)).unwrap();