diff --git a/crates/av1enc/README.md b/crates/av1enc/README.md index 88b2a6c6..5d1815e2 100644 --- a/crates/av1enc/README.md +++ b/crates/av1enc/README.md @@ -22,8 +22,8 @@ reconstruction, so prediction never drifts. `Encoder::av1c()` returns the `AV1CodecConfigurationRecord` (also the WebM `CodecPrivate`). - **Frames**: a `KEY_FRAME` every `keyint` frames, `INTER_FRAME`s otherwise. Inter frames use slot 0 (the previous reconstruction) for all seven references and refresh slot 0 only, so - `LAST_FRAME` is always the previous frame. `primary_ref_frame = PRIMARY_REF_NONE` (default - CDFs every frame), symbol-level CDF adaptation on, frame-end CDF update off. One tile, + `LAST_FRAME` is always the previous frame. CDFs adapt per symbol and are saved at the end of + every frame; inter frames load them from `LAST_FRAME` (`primary_ref_frame` = 0). One tile, `TX_MODE_LARGEST`, `reduced_tx_set`, no segmentation, no delta q / lf. - **Deblocking**: the normative loop filter (7.14) runs on the reconstruction; the frame level (the same for both luma directions and both chroma planes, no deltas, sharpness 0) is searched @@ -40,8 +40,9 @@ reconstruction, so prediction never drifts. from neighbour / co-located / parent candidates, a whole-sample diamond search and half- then quarter-sample refinement (`allow_high_precision_mv` = 0); the normative 8-tap `EIGHTTAP` prediction for luma and chroma; skip blocks; per-block intra / inter choice. -- **Residual**: one transform per plane and block (4x4 ... 32x32): `DCT_DCT` for luma, the - mode-dependent `DCT` / `ADST` combination for intra chroma; dead-zone quantisation with the +- **Residual**: one transform per plane and block (4x4 ... 32x32): for intra luma `DCT_DCT` or + the mode's preferred `ADST` combination (rate-distortion choice over the two best modes), the + mode-dependent type for intra chroma, `DCT_DCT` for inter blocks; dead-zone quantisation with the spec's `Dc_Qlookup` / `Ac_Qlookup`; the normative dequantisation and inverse transforms in the 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. diff --git a/crates/av1enc/src/cdf.rs b/crates/av1enc/src/cdf.rs index 8540e60a..0466a159 100644 --- a/crates/av1enc/src/cdf.rs +++ b/crates/av1enc/src/cdf.rs @@ -95,3 +95,70 @@ impl Cdfs { }) } } + +impl Cdfs { + /// The counter reset of load_cdfs( ): the last entry of every CDF array is set to 0. + pub fn clear_counts(&mut self) { + macro_rules! clr { + ($($f:ident),*) => { $( self.$f.clear_count(); )* }; + } + clr!( + y_mode, + uv_mode_cfl_not_allowed, + uv_mode_cfl_allowed, + angle_delta, + partition_w8, + partition_w16, + partition_w32, + partition_w64, + skip, + is_inter, + single_ref, + new_mv, + zero_mv, + ref_mv, + drl_mode, + mv_joint, + mv_class, + mv_class0_bit, + mv_fr, + mv_class0_fr, + mv_sign, + mv_bit, + intra_tx_type_set2, + inter_tx_type_set3, + txb_skip, + eob_pt_16, + eob_pt_32, + eob_pt_64, + eob_pt_128, + eob_pt_256, + eob_pt_512, + eob_pt_1024, + eob_extra, + dc_sign, + coeff_base_eob, + coeff_base, + coeff_br + ); + } +} + +/// Zeroes the adaptation counter (last element) of every innermost CDF array. +trait ClearCount { + fn clear_count(&mut self); +} + +impl ClearCount for [u16; N] { + fn clear_count(&mut self) { + self[N - 1] = 0; + } +} + +impl ClearCount for [T; M] { + fn clear_count(&mut self) { + for e in self.iter_mut() { + e.clear_count(); + } + } +} diff --git a/crates/av1enc/src/decide.rs b/crates/av1enc/src/decide.rs index 939cabaa..b877fc23 100644 --- a/crates/av1enc/src/decide.rs +++ b/crates/av1enc/src/decide.rs @@ -35,6 +35,8 @@ pub(crate) struct RefFrame { pub mvs: Vec<[i32; 2]>, /// The frame's loop filter level. pub lf_level: u32, + /// The CDFs saved at the end of the frame (loaded by the next frame). + pub cdfs: Box, } pub(crate) struct FrameGeom { @@ -339,19 +341,45 @@ impl<'a> Decider<'a> { bit_depth: bd, }; let ed = edges(&self.rec[0], &e); - let mut best = (f64::MAX, DC_PRED); - for &mode in INTRA_MODES_USED.iter() { + // the two best modes by SAD, then a rate-distortion choice of mode and transform type + let mut ranked: Vec<(f64, usize)> = INTRA_MODES_USED + .iter() + .map(|&mode| { + predict_intra(&ed, &e, mode, &mut pred); + let sad = self.sad(0, c * 4, r * 4, n, &pred[..n * n]) as f64; + (sad + self.lambda_sad * if mode == DC_PRED { 1.0 } else { MODE_BITS_INTRA }, mode) + }) + .collect(); + ranked.sort_by(|a, b| a.0.total_cmp(&b.0)); + let tx_y = tx_size_for(log2); + let mut best: Option<(f64, usize, usize, TxResult, Option>)> = None; + for &(_, mode) in ranked.iter().take(2) { predict_intra(&ed, &e, mode, &mut pred); - let sad = self.sad(0, c * 4, r * 4, n, &pred[..n * n]) as f64; - let cost = sad + self.lambda_sad * if mode == DC_PRED { 1.0 } else { MODE_BITS_INTRA }; - if cost < best.0 { - best = (cost, mode); - best_pred[..n * n].copy_from_slice(&pred[..n * n]); + let mut types = vec![DCT_DCT]; + let alt = MODE_TO_TXFM[mode] as usize; + if alt != DCT_DCT && tx_in_set(tx_set(tx_y, false), false, alt) { + types.push(alt); + } + for &t in &types { + let (res, recon) = self.try_tx(0, c * 4, r * 4, log2, t, &pred[..n * n], true); + let type_bits = if !res.nonzero || tx_set(tx_y, false) == TX_SET_DCTONLY { + 0.0 + } else if t == DCT_DCT { + 1.0 + } else { + 2.0 + }; + let mode_bits = if mode == DC_PRED { 1.0 } else { MODE_BITS_INTRA }; + let cost = res.dist + self.lambda * (res.bits + type_bits + mode_bits); + if best.as_ref().is_none_or(|b| cost < b.0) { + best_pred[..n * n].copy_from_slice(&pred[..n * n]); + best = Some((cost, mode, t, res, recon)); + } } } - let y_mode = best.1; - let tx_y = tx_size_for(log2); - let ty = self.code_tx(0, c * 4, r * 4, log2, DCT_DCT, &best_pred[..n * n], true); + let (_, y_mode, tx_type_y, ty, recon) = best.expect("an intra mode"); + self.put(0, c * 4, r * 4, n, recon.as_deref().unwrap_or(&best_pred[..n * n])); + let tx_type_y = if ty.nonzero { tx_type_y } else { DCT_DCT }; self.mark_decoded(0, r, c, n4); // chroma (one mode for both planes) let cn = n >> 1; @@ -391,12 +419,11 @@ impl<'a> Decider<'a> { tuv.push(self.code_tx(1 + i, c * 2, r * 2, clog2, uv_type, &pred[..cn * cn], true)); self.mark_decoded(1 + i, r, c, n4); } - let _ = tx_y; let skip = !ty.nonzero && !tuv[0].nonzero && !tuv[1].nonzero; let dist = ty.dist + tuv[0].dist + tuv[1].dist; let bits = if skip { 1.0 } else { ty.bits + tuv[0].bits + tuv[1].bits + 1.0 } + MODE_BITS_INTRA + 2.0; let [u, v]: [TxResult; 2] = tuv.try_into().ok().expect("two chroma planes"); - let leaf = Leaf { bsize, is_inter: false, y_mode, uv_mode, tx_type_y: DCT_DCT, mv: [0, 0], skip, coefs: [ty.levels, u.levels, v.levels] }; + let leaf = Leaf { bsize, is_inter: false, y_mode, uv_mode, tx_type_y, mv: [0, 0], skip, coefs: [ty.levels, u.levels, v.levels] }; (leaf, dist + self.lambda * bits) } @@ -688,6 +715,15 @@ impl<'a> Decider<'a> { /// Transform, quantise and reconstruct one square transform block of `plane` at (x, y) /// over the prediction `pred`; the reconstruction is written to `rec`. fn code_tx(&mut self, p: usize, x: usize, y: usize, log2: u32, tx_type: usize, pred: &[u16], intra: bool) -> TxResult { + let (res, recon) = self.try_tx(p, x, y, log2, tx_type, pred, intra); + let n = 1usize << log2; + self.put(p, x, y, n, recon.as_deref().unwrap_or(pred)); + res + } + + /// Transform and quantise one block; returns the result and its reconstruction (None: the + /// prediction) without writing it. + fn try_tx(&self, p: usize, x: usize, y: usize, log2: u32, tx_type: usize, pred: &[u16], intra: bool) -> (TxResult, Option>) { let n = 1usize << log2; let tx_sz = tx_size_for(log2); let bd = self.bd; @@ -715,19 +751,16 @@ impl<'a> Decider<'a> { } let zero_dist = self.sse(p, x, y, n, pred); if !nonzero { - self.put(p, x, y, n, pred); - return TxResult { levels, nonzero, dist: zero_dist, bits: 1.0 }; + return (TxResult { levels, nonzero, dist: zero_dist, bits: 1.0 }, None); } let recon = self.reconstruct(&levels, tx_sz, tx_type, pred, n, bd); let dist = self.sse(p, x, y, n, &recon); let bits = estimate_bits(&levels, tx_sz); if dist + self.lambda * bits >= zero_dist + self.lambda { levels.iter_mut().for_each(|l| *l = 0); - self.put(p, x, y, n, pred); - return TxResult { levels, nonzero: false, dist: zero_dist, bits: 1.0 }; + return (TxResult { levels, nonzero: false, dist: zero_dist, bits: 1.0 }, None); } - self.put(p, x, y, n, &recon); - TxResult { levels, nonzero: true, dist, bits } + (TxResult { levels, nonzero: true, dist, bits }, Some(recon)) } /// Dequantisation (7.12.3) and the inverse transform, added to the prediction. diff --git a/crates/av1enc/src/lib.rs b/crates/av1enc/src/lib.rs index 99212361..f1ab74a4 100644 --- a/crates/av1enc/src/lib.rs +++ b/crates/av1enc/src/lib.rs @@ -275,7 +275,7 @@ impl Encoder { w.f(1, 0); // disable_cdf_update w.f(1, 0); // frame_size_override_flag if !key { - w.f(3, 7); // primary_ref_frame = PRIMARY_REF_NONE + 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 ] @@ -288,7 +288,7 @@ impl Encoder { } else { w.f(1, 0); // render_and_frame_size_different } - w.f(1, 1); // disable_frame_end_update_cdf + 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| { @@ -364,7 +364,7 @@ impl Encoder { 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); + 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 { @@ -374,7 +374,7 @@ impl Encoder { } 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); - let tile = tw.finish(); + let (tile, cdfs) = tw.finish(); let mut w = BitWriter::new(); self.frame_header(&mut w, key, qidx, lf_level); w.byte_align(); @@ -383,7 +383,7 @@ impl Encoder { 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 }; + let rf = RefFrame { planes: dec.rec, mvs: dec.mvs, lf_level, cdfs }; (out, rf) } diff --git a/crates/av1enc/src/writer.rs b/crates/av1enc/src/writer.rs index 0f0262a0..903492b1 100644 --- a/crates/av1enc/src/writer.rs +++ b/crates/av1enc/src/writer.rs @@ -80,14 +80,23 @@ fn is_directional(mode: usize) -> bool { } impl<'a> TileWriter<'a> { - pub fn new(g: &'a FrameGeom, intra_only: bool, qidx: u8) -> Self { + /// `init`: the CDFs saved by the primary reference frame (load_cdfs), or None for the + /// defaults (init_non_coeff_cdfs / init_coeff_cdfs). + pub fn new(g: &'a FrameGeom, intra_only: bool, qidx: u8, init: Option<&Cdfs>) -> Self { let n = g.mi_cols * g.mi_rows; TileWriter { g, intra_only, qidx, sw: SymbolWriter::new(false), - cdf: Cdfs::new(qidx as u32), + cdf: match init { + Some(c) => { + let mut c = Box::new(c.clone()); + c.clear_counts(); + c + } + None => Cdfs::new(qidx as u32), + }, kf_y_mode: DEFAULT_INTRA_FRAME_Y_MODE_CDF, mi: Mi { cols: g.mi_cols, @@ -114,8 +123,9 @@ impl<'a> TileWriter<'a> { } } - pub fn finish(self) -> Vec { - self.sw.finish() + /// The tile data and the final CDFs (saved for the frame-end CDF update). + pub fn finish(self) -> (Vec, Box) { + (self.sw.finish(), self.cdf) } /// The mode info the loop filter needs. diff --git a/crates/av1enc/tests/ffmpeg_oracle.rs b/crates/av1enc/tests/ffmpeg_oracle.rs index dc13aed6..c1a087d2 100644 --- a/crates/av1enc/tests/ffmpeg_oracle.rs +++ b/crates/av1enc/tests/ffmpeg_oracle.rs @@ -117,6 +117,11 @@ fn libdav1d_decodes_bit_exact() { // all-intra at a fine quantizer (large coefficients, Golomb codes) let pics: Vec = (0..2).map(|t| noise(64, 48, 8, t as u32 + 3)).collect(); check("noise-intra", cfg(64, 48, 8, RateControl::ConstantQ(4), 1), &pics); + // the finest quantizer: the largest coefficients, 8 and 10 bit, intra and inter + let pics: Vec = (0..3).map(|t| noise(72, 40, 8, t as u32 + 11)).collect(); + check("noise-q1", cfg(72, 40, 8, RateControl::ConstantQ(1), 2), &pics); + let pics: Vec = (0..3).map(|t| noise(40, 72, 10, t as u32 + 21)).collect(); + check("noise-q1-10", cfg(40, 72, 10, RateControl::ConstantQ(1), 2), &pics); // odd sizes, 8 and 10 bit let pics: Vec = (0..4).map(|t| synth(33, 17, 8, t)).collect(); check("odd8", cfg(33, 17, 8, RateControl::ConstantQ(60), 3), &pics);