mirror of
https://github.com/storytold/effectcraft.git
synced 2026-11-10 20:41:25 +00:00
JonRC
6d5e84e7a5
Every MP4 export said BT.709 whatever the Output color space, in the colr box and in the bitstream, so Rec. 2020, P3, PQ and HLG renders were labelled as BT.709. Move the ProRes colr mapping to output_codes and use it for every movie path: the nclx colr box, FilmCraft's H.264 ColorConfig (VUI), and new colour fields in effectcraft-hevcenc (VUI) and effectcraft-av1enc (color_config), which wrote 1-1-1 as constants. The matrix stays BT.709, the coefficients every path converts with. Default renders don't change.
671 lines
27 KiB
Rust
671 lines
27 KiB
Rust
//! Clean-room pure-Rust AV1 encoder (Main profile, 4:2:0, 8 and 10 bit), implemented from the
|
||
//! AV1 Bitstream & Decoding Process Specification v1.0.0 with Errata 1.
|
||
//!
|
||
//! Key frames and single-reference inter frames (LAST_FRAME = the previous frame), 64x64
|
||
//! superblocks split down to 8x8 by rate-distortion decisions, intra modes DC / V / H / Paeth /
|
||
//! smooth, full-sample motion search with the spec's motion vector prediction, DCT / ADST
|
||
//! transforms up to 32x32, one tile per frame, in-loop filters off. See the crate README.
|
||
//!
|
||
//! The public API below is the contract `crates/export` codes against; keep it stable.
|
||
|
||
#![deny(clippy::unwrap_used, clippy::expect_used, clippy::panic, clippy::unimplemented, clippy::todo, clippy::unreachable)]
|
||
|
||
mod bitw;
|
||
mod cdf;
|
||
mod decide;
|
||
mod loopfilter;
|
||
mod pred;
|
||
mod symbol;
|
||
mod tables;
|
||
mod transform;
|
||
mod writer;
|
||
|
||
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.
|
||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||
pub enum RateControl {
|
||
/// Constant quantizer index (`base_q_idx`, 1–255; lower is better).
|
||
ConstantQ(u8),
|
||
/// Average bitrate target in kbit/s (frame-level quantizer adaptation).
|
||
Bitrate { kbps: u32 },
|
||
}
|
||
|
||
/// Encoder settings. Main profile (seq_profile 0), 4:2:0.
|
||
#[derive(Clone, Debug, PartialEq)]
|
||
pub struct EncoderConfig {
|
||
pub width: u32,
|
||
pub height: u32,
|
||
pub fps_num: u32,
|
||
pub fps_den: u32,
|
||
/// 8 or 10.
|
||
pub bit_depth: u8,
|
||
/// `seq_level_idx` (e.g. 8 = level 4.0); `None` picks the lowest level that fits.
|
||
pub level_idx: Option<u8>,
|
||
pub rate: RateControl,
|
||
/// Frames between key frames (1 = all key frames).
|
||
pub keyint: u32,
|
||
/// BT.709 limited range is signalled unless this is set.
|
||
pub full_range: bool,
|
||
/// The colour description of `color_config` (ITU-T H.273 code points); BT.709 by default.
|
||
pub color_primaries: u8,
|
||
pub transfer_characteristics: u8,
|
||
pub matrix_coefficients: u8,
|
||
}
|
||
|
||
impl EncoderConfig {
|
||
pub fn new(width: u32, height: u32, fps_num: u32, fps_den: u32) -> Self {
|
||
EncoderConfig {
|
||
width,
|
||
height,
|
||
fps_num,
|
||
fps_den,
|
||
bit_depth: 8,
|
||
level_idx: None,
|
||
rate: RateControl::ConstantQ(100),
|
||
keyint: 60,
|
||
full_range: false,
|
||
color_primaries: 1,
|
||
transfer_characteristics: 1,
|
||
matrix_coefficients: 1,
|
||
}
|
||
}
|
||
}
|
||
|
||
/// One input picture, 4:2:0, samples at the configured bit depth.
|
||
#[derive(Clone, Copy, Debug)]
|
||
pub struct Frame<'a> {
|
||
pub y: &'a [u16],
|
||
pub u: &'a [u16],
|
||
pub v: &'a [u16],
|
||
pub y_stride: usize,
|
||
pub uv_stride: usize,
|
||
}
|
||
|
||
/// One temporal unit as stored in MP4 / WebM samples: OBUs with `obu_has_size_field` = 1 and
|
||
/// no temporal delimiter (key frames start with the sequence header OBU).
|
||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||
pub struct Packet {
|
||
pub data: Vec<u8>,
|
||
pub keyframe: bool,
|
||
}
|
||
|
||
#[derive(Debug, PartialEq, Eq)]
|
||
pub enum Error {
|
||
InvalidConfig(String),
|
||
}
|
||
|
||
impl std::fmt::Display for Error {
|
||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||
match self {
|
||
Error::InvalidConfig(s) => write!(f, "invalid AV1 encoder config: {s}"),
|
||
}
|
||
}
|
||
}
|
||
|
||
impl std::error::Error for Error {}
|
||
|
||
const OBU_SEQUENCE_HEADER: u8 = 1;
|
||
const OBU_FRAME: u8 = 6;
|
||
|
||
/// Annex A level limits: (seq_level_idx, MaxPicSize, MaxHSize, MaxVSize, MaxDisplayRate,
|
||
/// MaxBitrate (Main tier, Mbit/s)).
|
||
const LEVELS: [(u8, u64, u32, u32, u64, u64); 14] = [
|
||
(0, 147_456, 2048, 1152, 4_423_680, 1_500_000),
|
||
(1, 278_784, 2816, 1584, 8_363_520, 3_000_000),
|
||
(4, 665_856, 4352, 2448, 19_975_680, 6_000_000),
|
||
(5, 1_065_024, 5504, 3096, 31_950_720, 10_000_000),
|
||
(8, 2_359_296, 6144, 3456, 70_778_880, 12_000_000),
|
||
(9, 2_359_296, 6144, 3456, 141_557_760, 20_000_000),
|
||
(12, 8_912_896, 8192, 4352, 267_386_880, 30_000_000),
|
||
(13, 8_912_896, 8192, 4352, 534_773_760, 40_000_000),
|
||
(14, 8_912_896, 8192, 4352, 1_069_547_520, 60_000_000),
|
||
(15, 8_912_896, 8192, 4352, 1_069_547_520, 60_000_000),
|
||
(16, 35_651_584, 16384, 8704, 1_069_547_520, 60_000_000),
|
||
(17, 35_651_584, 16384, 8704, 2_139_095_040, 100_000_000),
|
||
(18, 35_651_584, 16384, 8704, 4_278_190_080, 160_000_000),
|
||
(19, 35_651_584, 16384, 8704, 4_278_190_080, 160_000_000),
|
||
];
|
||
|
||
/// The lowest level whose limits fit the stream (31 = no level limits when none does).
|
||
fn pick_level(cfg: &EncoderConfig) -> u8 {
|
||
let pic = cfg.width as u64 * cfg.height as u64;
|
||
let rate = pic * cfg.fps_num as u64 / cfg.fps_den.max(1) as u64;
|
||
let bitrate = match cfg.rate {
|
||
RateControl::Bitrate { kbps } => kbps as u64 * 1000,
|
||
RateControl::ConstantQ(_) => 0,
|
||
};
|
||
for &(idx, max_pic, max_h, max_v, max_rate, max_br) in LEVELS.iter() {
|
||
if pic <= max_pic && cfg.width <= max_h && cfg.height <= max_v && rate <= max_rate && bitrate <= max_br {
|
||
return idx;
|
||
}
|
||
}
|
||
31
|
||
}
|
||
|
||
/// Frame-level rate control state for `RateControl::Bitrate`.
|
||
struct RateState {
|
||
/// Bits per frame at the target rate.
|
||
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,
|
||
geom: FrameGeom,
|
||
level: u8,
|
||
seq_obu: Vec<u8>,
|
||
frame_num: u64,
|
||
refr: Option<RefFrame>,
|
||
rc: Option<RateState>,
|
||
}
|
||
|
||
impl Encoder {
|
||
pub fn new(cfg: EncoderConfig) -> Result<Encoder, Error> {
|
||
if cfg.width == 0 || cfg.height == 0 {
|
||
return Err(Error::InvalidConfig("empty frame".into()));
|
||
}
|
||
if cfg.width > 4096 || cfg.height > 4096 {
|
||
return Err(Error::InvalidConfig(format!("{}x{} exceeds the single-tile limit (4096 samples per side)", cfg.width, cfg.height)));
|
||
}
|
||
if cfg.bit_depth != 8 && cfg.bit_depth != 10 {
|
||
return Err(Error::InvalidConfig(format!("bit depth {} (8 or 10 supported)", cfg.bit_depth)));
|
||
}
|
||
if cfg.fps_num == 0 || cfg.fps_den == 0 {
|
||
return Err(Error::InvalidConfig("frame rate".into()));
|
||
}
|
||
let geom = FrameGeom::new(cfg.width as usize, cfg.height as usize);
|
||
if geom.sb_cols * geom.sb_rows > 2304 {
|
||
return Err(Error::InvalidConfig(format!("{}x{} exceeds the single-tile area limit (MAX_TILE_AREA)", cfg.width, cfg.height)));
|
||
}
|
||
if let RateControl::Bitrate { kbps } = cfg.rate
|
||
&& kbps == 0
|
||
{
|
||
return Err(Error::InvalidConfig("zero bitrate".into()));
|
||
}
|
||
let level = cfg.level_idx.unwrap_or_else(|| pick_level(&cfg));
|
||
if level > 31 {
|
||
return Err(Error::InvalidConfig(format!("seq_level_idx {level}")));
|
||
}
|
||
let rc = match cfg.rate {
|
||
RateControl::ConstantQ(_) => None,
|
||
RateControl::Bitrate { kbps } => {
|
||
let frame_bits = kbps as f64 * 1000.0 * cfg.fps_den as f64 / cfg.fps_num as f64;
|
||
// 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, 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 };
|
||
let payload = enc.sequence_header_payload();
|
||
obu(OBU_SEQUENCE_HEADER, &payload, &mut enc.seq_obu);
|
||
Ok(enc)
|
||
}
|
||
|
||
pub fn config(&self) -> &EncoderConfig {
|
||
&self.cfg
|
||
}
|
||
|
||
/// The sequence header OBU (with size field).
|
||
pub fn sequence_header_obu(&self) -> Vec<u8> {
|
||
self.seq_obu.clone()
|
||
}
|
||
|
||
/// AV1CodecConfigurationRecord (`av1C` payload, also the WebM CodecPrivate).
|
||
pub fn av1c(&self) -> Vec<u8> {
|
||
let high = (self.cfg.bit_depth > 8) as u8;
|
||
let mut v = vec![
|
||
0x81, // marker (1), version (1)
|
||
self.level & 0x1f, // seq_profile 0, seq_level_idx_0
|
||
(high << 6) | (1 << 3) | (1 << 2), // tier 0, high_bitdepth, twelve_bit 0, mono 0, ss_x 1, ss_y 1, csp 0
|
||
0, // no initial_presentation_delay
|
||
];
|
||
v.extend_from_slice(&self.seq_obu);
|
||
v
|
||
}
|
||
|
||
/// sequence_header_obu( ) payload (5.5) including trailing bits.
|
||
fn sequence_header_payload(&self) -> Vec<u8> {
|
||
let c = &self.cfg;
|
||
let mut w = BitWriter::new();
|
||
w.f(3, 0); // seq_profile
|
||
w.f(1, 0); // still_picture
|
||
w.f(1, 0); // reduced_still_picture_header
|
||
w.f(1, 0); // timing_info_present_flag
|
||
w.f(1, 0); // initial_display_delay_present_flag
|
||
w.f(5, 0); // operating_points_cnt_minus_1
|
||
w.f(12, 0); // operating_point_idc[ 0 ]
|
||
w.f(5, self.level as u32); // seq_level_idx[ 0 ]
|
||
if self.level > 7 {
|
||
w.f(1, 0); // seq_tier[ 0 ]
|
||
}
|
||
let wbits = 32 - (c.width - 1).max(1).leading_zeros();
|
||
let hbits = 32 - (c.height - 1).max(1).leading_zeros();
|
||
w.f(4, wbits - 1);
|
||
w.f(4, hbits - 1);
|
||
w.f(wbits, c.width - 1);
|
||
w.f(hbits, c.height - 1);
|
||
w.f(1, 0); // frame_id_numbers_present_flag
|
||
w.f(1, 0); // use_128x128_superblock
|
||
w.f(1, 0); // enable_filter_intra
|
||
w.f(1, 0); // enable_intra_edge_filter
|
||
w.f(1, 0); // enable_interintra_compound
|
||
w.f(1, 0); // enable_masked_compound
|
||
w.f(1, 0); // enable_warped_motion
|
||
w.f(1, 0); // enable_dual_filter
|
||
w.f(1, 0); // enable_order_hint
|
||
w.f(1, 0); // seq_choose_screen_content_tools
|
||
w.f(1, 0); // seq_force_screen_content_tools
|
||
w.f(1, 0); // enable_superres
|
||
w.f(1, 0); // enable_cdef
|
||
w.f(1, 0); // enable_restoration
|
||
// color_config( )
|
||
w.f(1, (c.bit_depth > 8) as u32); // high_bitdepth
|
||
w.f(1, 0); // mono_chrome
|
||
w.f(1, 1); // color_description_present_flag
|
||
w.f(8, c.color_primaries as u32); // color_primaries
|
||
w.f(8, c.transfer_characteristics as u32); // transfer_characteristics
|
||
w.f(8, c.matrix_coefficients as u32); // matrix_coefficients
|
||
w.f(1, c.full_range as u32); // color_range
|
||
w.f(2, 0); // chroma_sample_position: unknown
|
||
w.f(1, 0); // separate_uv_delta_q
|
||
w.f(1, 0); // film_grain_params_present
|
||
w.trailing_bits();
|
||
w.into_bytes()
|
||
}
|
||
|
||
/// uncompressed_header( ) (5.9) for a shown KEY_FRAME or an INTER_FRAME referencing slot 0.
|
||
fn frame_header(&self, w: &mut BitWriter, key: bool, qidx: u8, lf_level: u32) {
|
||
w.f(1, 0); // show_existing_frame
|
||
w.f(2, if key { 0 } else { 1 }); // frame_type
|
||
w.f(1, 1); // show_frame
|
||
if !key {
|
||
w.f(1, 0); // error_resilient_mode
|
||
}
|
||
w.f(1, 0); // disable_cdf_update
|
||
w.f(1, 0); // frame_size_override_flag
|
||
if !key {
|
||
w.f(3, 0); // primary_ref_frame: LAST_FRAME's saved CDFs
|
||
w.f(8, 1); // refresh_frame_flags: slot 0 becomes LAST for the next frame
|
||
for _ in 0..7 {
|
||
w.f(3, 0); // ref_frame_idx[ i ]
|
||
}
|
||
w.f(1, 0); // render_and_frame_size_different
|
||
w.f(1, 0); // allow_high_precision_mv
|
||
w.f(1, 0); // is_filter_switchable
|
||
w.f(2, 0); // interpolation_filter = EIGHTTAP
|
||
w.f(1, 0); // is_motion_mode_switchable
|
||
} else {
|
||
w.f(1, 0); // render_and_frame_size_different
|
||
}
|
||
w.f(1, 0); // disable_frame_end_update_cdf
|
||
// tile_info( ): uniform spacing, one tile
|
||
let g = &self.geom;
|
||
let tile_log2 = |blk: usize, target: usize| {
|
||
let mut k = 0;
|
||
while (blk << k) < target {
|
||
k += 1;
|
||
}
|
||
k
|
||
};
|
||
w.f(1, 1); // uniform_tile_spacing_flag
|
||
if tile_log2(1, g.sb_cols.min(64)) > 0 {
|
||
w.f(1, 0); // increment_tile_cols_log2
|
||
}
|
||
if tile_log2(1, g.sb_rows.min(64)) > 0 {
|
||
w.f(1, 0); // increment_tile_rows_log2
|
||
}
|
||
// quantization_params( )
|
||
w.f(8, qidx as u32);
|
||
w.f(1, 0); // DeltaQYDc
|
||
w.f(1, 0); // DeltaQUDc
|
||
w.f(1, 0); // DeltaQUAc
|
||
w.f(1, 0); // using_qmatrix
|
||
w.f(1, 0); // segmentation_enabled
|
||
w.f(1, 0); // delta_q_present (base_q_idx > 0)
|
||
// loop_filter_params( ): the same level for every plane and direction
|
||
w.f(6, lf_level);
|
||
w.f(6, lf_level);
|
||
if lf_level != 0 {
|
||
w.f(6, lf_level);
|
||
w.f(6, lf_level);
|
||
}
|
||
w.f(3, 0); // loop_filter_sharpness
|
||
w.f(1, 0); // loop_filter_delta_enabled
|
||
w.f(1, 0); // tx_mode_select (TX_MODE_LARGEST)
|
||
if !key {
|
||
w.f(1, 0); // reference_select
|
||
}
|
||
w.f(1, 1); // reduced_tx_set
|
||
if !key {
|
||
for _ in 0..7 {
|
||
w.f(1, 0); // is_global
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Copies a picture into superblock-padded planes (edge replication).
|
||
fn load_source(&self, frame: &Frame) -> [Plane; 3] {
|
||
let maxv = (1u16 << self.cfg.bit_depth) - 1;
|
||
std::array::from_fn(|p| {
|
||
let (vw, vh) = self.geom.visible(p);
|
||
let (pw, ph) = self.geom.padded(p);
|
||
let (data, stride) = match p {
|
||
0 => (frame.y, frame.y_stride),
|
||
1 => (frame.u, frame.uv_stride),
|
||
_ => (frame.v, frame.uv_stride),
|
||
};
|
||
assert!(stride >= vw && data.len() >= stride * (vh - 1) + vw, "plane {p} is too small for {vw}x{vh}");
|
||
let mut pl = Plane::new(pw, ph);
|
||
for y in 0..ph {
|
||
let sy = y.min(vh - 1);
|
||
let src = &data[sy * stride..sy * stride + vw];
|
||
let row = &mut pl.data[y * pw..y * pw + pw];
|
||
for (x, d) in row.iter_mut().enumerate() {
|
||
*d = src[x.min(vw - 1)].min(maxv);
|
||
}
|
||
}
|
||
pl
|
||
})
|
||
}
|
||
|
||
/// Codes one frame at `qidx`; returns the frame OBU and the reconstruction.
|
||
fn encode_frame(&self, src: &[Plane; 3], key: bool, qidx: u8) -> (Vec<u8>, RefFrame) {
|
||
let g = &self.geom;
|
||
let refr = if key { None } else { self.refr.as_ref() };
|
||
let mut dec = Decider::new(g, self.cfg.bit_depth as u32, qidx, key, src, refr);
|
||
let mut tw = TileWriter::new(g, key, qidx, refr.map(|r| &*r.cdfs));
|
||
for sbr in 0..g.sb_rows {
|
||
tw.start_sb_row();
|
||
for sbc in 0..g.sb_cols {
|
||
let node = dec.superblock(sbr * 16, sbc * 16);
|
||
tw.superblock(sbr * 16, sbc * 16, &node);
|
||
}
|
||
}
|
||
let lf_level = self.pick_filter_level(&dec.rec, &tw.lf_info(), src, qidx);
|
||
loop_filter(&mut dec.rec, &tw.lf_info(), [lf_level; 4], self.cfg.bit_depth as u32, false, 1);
|
||
let (tile, cdfs) = tw.finish();
|
||
let mut w = BitWriter::new();
|
||
self.frame_header(&mut w, key, qidx, lf_level);
|
||
w.byte_align();
|
||
// tile_group_obu( ): a single tile, no tile_start_and_end_present_flag
|
||
let mut payload = w.into_bytes();
|
||
payload.extend_from_slice(&tile);
|
||
let mut out = Vec::with_capacity(payload.len() + 8);
|
||
obu(OBU_FRAME, &payload, &mut out);
|
||
let rf = RefFrame { planes: dec.rec, mvs: dec.mvs, lf_level, cdfs };
|
||
(out, rf)
|
||
}
|
||
|
||
/// The loop filter level minimising the luma error (searched around a guess from the
|
||
/// quantizer, on the luma of a subset of the superblock rows).
|
||
fn pick_filter_level(&self, rec: &[Plane; 3], info: &LfInfo, src: &[Plane; 3], qidx: u8) -> u32 {
|
||
let (w, h) = (self.geom.width, self.geom.height);
|
||
let sb_row_step = if self.geom.sb_rows >= 9 { 3 } else { 1 };
|
||
let sse = |pl: &Plane| -> u64 {
|
||
let mut s = 0u64;
|
||
for y in 0..h {
|
||
if !(y >> 6).is_multiple_of(sb_row_step) {
|
||
continue;
|
||
}
|
||
for (a, b) in pl.row(y)[..w].iter().zip(&src[0].row(y)[..w]) {
|
||
let d = *a as i64 - *b as i64;
|
||
s += (d * d) as u64;
|
||
}
|
||
}
|
||
s
|
||
};
|
||
let bd = self.cfg.bit_depth as u32;
|
||
let mut cache: Vec<(u32, u64)> = vec![(0, sse(&rec[0]))];
|
||
// one working copy of the luma plane for all candidate levels
|
||
let mut planes = [rec[0].clone(), Plane::new(0, 0), Plane::new(0, 0)];
|
||
let mut eval = |lvl: u32| -> u64 {
|
||
if let Some(&(_, e)) = cache.iter().find(|(l, _)| *l == lvl) {
|
||
return e;
|
||
}
|
||
planes[0].data.copy_from_slice(&rec[0].data);
|
||
loop_filter(&mut planes, info, [lvl; 4], bd, true, sb_row_step);
|
||
let e = sse(&planes[0]);
|
||
cache.push((lvl, e));
|
||
e
|
||
};
|
||
let guess = ((qidx as u32 * 3) / 16).clamp(2, 50);
|
||
let mut best = (0u32, eval(0));
|
||
let mut step = (guess / 2).max(1);
|
||
let mut center = guess;
|
||
let e = eval(center);
|
||
if e < best.1 {
|
||
best = (center, e);
|
||
}
|
||
while step >= 2 {
|
||
for lvl in [center.saturating_sub(step), (center + step).min(63)] {
|
||
let e = eval(lvl);
|
||
if e < best.1 {
|
||
best = (lvl, e);
|
||
}
|
||
}
|
||
center = best.0.max(1);
|
||
step /= 2;
|
||
}
|
||
best.0
|
||
}
|
||
|
||
/// Encodes one frame.
|
||
pub fn encode(&mut self, frame: &Frame) -> Packet {
|
||
let key = self.frame_num.is_multiple_of(self.cfg.keyint.max(1) as u64);
|
||
let src = self.load_source(frame);
|
||
let (frame_obu, rf) = match self.cfg.rate {
|
||
RateControl::ConstantQ(q) => self.encode_frame(&src, key, q.max(1)),
|
||
RateControl::Bitrate { .. } => match self.rc.take() {
|
||
Some(mut rc) => {
|
||
let out = self.encode_rate_controlled(&src, key, &mut rc);
|
||
self.rc = Some(rc);
|
||
out
|
||
}
|
||
// The state is created with the encoder for bitrate mode; without it, encode at a
|
||
// middle quantizer rather than fail.
|
||
None => self.encode_frame(&src, key, 128),
|
||
},
|
||
};
|
||
self.refr = Some(rf);
|
||
self.frame_num += 1;
|
||
let mut data = Vec::with_capacity(frame_obu.len() + self.seq_obu.len());
|
||
if key {
|
||
data.extend_from_slice(&self.seq_obu);
|
||
}
|
||
data.extend_from_slice(&frame_obu);
|
||
Packet { data, keyframe: key }
|
||
}
|
||
|
||
fn encode_rate_controlled(&mut self, src: &[Plane; 3], key: bool, rc: &mut RateState) -> (Vec<u8>, RefFrame) {
|
||
let first = self.frame_num == 0;
|
||
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();
|
||
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;
|
||
}
|
||
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, q_used)) = best else { return self.encode_frame(src, key, q.round().clamp(1.0, 255.0) as u8) };
|
||
rc.debt += bits - budget;
|
||
rc.update(key, first, q_used, bits, target);
|
||
(obu, rf)
|
||
}
|
||
|
||
/// The loop filter level of the last encoded frame. For tests.
|
||
#[doc(hidden)]
|
||
pub fn last_loop_filter_level(&self) -> Option<u32> {
|
||
self.refr.as_ref().map(|r| r.lf_level)
|
||
}
|
||
|
||
/// The reconstruction of the last encoded frame (visible area, planes Y, U, V), equal to
|
||
/// what a decoder outputs. For tests.
|
||
#[doc(hidden)]
|
||
pub fn last_reconstruction(&self) -> Option<[Vec<u16>; 3]> {
|
||
let rf = self.refr.as_ref()?;
|
||
Some(std::array::from_fn(|p| {
|
||
let (vw, vh) = self.geom.visible(p);
|
||
let pl = &rf.planes[p];
|
||
let mut v = Vec::with_capacity(vw * vh);
|
||
for y in 0..vh {
|
||
v.extend_from_slice(&pl.row(y)[..vw]);
|
||
}
|
||
v
|
||
}))
|
||
}
|
||
}
|