M13.4: av1enc key and inter frames, bit-exact with libdav1d
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Brandon Thomas
parent
7f97c25c3e
commit
4130172bbe
|
|
@ -0,0 +1,67 @@
|
|||
# effectcraft-av1enc
|
||||
|
||||
Clean-room, pure-Rust AV1 encoder for EffectCraft's MP4 / WebM export (layer L0, `std` only,
|
||||
no `unsafe`, builds for `wasm32-unknown-unknown`).
|
||||
|
||||
Output is decoded **bit-exactly** by libdav1d (`ffmpeg -c:v libdav1d`, checked in
|
||||
`tests/ffmpeg_oracle.rs` when ffmpeg is installed) and by FilmCraft's spec-derived
|
||||
`filmcraft-av1` decoder (`tests/roundtrip.rs`): every decoded frame equals the encoder's own
|
||||
reconstruction, so prediction never drifts.
|
||||
|
||||
## What is implemented
|
||||
|
||||
- **Stream**: Main profile (`seq_profile` 0), 4:2:0, 8 or 10 bit, BT.709 colour description,
|
||||
limited or full range. Sequence header with one operating point; `seq_level_idx` given or the
|
||||
lowest Annex A level whose MaxPicSize / MaxH/VSize / MaxDisplayRate (and, in bitrate mode, Main
|
||||
tier MaxBitrate) fit. Order hints, film grain, superres, CDEF, loop restoration, intra edge
|
||||
filter, filter intra, inter-intra, masked / distance compound, warped motion and dual filter
|
||||
are disabled in the sequence header.
|
||||
- **Packets**: one temporal unit per frame as stored in MP4 / Matroska samples: OBUs with size
|
||||
fields, no temporal delimiter; key frames start with the sequence header OBU, followed by one
|
||||
`OBU_FRAME` (frame header + tile group). Every frame is shown, in input order.
|
||||
`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, loop
|
||||
filter level 0, `TX_MODE_LARGEST`, `reduced_tx_set`, no segmentation, no delta q / lf.
|
||||
- **Tiles**: the multi-symbol arithmetic encoder is the exact inverse of the spec's symbol
|
||||
decoder (8.2), including the exit process padding; CDF adaptation as in 8.2.6.
|
||||
- **Partitions**: 64x64 superblocks, `PARTITION_NONE` / `PARTITION_SPLIT` down to 8x8 chosen by
|
||||
a rate-distortion cost (SSE + λ·estimated bits); blocks crossing the frame edge are split as
|
||||
the syntax requires. 64x64 blocks are only used as inter skip blocks.
|
||||
- **Intra**: DC, V, H, Paeth, smooth, smooth-V, smooth-H for luma and chroma (chosen by SAD),
|
||||
with the spec's edge preparation and above-right / below-left availability.
|
||||
- **Inter**: single reference (`LAST_FRAME`), `NEWMV` / `NEARESTMV` / `NEARMV` / `GLOBALMV`
|
||||
chosen from the spec's motion vector prediction stack (7.10.2, mirrored exactly), full-sample
|
||||
motion search (candidates + diamond search), the normative 8-tap `EIGHTTAP` prediction for
|
||||
luma and (half-sample) 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
|
||||
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.
|
||||
- **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).
|
||||
|
||||
## Limits
|
||||
|
||||
- One tile: width and height up to 4096 and at most 2304 superblocks (4096x2304 samples).
|
||||
- No in-loop filters (deblocking, CDEF, loop restoration), so compression at low rates is below
|
||||
mature encoders; no sub-sample motion search, no compound prediction, no 4xN blocks.
|
||||
- 4:2:0 only (Main profile), no 12 bit.
|
||||
|
||||
## Specification and provenance
|
||||
|
||||
Implemented from the **AV1 Bitstream & Decoding Process Specification, version 1.0.0 with
|
||||
Errata 1** (AOMediaCodec, 2019; the edition FilmCraft's decoder follows) and the **AV1 Codec
|
||||
ISO Media File Format Binding** (v1.2.0) for `av1C`.
|
||||
|
||||
The constant tables (default CDFs, scan orders, quantizer lookup tables, interpolation filter
|
||||
taps, smooth weights, block / transform size tables, ...) and the decoder-side processes the
|
||||
encoder must reproduce exactly (inverse transforms, intra and inter prediction rounding, context
|
||||
derivations, motion vector prediction) are mirrored from FilmCraft's first-party, spec-derived
|
||||
`filmcraft-av1` decoder (MIT OR Apache-2.0), which is also the round-trip oracle. No source of
|
||||
libaom, dav1d, rav1e, SVT-AV1, ffmpeg or any other third-party encoder or decoder was consulted;
|
||||
ffmpeg / libdav1d is used only as an external test oracle.
|
||||
|
|
@ -0,0 +1,87 @@
|
|||
//! MSB-first bit writer for the OBU headers (spec 4.10 `f(n)` descriptors) and `leb128()`.
|
||||
|
||||
#[derive(Default)]
|
||||
pub(crate) struct BitWriter {
|
||||
pub buf: Vec<u8>,
|
||||
nbits: usize,
|
||||
}
|
||||
|
||||
impl BitWriter {
|
||||
pub fn new() -> Self {
|
||||
BitWriter::default()
|
||||
}
|
||||
|
||||
/// `f(n)`: writes the low `n` bits of `v`, most significant first.
|
||||
pub fn f(&mut self, n: u32, v: u32) {
|
||||
for i in (0..n).rev() {
|
||||
self.bit((v >> i) & 1 != 0);
|
||||
}
|
||||
}
|
||||
|
||||
pub fn bit(&mut self, b: bool) {
|
||||
if self.nbits.is_multiple_of(8) {
|
||||
self.buf.push(0);
|
||||
}
|
||||
if b {
|
||||
let last = self.buf.len() - 1;
|
||||
self.buf[last] |= 0x80 >> (self.nbits % 8);
|
||||
}
|
||||
self.nbits += 1;
|
||||
}
|
||||
|
||||
/// `byte_alignment()`: zero bits up to the next byte boundary.
|
||||
pub fn byte_align(&mut self) {
|
||||
while !self.nbits.is_multiple_of(8) {
|
||||
self.bit(false);
|
||||
}
|
||||
}
|
||||
|
||||
/// `trailing_bits()`: a one bit followed by zero bits up to the byte boundary.
|
||||
pub fn trailing_bits(&mut self) {
|
||||
self.bit(true);
|
||||
self.byte_align();
|
||||
}
|
||||
|
||||
pub fn into_bytes(self) -> Vec<u8> {
|
||||
self.buf
|
||||
}
|
||||
}
|
||||
|
||||
/// `leb128()` encoding of an OBU size.
|
||||
pub(crate) fn leb128(mut v: usize, out: &mut Vec<u8>) {
|
||||
loop {
|
||||
let mut b = (v & 0x7f) as u8;
|
||||
v >>= 7;
|
||||
if v != 0 {
|
||||
b |= 0x80;
|
||||
}
|
||||
out.push(b);
|
||||
if v == 0 {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// An OBU with `obu_has_size_field` = 1 and no extension header.
|
||||
pub(crate) fn obu(obu_type: u8, payload: &[u8], out: &mut Vec<u8>) {
|
||||
out.push((obu_type << 3) | 0x02);
|
||||
leb128(payload.len(), out);
|
||||
out.extend_from_slice(payload);
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn bits_and_leb128() {
|
||||
let mut w = BitWriter::new();
|
||||
w.f(3, 0b101);
|
||||
w.f(6, 0b110011);
|
||||
w.trailing_bits();
|
||||
assert_eq!(w.into_bytes(), vec![0b1011_1001, 0b1100_0000]);
|
||||
let mut v = Vec::new();
|
||||
leb128(300, &mut v);
|
||||
assert_eq!(v, vec![0xAC, 0x02]);
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,97 @@
|
|||
//! The adaptive CDFs this encoder codes with (the subset of spec 6.8.2 init_non_coeff_cdfs /
|
||||
//! init_coeff_cdfs that its syntax uses). Layout as in filmcraft-av1: each innermost array
|
||||
//! holds N cumulative values (the last is 32768) plus the adaptation counter.
|
||||
|
||||
use crate::tables::*;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct Cdfs {
|
||||
pub y_mode: [[u16; 14]; 4],
|
||||
pub uv_mode_cfl_not_allowed: [[u16; 14]; 13],
|
||||
pub uv_mode_cfl_allowed: [[u16; 15]; 13],
|
||||
pub angle_delta: [[u16; 8]; 8],
|
||||
pub partition_w8: [[u16; 5]; 4],
|
||||
pub partition_w16: [[u16; 11]; 4],
|
||||
pub partition_w32: [[u16; 11]; 4],
|
||||
pub partition_w64: [[u16; 11]; 4],
|
||||
pub skip: [[u16; 3]; 3],
|
||||
pub is_inter: [[u16; 3]; 4],
|
||||
pub single_ref: [[[u16; 3]; 6]; 3],
|
||||
pub new_mv: [[u16; 3]; 6],
|
||||
pub zero_mv: [[u16; 3]; 2],
|
||||
pub ref_mv: [[u16; 3]; 6],
|
||||
pub drl_mode: [[u16; 3]; 3],
|
||||
pub mv_joint: [u16; 5],
|
||||
pub mv_class: [[u16; 12]; 2],
|
||||
pub mv_class0_bit: [[u16; 3]; 2],
|
||||
pub mv_fr: [[u16; 5]; 2],
|
||||
pub mv_class0_fr: [[[u16; 5]; 2]; 2],
|
||||
pub mv_sign: [[u16; 3]; 2],
|
||||
pub mv_bit: [[[u16; 3]; 10]; 2],
|
||||
pub intra_tx_type_set2: [[[u16; 6]; 13]; 3],
|
||||
pub inter_tx_type_set3: [[u16; 3]; 4],
|
||||
pub txb_skip: [[[u16; 3]; 13]; 5],
|
||||
pub eob_pt_16: [[[u16; 6]; 2]; 2],
|
||||
pub eob_pt_32: [[[u16; 7]; 2]; 2],
|
||||
pub eob_pt_64: [[[u16; 8]; 2]; 2],
|
||||
pub eob_pt_128: [[[u16; 9]; 2]; 2],
|
||||
pub eob_pt_256: [[[u16; 10]; 2]; 2],
|
||||
pub eob_pt_512: [[u16; 11]; 2],
|
||||
pub eob_pt_1024: [[u16; 12]; 2],
|
||||
pub eob_extra: [[[[u16; 3]; 9]; 2]; 5],
|
||||
pub dc_sign: [[[u16; 3]; 3]; 2],
|
||||
pub coeff_base_eob: [[[[u16; 4]; 4]; 2]; 5],
|
||||
pub coeff_base: [[[[u16; 5]; 42]; 2]; 5],
|
||||
pub coeff_br: [[[[u16; 5]; 21]; 2]; 5],
|
||||
}
|
||||
|
||||
impl Cdfs {
|
||||
/// init_non_coeff_cdfs( ) followed by init_coeff_cdfs( ) for `base_q_idx`.
|
||||
pub fn new(base_q_idx: u32) -> Box<Cdfs> {
|
||||
let idx = match base_q_idx {
|
||||
0..=20 => 0,
|
||||
21..=60 => 1,
|
||||
61..=120 => 2,
|
||||
_ => 3,
|
||||
};
|
||||
Box::new(Cdfs {
|
||||
y_mode: DEFAULT_Y_MODE_CDF,
|
||||
uv_mode_cfl_not_allowed: DEFAULT_UV_MODE_CFL_NOT_ALLOWED_CDF,
|
||||
uv_mode_cfl_allowed: DEFAULT_UV_MODE_CFL_ALLOWED_CDF,
|
||||
angle_delta: DEFAULT_ANGLE_DELTA_CDF,
|
||||
partition_w8: DEFAULT_PARTITION_W8_CDF,
|
||||
partition_w16: DEFAULT_PARTITION_W16_CDF,
|
||||
partition_w32: DEFAULT_PARTITION_W32_CDF,
|
||||
partition_w64: DEFAULT_PARTITION_W64_CDF,
|
||||
skip: DEFAULT_SKIP_CDF,
|
||||
is_inter: DEFAULT_IS_INTER_CDF,
|
||||
single_ref: DEFAULT_SINGLE_REF_CDF,
|
||||
new_mv: DEFAULT_NEW_MV_CDF,
|
||||
zero_mv: DEFAULT_ZERO_MV_CDF,
|
||||
ref_mv: DEFAULT_REF_MV_CDF,
|
||||
drl_mode: DEFAULT_DRL_MODE_CDF,
|
||||
mv_joint: DEFAULT_MV_JOINT_CDF,
|
||||
mv_class: DEFAULT_MV_CLASS_CDF,
|
||||
mv_class0_bit: [DEFAULT_MV_CLASS0_BIT_CDF; 2],
|
||||
mv_fr: DEFAULT_MV_FR_CDF,
|
||||
mv_class0_fr: DEFAULT_MV_CLASS0_FR_CDF,
|
||||
mv_sign: [DEFAULT_MV_SIGN_CDF; 2],
|
||||
mv_bit: [DEFAULT_MV_BIT_CDF; 2],
|
||||
intra_tx_type_set2: DEFAULT_INTRA_TX_TYPE_SET2_CDF,
|
||||
inter_tx_type_set3: DEFAULT_INTER_TX_TYPE_SET3_CDF,
|
||||
txb_skip: DEFAULT_TXB_SKIP_CDF[idx],
|
||||
eob_pt_16: DEFAULT_EOB_PT_16_CDF[idx],
|
||||
eob_pt_32: DEFAULT_EOB_PT_32_CDF[idx],
|
||||
eob_pt_64: DEFAULT_EOB_PT_64_CDF[idx],
|
||||
eob_pt_128: DEFAULT_EOB_PT_128_CDF[idx],
|
||||
eob_pt_256: DEFAULT_EOB_PT_256_CDF[idx],
|
||||
eob_pt_512: DEFAULT_EOB_PT_512_CDF[idx],
|
||||
eob_pt_1024: DEFAULT_EOB_PT_1024_CDF[idx],
|
||||
eob_extra: DEFAULT_EOB_EXTRA_CDF[idx],
|
||||
dc_sign: DEFAULT_DC_SIGN_CDF[idx],
|
||||
coeff_base_eob: DEFAULT_COEFF_BASE_EOB_CDF[idx],
|
||||
coeff_base: DEFAULT_COEFF_BASE_CDF[idx],
|
||||
coeff_br: DEFAULT_COEFF_BR_CDF[idx],
|
||||
})
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,677 @@
|
|||
//! Mode decision and reconstruction (not normative): partitions, intra modes, motion vectors,
|
||||
//! transform coefficients. Every prediction and reconstruction step uses the normative
|
||||
//! processes, so the reconstruction built here is exactly what a decoder produces; the syntax
|
||||
//! is written afterwards by `writer.rs`.
|
||||
|
||||
use crate::pred::*;
|
||||
use crate::tables::*;
|
||||
use crate::transform::{forward_2d, inverse_2d};
|
||||
|
||||
/// Decisions for one coded block.
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct Leaf {
|
||||
pub bsize: usize,
|
||||
pub is_inter: bool,
|
||||
pub y_mode: usize,
|
||||
pub uv_mode: usize,
|
||||
pub tx_type_y: usize,
|
||||
/// Motion vector [row, col] in 1/8 luma samples (inter blocks).
|
||||
pub mv: [i32; 2],
|
||||
pub skip: bool,
|
||||
/// Quantised coefficients of the single transform block of each plane (raster order).
|
||||
pub coefs: [Vec<i32>; 3],
|
||||
}
|
||||
|
||||
pub(crate) enum Node {
|
||||
Out,
|
||||
Leaf(Box<Leaf>),
|
||||
Split(Box<[Node; 4]>),
|
||||
}
|
||||
|
||||
/// A reconstructed frame kept as a reference.
|
||||
pub(crate) struct RefFrame {
|
||||
pub planes: [Plane; 3],
|
||||
/// Motion vectors chosen for the frame (per 4x4), used as search candidates.
|
||||
pub mvs: Vec<[i32; 2]>,
|
||||
}
|
||||
|
||||
pub(crate) struct FrameGeom {
|
||||
pub width: usize,
|
||||
pub height: usize,
|
||||
pub mi_cols: usize,
|
||||
pub mi_rows: usize,
|
||||
pub sb_cols: usize,
|
||||
pub sb_rows: usize,
|
||||
}
|
||||
|
||||
impl FrameGeom {
|
||||
pub fn new(width: usize, height: usize) -> FrameGeom {
|
||||
let mi_cols = 2 * width.div_ceil(8);
|
||||
let mi_rows = 2 * height.div_ceil(8);
|
||||
FrameGeom { width, height, mi_cols, mi_rows, sb_cols: mi_cols.div_ceil(16), sb_rows: mi_rows.div_ceil(16) }
|
||||
}
|
||||
/// Padded plane size (whole superblocks).
|
||||
pub fn padded(&self, plane: usize) -> (usize, usize) {
|
||||
let s = (plane > 0) as usize;
|
||||
((self.sb_cols * 64) >> s, (self.sb_rows * 64) >> s)
|
||||
}
|
||||
/// Visible plane size.
|
||||
pub fn visible(&self, plane: usize) -> (usize, usize) {
|
||||
if plane == 0 { (self.width, self.height) } else { (self.width.div_ceil(2), self.height.div_ceil(2)) }
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn tx_size_for(log2: u32) -> usize {
|
||||
match log2 {
|
||||
2 => TX_4X4,
|
||||
3 => TX_8X8,
|
||||
4 => TX_16X16,
|
||||
5 => TX_32X32,
|
||||
_ => TX_64X64,
|
||||
}
|
||||
}
|
||||
|
||||
/// get_tx_set( ) for the square sizes used here with reduced_tx_set = 1.
|
||||
pub(crate) fn tx_set(tx_sz: usize, is_inter: bool) -> usize {
|
||||
let sqr_up = TX_SIZE_SQR_UP[tx_sz] as usize;
|
||||
if sqr_up > TX_32X32 {
|
||||
return TX_SET_DCTONLY;
|
||||
}
|
||||
if is_inter {
|
||||
TX_SET_INTER_3
|
||||
} else if sqr_up == TX_32X32 {
|
||||
TX_SET_DCTONLY
|
||||
} else {
|
||||
TX_SET_INTRA_2
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn tx_in_set(set: usize, is_inter: bool, tx_type: usize) -> bool {
|
||||
if is_inter { TX_TYPE_IN_SET_INTER[set][tx_type] == 1 } else { TX_TYPE_IN_SET_INTRA[set][tx_type] == 1 }
|
||||
}
|
||||
|
||||
/// The chroma transform type of an intra block (compute_tx_type for plane > 0).
|
||||
pub(crate) fn intra_uv_tx_type(uv_mode: usize, uv_tx: usize) -> usize {
|
||||
let t = MODE_TO_TXFM[uv_mode] as usize;
|
||||
if tx_in_set(tx_set(uv_tx, false), false, t) { t } else { DCT_DCT }
|
||||
}
|
||||
|
||||
struct TxResult {
|
||||
levels: Vec<i32>,
|
||||
nonzero: bool,
|
||||
dist: f64,
|
||||
bits: f64,
|
||||
}
|
||||
|
||||
/// Saved encoder state of a block area (reconstruction, BlockDecoded flags).
|
||||
struct Snapshot {
|
||||
planes: [Vec<u16>; 3],
|
||||
decoded: Box<[[[bool; 35]; 35]; 3]>,
|
||||
}
|
||||
|
||||
pub(crate) struct Decider<'a> {
|
||||
pub g: &'a FrameGeom,
|
||||
pub bd: u32,
|
||||
pub intra_only: bool,
|
||||
pub src: &'a [Plane; 3],
|
||||
pub rec: [Plane; 3],
|
||||
pub refr: Option<&'a RefFrame>,
|
||||
decoded: Box<[[[bool; 35]; 35]; 3]>,
|
||||
/// Decision motion vectors per 4x4 (search candidates).
|
||||
pub mvs: Vec<[i32; 2]>,
|
||||
dc_q: f32,
|
||||
ac_q: f32,
|
||||
dc_qi: i64,
|
||||
ac_qi: i64,
|
||||
lambda: f64,
|
||||
lambda_sad: f64,
|
||||
sb_r: usize,
|
||||
sb_c: usize,
|
||||
}
|
||||
|
||||
const MODE_BITS_INTRA: f64 = 3.0;
|
||||
|
||||
impl<'a> Decider<'a> {
|
||||
pub fn new(g: &'a FrameGeom, bd: u32, qidx: u8, intra_only: bool, src: &'a [Plane; 3], refr: Option<&'a RefFrame>) -> Self {
|
||||
let bdi = ((bd - 8) >> 1) as usize;
|
||||
let dc_qi = DC_QLOOKUP[bdi][qidx as usize] as i64;
|
||||
let ac_qi = AC_QLOOKUP[bdi][qidx as usize] as i64;
|
||||
// Lagrangian multiplier in squared-sample units of the coded bit depth.
|
||||
let step = ac_qi as f64 / 8.0;
|
||||
let lambda = 0.22 * step * step;
|
||||
let rec = std::array::from_fn(|p| {
|
||||
let (w, h) = g.padded(p);
|
||||
Plane::new(w, h)
|
||||
});
|
||||
Decider {
|
||||
g,
|
||||
bd,
|
||||
intra_only,
|
||||
src,
|
||||
rec,
|
||||
refr,
|
||||
decoded: Box::new([[[false; 35]; 35]; 3]),
|
||||
mvs: vec![[0, 0]; g.mi_cols * g.mi_rows],
|
||||
dc_q: dc_qi as f32,
|
||||
ac_q: ac_qi as f32,
|
||||
dc_qi,
|
||||
ac_qi,
|
||||
lambda,
|
||||
lambda_sad: lambda.sqrt(),
|
||||
sb_r: 0,
|
||||
sb_c: 0,
|
||||
}
|
||||
}
|
||||
|
||||
/// Decides one 64x64 superblock (reconstruction is written to `rec`).
|
||||
pub fn superblock(&mut self, r: usize, c: usize) -> Node {
|
||||
self.sb_r = r;
|
||||
self.sb_c = c;
|
||||
// clear_block_decoded_flags( r, c, sbSize4 )
|
||||
for plane in 0..3 {
|
||||
let s = (plane > 0) as usize;
|
||||
let sb_w4 = (self.g.mi_cols as isize - c as isize) >> s;
|
||||
let sb_h4 = (self.g.mi_rows as isize - r as isize) >> s;
|
||||
let bd = &mut self.decoded[plane];
|
||||
for y in -1..=((16 >> s) as isize) {
|
||||
for x in -1..=((16 >> s) as isize) {
|
||||
let v = if y < 0 && x < sb_w4 { true } else { x < 0 && y < sb_h4 };
|
||||
bd[(y + 1) as usize][(x + 1) as usize] = v;
|
||||
}
|
||||
}
|
||||
bd[(16 >> s) + 1][0] = false;
|
||||
}
|
||||
self.partition(r, c, BLOCK_64X64, None).0
|
||||
}
|
||||
|
||||
fn decoded_at(&self, plane: usize, y: isize, x: isize) -> bool {
|
||||
let (yy, xx) = (y + 1, x + 1);
|
||||
if yy < 0 || xx < 0 || yy > 34 || xx > 34 {
|
||||
return false;
|
||||
}
|
||||
self.decoded[plane][yy as usize][xx as usize]
|
||||
}
|
||||
|
||||
/// Marks the transform block of `plane` covering the block at (r, c) as decoded.
|
||||
fn mark_decoded(&mut self, plane: usize, r: usize, c: usize, n4: usize) {
|
||||
let s = (plane > 0) as usize;
|
||||
let by = (r & 15) >> s;
|
||||
let bx = (c & 15) >> s;
|
||||
let step = (n4 >> s).max(1);
|
||||
for i in 0..step {
|
||||
for j in 0..step {
|
||||
if by + i + 1 < 35 && bx + j + 1 < 35 {
|
||||
self.decoded[plane][by + i + 1][bx + j + 1] = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn snapshot(&self, r: usize, c: usize, n4: usize) -> Snapshot {
|
||||
let planes = std::array::from_fn(|p| {
|
||||
let s = (p > 0) as usize;
|
||||
let (x0, y0, n) = ((c * 4) >> s, (r * 4) >> s, (n4 * 4) >> s);
|
||||
let pl = &self.rec[p];
|
||||
let mut v = Vec::with_capacity(n * n);
|
||||
for y in y0..y0 + n {
|
||||
v.extend_from_slice(&pl.data[y * pl.stride + x0..y * pl.stride + x0 + n]);
|
||||
}
|
||||
v
|
||||
});
|
||||
Snapshot { planes, decoded: self.decoded.clone() }
|
||||
}
|
||||
|
||||
fn restore(&mut self, snap: &Snapshot, r: usize, c: usize, n4: usize) {
|
||||
for p in 0..3 {
|
||||
let s = (p > 0) as usize;
|
||||
let (x0, y0, n) = ((c * 4) >> s, (r * 4) >> s, (n4 * 4) >> s);
|
||||
let pl = &mut self.rec[p];
|
||||
for (i, y) in (y0..y0 + n).enumerate() {
|
||||
pl.data[y * pl.stride + x0..y * pl.stride + x0 + n].copy_from_slice(&snap.planes[p][i * n..i * n + n]);
|
||||
}
|
||||
}
|
||||
self.decoded.clone_from(&snap.decoded);
|
||||
}
|
||||
|
||||
fn set_mvs(&mut self, r: usize, c: usize, n4: usize, mv: [i32; 2]) {
|
||||
for y in r..(r + n4).min(self.g.mi_rows) {
|
||||
for x in c..(c + n4).min(self.g.mi_cols) {
|
||||
self.mvs[y * self.g.mi_cols + x] = mv;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn partition(&mut self, r: usize, c: usize, bsize: usize, parent_mv: Option<[i32; 2]>) -> (Node, f64) {
|
||||
let g = self.g;
|
||||
if r >= g.mi_rows || c >= g.mi_cols {
|
||||
return (Node::Out, 0.0);
|
||||
}
|
||||
let n4 = NUM_4X4_BLOCKS_WIDE[bsize] as usize;
|
||||
let half = n4 >> 1;
|
||||
let has_rows = r + half < g.mi_rows;
|
||||
let has_cols = c + half < g.mi_cols;
|
||||
let can_none = has_rows && has_cols && (bsize != BLOCK_64X64 || !self.intra_only);
|
||||
let can_split = bsize > BLOCK_8X8;
|
||||
let sub = PARTITION_SUBSIZE[PARTITION_SPLIT][bsize] as usize;
|
||||
if !can_none {
|
||||
return self.split(r, c, sub, half, parent_mv);
|
||||
}
|
||||
let before = self.snapshot(r, c, n4);
|
||||
let (leaf, cost_none) = self.leaf(r, c, bsize, parent_mv);
|
||||
let cost_none = cost_none + self.lambda * 1.0;
|
||||
if !can_split {
|
||||
return (Node::Leaf(Box::new(leaf)), cost_none);
|
||||
}
|
||||
// a cheap skip block rarely gains from splitting
|
||||
if leaf.skip && cost_none < self.lambda * 8.0 {
|
||||
return (Node::Leaf(Box::new(leaf)), cost_none);
|
||||
}
|
||||
let after_none = self.snapshot(r, c, n4);
|
||||
self.restore(&before, r, c, n4);
|
||||
let pm = if leaf.is_inter { Some(leaf.mv) } else { parent_mv };
|
||||
let (node, cost_split) = self.split(r, c, sub, half, pm);
|
||||
let cost_split = cost_split + self.lambda * 2.0;
|
||||
if cost_split < cost_none {
|
||||
(node, cost_split)
|
||||
} else {
|
||||
self.restore(&after_none, r, c, n4);
|
||||
self.set_mvs(r, c, n4, if leaf.is_inter { leaf.mv } else { [0, 0] });
|
||||
(Node::Leaf(Box::new(leaf)), cost_none)
|
||||
}
|
||||
}
|
||||
|
||||
fn split(&mut self, r: usize, c: usize, sub: usize, half: usize, pm: Option<[i32; 2]>) -> (Node, f64) {
|
||||
let (n0, c0) = self.partition(r, c, sub, pm);
|
||||
let (n1, c1) = self.partition(r, c + half, sub, pm);
|
||||
let (n2, c2) = self.partition(r + half, c, sub, pm);
|
||||
let (n3, c3) = self.partition(r + half, c + half, sub, pm);
|
||||
(Node::Split(Box::new([n0, n1, n2, n3])), c0 + c1 + c2 + c3)
|
||||
}
|
||||
|
||||
/// Best coding of the block as a single (NONE partition) block.
|
||||
fn leaf(&mut self, r: usize, c: usize, bsize: usize, parent_mv: Option<[i32; 2]>) -> (Leaf, f64) {
|
||||
let n4 = NUM_4X4_BLOCKS_WIDE[bsize] as usize;
|
||||
if self.intra_only {
|
||||
let res = self.intra_leaf(r, c, bsize);
|
||||
self.set_mvs(r, c, n4, [0, 0]);
|
||||
return res;
|
||||
}
|
||||
let before = self.snapshot(r, c, n4);
|
||||
let (inter, cost_inter) = self.inter_leaf(r, c, bsize, parent_mv);
|
||||
if bsize == BLOCK_64X64 || inter.skip {
|
||||
self.set_mvs(r, c, n4, inter.mv);
|
||||
return (inter, cost_inter);
|
||||
}
|
||||
let after_inter = self.snapshot(r, c, n4);
|
||||
self.restore(&before, r, c, n4);
|
||||
let (intra, cost_intra) = self.intra_leaf(r, c, bsize);
|
||||
if cost_intra < cost_inter {
|
||||
self.set_mvs(r, c, n4, [0, 0]);
|
||||
(intra, cost_intra)
|
||||
} else {
|
||||
self.restore(&after_inter, r, c, n4);
|
||||
self.set_mvs(r, c, n4, inter.mv);
|
||||
(inter, cost_inter)
|
||||
}
|
||||
}
|
||||
|
||||
fn intra_leaf(&mut self, r: usize, c: usize, bsize: usize) -> (Leaf, f64) {
|
||||
let n4 = NUM_4X4_BLOCKS_WIDE[bsize] as usize;
|
||||
let log2 = (n4 * 4).trailing_zeros();
|
||||
let (sbr, sbc) = ((r & 15) as isize, (c & 15) as isize);
|
||||
let bd = self.bd;
|
||||
let mut pred = vec![0u16; 32 * 32];
|
||||
let mut best_pred = vec![0u16; 32 * 32];
|
||||
let n = 1usize << log2;
|
||||
// luma
|
||||
let e = IntraEdge {
|
||||
x: c * 4,
|
||||
y: r * 4,
|
||||
have_left: c > 0,
|
||||
have_above: r > 0,
|
||||
have_above_right: self.decoded_at(0, sbr - 1, sbc + n4 as isize),
|
||||
have_below_left: self.decoded_at(0, sbr + n4 as isize, sbc - 1),
|
||||
log2,
|
||||
max_x: (self.g.mi_cols * 4) as i32 - 1,
|
||||
max_y: (self.g.mi_rows * 4) as i32 - 1,
|
||||
bit_depth: bd,
|
||||
};
|
||||
let ed = edges(&self.rec[0], &e);
|
||||
let mut best = (f64::MAX, DC_PRED);
|
||||
for &mode in INTRA_MODES_USED.iter() {
|
||||
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 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);
|
||||
self.mark_decoded(0, r, c, n4);
|
||||
// chroma (one mode for both planes)
|
||||
let cn = n >> 1;
|
||||
let clog2 = log2 - 1;
|
||||
let (csbr, csbc) = (sbr >> 1, sbc >> 1);
|
||||
let ce: [IntraEdge; 2] = std::array::from_fn(|i| IntraEdge {
|
||||
x: c * 2,
|
||||
y: r * 2,
|
||||
have_left: c > 0,
|
||||
have_above: r > 0,
|
||||
have_above_right: self.decoded_at(1 + i, csbr - 1, csbc + (n4 >> 1) as isize),
|
||||
have_below_left: self.decoded_at(1 + i, csbr + (n4 >> 1) as isize, csbc - 1),
|
||||
log2: clog2,
|
||||
max_x: (self.g.mi_cols * 2) as i32 - 1,
|
||||
max_y: (self.g.mi_rows * 2) as i32 - 1,
|
||||
bit_depth: bd,
|
||||
});
|
||||
let ced = [edges(&self.rec[1], &ce[0]), edges(&self.rec[2], &ce[1])];
|
||||
let mut best_uv = (f64::MAX, DC_PRED);
|
||||
for &mode in INTRA_MODES_USED.iter() {
|
||||
let mut sad = 0u64;
|
||||
for i in 0..2 {
|
||||
predict_intra(&ced[i], &ce[i], mode, &mut pred);
|
||||
sad += self.sad(1 + i, c * 2, r * 2, cn, &pred[..cn * cn]);
|
||||
}
|
||||
let cost = sad as f64 + self.lambda_sad * if mode == DC_PRED { 1.0 } else { 2.5 };
|
||||
if cost < best_uv.0 {
|
||||
best_uv = (cost, mode);
|
||||
}
|
||||
}
|
||||
let uv_mode = best_uv.1;
|
||||
let uv_tx = tx_size_for(clog2);
|
||||
let uv_type = intra_uv_tx_type(uv_mode, uv_tx);
|
||||
let mut tuv = Vec::with_capacity(2);
|
||||
for i in 0..2 {
|
||||
predict_intra(&ced[i], &ce[i], uv_mode, &mut pred);
|
||||
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] };
|
||||
(leaf, dist + self.lambda * bits)
|
||||
}
|
||||
|
||||
fn inter_leaf(&mut self, r: usize, c: usize, bsize: usize, parent_mv: Option<[i32; 2]>) -> (Leaf, f64) {
|
||||
let n4 = NUM_4X4_BLOCKS_WIDE[bsize] as usize;
|
||||
let n = n4 * 4;
|
||||
let log2 = n.trailing_zeros();
|
||||
let (mv, mv_bits) = self.motion_search(r, c, n, parent_mv);
|
||||
let refr = self.refr.expect("reference frame");
|
||||
let mut pred = vec![0u16; n * n];
|
||||
let mut tx = Vec::with_capacity(3);
|
||||
let mut dist = 0.0;
|
||||
for p in 0..3 {
|
||||
let s = (p > 0) as usize;
|
||||
let (vw, vh) = self.g.visible(p);
|
||||
let pn = n >> s;
|
||||
let (x, y) = ((c * 4) >> s, (r * 4) >> s);
|
||||
predict_inter(&refr.planes[p], vw as i32 - 1, vh as i32 - 1, x, y, pn, pn, mv, s, self.bd, &mut pred[..pn * pn]);
|
||||
if bsize == BLOCK_64X64 {
|
||||
// coded as a skip block: the prediction is the reconstruction
|
||||
dist += self.sse(p, x, y, pn, &pred[..pn * pn]);
|
||||
self.put(p, x, y, pn, &pred[..pn * pn]);
|
||||
} else {
|
||||
let t = self.code_tx(p, x, y, log2 - s as u32, DCT_DCT, &pred[..pn * pn], false);
|
||||
dist += t.dist;
|
||||
tx.push(t);
|
||||
}
|
||||
self.mark_decoded(p, r, c, n4);
|
||||
}
|
||||
let mode_bits = if mv == [0, 0] { 1.5 } else { 2.0 + mv_bits };
|
||||
if bsize == BLOCK_64X64 {
|
||||
let leaf = Leaf { bsize, is_inter: true, y_mode: NEWMV, uv_mode: DC_PRED, tx_type_y: DCT_DCT, mv, skip: true, coefs: Default::default() };
|
||||
return (leaf, dist + self.lambda * (mode_bits + 1.0));
|
||||
}
|
||||
let skip = tx.iter().all(|t| !t.nonzero);
|
||||
let bits = mode_bits + if skip { 1.0 } else { 1.0 + tx.iter().map(|t| t.bits).sum::<f64>() };
|
||||
let [y, u, v]: [TxResult; 3] = tx.try_into().ok().expect("three planes");
|
||||
let leaf = Leaf { bsize, is_inter: true, y_mode: NEWMV, uv_mode: DC_PRED, tx_type_y: DCT_DCT, mv, skip, coefs: [y.levels, u.levels, v.levels] };
|
||||
(leaf, dist + self.lambda * bits)
|
||||
}
|
||||
|
||||
/// Full-sample motion search: candidate vectors then a shrinking diamond search.
|
||||
/// Returns the vector and an estimate of its coding cost in bits.
|
||||
fn motion_search(&self, r: usize, c: usize, n: usize, parent_mv: Option<[i32; 2]>) -> ([i32; 2], f64) {
|
||||
let g = self.g;
|
||||
let refr = self.refr.expect("reference frame");
|
||||
let mi = |rr: usize, cc: usize| self.mvs[rr * g.mi_cols + cc];
|
||||
let mut pred_mv = [0, 0];
|
||||
let mut cands: Vec<[i32; 2]> = vec![[0, 0]];
|
||||
if c > 0 {
|
||||
pred_mv = mi(r, c - 1);
|
||||
cands.push(pred_mv);
|
||||
}
|
||||
if r > 0 {
|
||||
cands.push(mi(r - 1, c));
|
||||
let cr = c + n / 4;
|
||||
if cr < g.mi_cols {
|
||||
cands.push(mi(r - 1, cr));
|
||||
}
|
||||
}
|
||||
if let Some(p) = parent_mv {
|
||||
cands.push(p);
|
||||
}
|
||||
cands.push(refr.mvs[r * g.mi_cols + c]);
|
||||
let (x0, y0) = ((c * 4) as i32, (r * 4) as i32);
|
||||
let (w, h) = (g.width as i32, g.height as i32);
|
||||
// keep the referenced block within 64 samples of the picture
|
||||
let lim_x = (-x0 - n as i32 - 64, w - x0 + 64);
|
||||
let lim_y = (-y0 - n as i32 - 64, h - y0 + 64);
|
||||
let valid = |m: [i32; 2]| {
|
||||
let (dy, dx) = (m[0] >> 3, m[1] >> 3);
|
||||
dx >= lim_x.0 && dx <= lim_x.1 && dy >= lim_y.0 && dy <= lim_y.1
|
||||
};
|
||||
let mv_bits = |m: [i32; 2]| -> f64 {
|
||||
let comp = |d: i32| if d == 0 { 0.5 } else { 2.0 + 2.0 * ((d.abs() / 8) as f64 + 1.0).log2() };
|
||||
comp(m[0] - pred_mv[0]) + comp(m[1] - pred_mv[1])
|
||||
};
|
||||
let cost =
|
||||
|m: [i32; 2]| -> f64 { self.sad_ref(&refr.planes[0], x0 + (m[1] >> 3), y0 + (m[0] >> 3), c * 4, r * 4, n) as f64 + self.lambda_sad * mv_bits(m) };
|
||||
let mut best = ([0, 0], f64::MAX);
|
||||
for m in cands {
|
||||
let m = [m[0] & !7, m[1] & !7];
|
||||
if !valid(m) {
|
||||
continue;
|
||||
}
|
||||
let cst = cost(m);
|
||||
if cst < best.1 {
|
||||
best = (m, cst);
|
||||
}
|
||||
}
|
||||
let mut step = 16 * 8;
|
||||
while step >= 8 {
|
||||
let mut moved = true;
|
||||
let mut iters = 0;
|
||||
while moved && iters < 16 {
|
||||
moved = false;
|
||||
iters += 1;
|
||||
let center = best.0;
|
||||
for d in [[-step, 0], [step, 0], [0, -step], [0, step]] {
|
||||
let m = [center[0] + d[0], center[1] + d[1]];
|
||||
if !valid(m) {
|
||||
continue;
|
||||
}
|
||||
let cst = cost(m);
|
||||
if cst < best.1 {
|
||||
best = (m, cst);
|
||||
moved = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
step >>= 1;
|
||||
}
|
||||
(best.0, mv_bits(best.0))
|
||||
}
|
||||
|
||||
/// SAD of the source block at (x, y) (luma) against the reference block at (rx, ry).
|
||||
fn sad_ref(&self, refp: &Plane, rx: i32, ry: i32, x: usize, y: usize, n: usize) -> u64 {
|
||||
let (vw, vh) = (self.g.width, self.g.height);
|
||||
let w = n.min(vw.saturating_sub(x));
|
||||
let h = n.min(vh.saturating_sub(y));
|
||||
let src = &self.src[0];
|
||||
let (lx, ly) = (vw as i32 - 1, vh as i32 - 1);
|
||||
let mut sad = 0u64;
|
||||
let inside = rx >= 0 && ry >= 0 && rx + w as i32 - 1 <= lx && ry + h as i32 - 1 <= ly;
|
||||
for i in 0..h {
|
||||
let s = &src.row(y + i)[x..x + w];
|
||||
let rr = refp.row((ry + i as i32).clamp(0, ly) as usize);
|
||||
if inside {
|
||||
let rrow = &rr[rx as usize..rx as usize + w];
|
||||
sad += s.iter().zip(rrow).map(|(&a, &b)| (a as i32 - b as i32).unsigned_abs() as u64).sum::<u64>();
|
||||
} else {
|
||||
for j in 0..w {
|
||||
let b = rr[(rx + j as i32).clamp(0, lx) as usize];
|
||||
sad += (s[j] as i32 - b as i32).unsigned_abs() as u64;
|
||||
}
|
||||
}
|
||||
}
|
||||
sad
|
||||
}
|
||||
|
||||
/// SAD of a prediction against the source over the visible part of the block.
|
||||
fn sad(&self, p: usize, x: usize, y: usize, n: usize, pred: &[u16]) -> u64 {
|
||||
let (vw, vh) = self.g.visible(p);
|
||||
let w = n.min(vw.saturating_sub(x));
|
||||
let h = n.min(vh.saturating_sub(y));
|
||||
let src = &self.src[p];
|
||||
let mut s = 0u64;
|
||||
for i in 0..h {
|
||||
let row = &src.row(y + i)[x..x + w];
|
||||
s += row.iter().zip(&pred[i * n..i * n + w]).map(|(&a, &b)| (a as i32 - b as i32).unsigned_abs() as u64).sum::<u64>();
|
||||
}
|
||||
s
|
||||
}
|
||||
|
||||
fn sse(&self, p: usize, x: usize, y: usize, n: usize, blk: &[u16]) -> f64 {
|
||||
let (vw, vh) = self.g.visible(p);
|
||||
let w = n.min(vw.saturating_sub(x));
|
||||
let h = n.min(vh.saturating_sub(y));
|
||||
let src = &self.src[p];
|
||||
let mut s = 0u64;
|
||||
for i in 0..h {
|
||||
let row = &src.row(y + i)[x..x + w];
|
||||
s += row
|
||||
.iter()
|
||||
.zip(&blk[i * n..i * n + w])
|
||||
.map(|(&a, &b)| {
|
||||
let d = a as i64 - b as i64;
|
||||
(d * d) as u64
|
||||
})
|
||||
.sum::<u64>();
|
||||
}
|
||||
s as f64
|
||||
}
|
||||
|
||||
fn put(&mut self, p: usize, x: usize, y: usize, n: usize, blk: &[u16]) {
|
||||
let pl = &mut self.rec[p];
|
||||
for i in 0..n {
|
||||
let o = (y + i) * pl.stride + x;
|
||||
pl.data[o..o + n].copy_from_slice(&blk[i * n..i * n + n]);
|
||||
}
|
||||
}
|
||||
|
||||
/// 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 n = 1usize << log2;
|
||||
let tx_sz = tx_size_for(log2);
|
||||
let bd = self.bd;
|
||||
let src = &self.src[p];
|
||||
let mut res = vec![0i32; n * n];
|
||||
for i in 0..n {
|
||||
let row = &src.row(y + i)[x..x + n];
|
||||
for j in 0..n {
|
||||
res[i * n + j] = row[j] as i32 - pred[i * n + j] as i32;
|
||||
}
|
||||
}
|
||||
let mut coef = vec![0f32; n * n];
|
||||
forward_2d(&res, &mut coef, n, tx_type);
|
||||
// dead-zone quantiser: rounding offset below one half
|
||||
let rnd = if intra { 0.40f32 } else { 0.33f32 };
|
||||
let mut levels = vec![0i32; n * n];
|
||||
let mut nonzero = false;
|
||||
for (k, (&cf, l)) in coef.iter().zip(levels.iter_mut()).enumerate() {
|
||||
let q = if k == 0 { self.dc_q } else { self.ac_q };
|
||||
let a = (cf.abs() / q + rnd).floor() as i32;
|
||||
if a != 0 {
|
||||
*l = if cf < 0.0 { -a.min(1 << 16) } else { a.min(1 << 16) };
|
||||
nonzero = true;
|
||||
}
|
||||
}
|
||||
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 };
|
||||
}
|
||||
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 };
|
||||
}
|
||||
self.put(p, x, y, n, &recon);
|
||||
TxResult { levels, nonzero: true, dist, bits }
|
||||
}
|
||||
|
||||
/// Dequantisation (7.12.3) and the inverse transform, added to the prediction.
|
||||
fn reconstruct(&self, levels: &[i32], tx_sz: usize, tx_type: usize, pred: &[u16], n: usize, bd: u32) -> Vec<u16> {
|
||||
let denom: i64 = if tx_sz == TX_32X32 { 2 } else { 1 };
|
||||
let lim = 1i64 << (7 + bd);
|
||||
let mut deq = vec![0i32; n * n];
|
||||
for (k, (&l, d)) in levels.iter().zip(deq.iter_mut()).enumerate() {
|
||||
if l == 0 {
|
||||
continue;
|
||||
}
|
||||
let q = if k == 0 { self.dc_qi } else { self.ac_qi };
|
||||
let dq = l as i64 * q;
|
||||
let mag = (dq.abs() & 0xFF_FFFF) / denom;
|
||||
*d = (if dq < 0 { -mag } else { mag }).clamp(-lim, lim - 1) as i32;
|
||||
}
|
||||
let mut resid = vec![0i32; n * n];
|
||||
inverse_2d(&deq, &mut resid, tx_sz, tx_type, bd);
|
||||
let max = (1i32 << bd) - 1;
|
||||
pred.iter().zip(&resid).map(|(&p, &r)| (p as i32 + r).clamp(0, max) as u16).collect()
|
||||
}
|
||||
}
|
||||
|
||||
/// Rough coefficient coding cost in bits.
|
||||
fn estimate_bits(levels: &[i32], tx_sz: usize) -> f64 {
|
||||
let scan = scan_for(tx_sz);
|
||||
let mut eob = 0;
|
||||
for (i, &pos) in scan.iter().enumerate() {
|
||||
if levels[pos as usize] != 0 {
|
||||
eob = i + 1;
|
||||
}
|
||||
}
|
||||
if eob == 0 {
|
||||
return 1.0;
|
||||
}
|
||||
let mut bits = 3.0 + (eob as f64).log2();
|
||||
for &pos in &scan[..eob] {
|
||||
let a = levels[pos as usize].unsigned_abs();
|
||||
bits += match a {
|
||||
0 => 0.7,
|
||||
1 => 2.5,
|
||||
_ => 3.5 + 1.6 * (a as f64).log2(),
|
||||
};
|
||||
}
|
||||
bits
|
||||
}
|
||||
|
||||
pub(crate) fn scan_for(tx_sz: usize) -> &'static [u16] {
|
||||
match tx_sz {
|
||||
TX_4X4 => &DEFAULT_SCAN_4X4,
|
||||
TX_8X8 => &DEFAULT_SCAN_8X8,
|
||||
TX_16X16 => &DEFAULT_SCAN_16X16,
|
||||
_ => &DEFAULT_SCAN_32X32,
|
||||
}
|
||||
}
|
||||
|
|
@ -1,7 +1,27 @@
|
|||
//! Clean-room pure-Rust AV1 encoder — API skeleton (implementation in progress).
|
||||
//! 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.
|
||||
|
||||
mod bitw;
|
||||
mod cdf;
|
||||
mod decide;
|
||||
mod pred;
|
||||
mod symbol;
|
||||
mod tables;
|
||||
mod transform;
|
||||
mod writer;
|
||||
|
||||
use bitw::{BitWriter, obu};
|
||||
use decide::{Decider, FrameGeom, RefFrame};
|
||||
use pred::Plane;
|
||||
use writer::TileWriter;
|
||||
|
||||
/// Rate control.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum RateControl {
|
||||
|
|
@ -68,9 +88,63 @@ impl std::fmt::Display for Error {
|
|||
|
||||
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,
|
||||
/// Accumulated (spent - budget) bits.
|
||||
debt: f64,
|
||||
}
|
||||
|
||||
/// 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 {
|
||||
|
|
@ -78,21 +152,302 @@ impl Encoder {
|
|||
if cfg.width == 0 || cfg.height == 0 {
|
||||
return Err(Error::InvalidConfig("empty frame".into()));
|
||||
}
|
||||
Ok(Encoder { cfg })
|
||||
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, 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> {
|
||||
vec![]
|
||||
self.seq_obu.clone()
|
||||
}
|
||||
|
||||
/// AV1CodecConfigurationRecord (`av1C` payload, also the WebM CodecPrivate).
|
||||
pub fn av1c(&self) -> Vec<u8> {
|
||||
vec![]
|
||||
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, 1); // color_primaries: BT.709
|
||||
w.f(8, 1); // transfer_characteristics: BT.709
|
||||
w.f(8, 1); // matrix_coefficients: BT.709
|
||||
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) {
|
||||
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, 7); // primary_ref_frame = PRIMARY_REF_NONE
|
||||
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, 1); // 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( )
|
||||
w.f(6, 0);
|
||||
w.f(6, 0);
|
||||
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);
|
||||
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 tile = tw.finish();
|
||||
let mut w = BitWriter::new();
|
||||
self.frame_header(&mut w, key, qidx);
|
||||
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 };
|
||||
(out, rf)
|
||||
}
|
||||
|
||||
/// Encodes one frame.
|
||||
pub fn encode(&mut self, _frame: &Frame) -> Packet {
|
||||
Packet { data: vec![], keyframe: true }
|
||||
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 { .. } => self.encode_rate_controlled(&src, key),
|
||||
};
|
||||
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) -> (Vec<u8>, RefFrame) {
|
||||
let first = self.frame_num == 0;
|
||||
let rc = self.rc.as_ref().expect("rate control state");
|
||||
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 qi = q.round().clamp(1.0, 255.0) as u8;
|
||||
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));
|
||||
}
|
||||
if pass + 1 < passes {
|
||||
q = (q + 30.0 * err).clamp(1.0, 255.0);
|
||||
}
|
||||
}
|
||||
let (obu, rf, bits) = best.expect("encoded");
|
||||
let rc = self.rc.as_mut().expect("rate control state");
|
||||
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);
|
||||
}
|
||||
(obu, rf)
|
||||
}
|
||||
|
||||
/// 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
|
||||
}))
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,262 @@
|
|||
//! Prediction processes, mirrored from the spec (7.11.2 intra prediction for the non-directional
|
||||
//! modes plus V_PRED / H_PRED without edge filtering, 7.11.3 block inter prediction without
|
||||
//! reference scaling) so that the encoder predicts exactly like a decoder.
|
||||
|
||||
use crate::tables::*;
|
||||
use crate::transform::round2;
|
||||
|
||||
/// A plane buffer (padded to whole superblocks).
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct Plane {
|
||||
pub data: Vec<u16>,
|
||||
pub stride: usize,
|
||||
}
|
||||
|
||||
impl Plane {
|
||||
pub fn new(w: usize, h: usize) -> Plane {
|
||||
Plane { data: vec![0; w * h], stride: w }
|
||||
}
|
||||
#[inline(always)]
|
||||
pub fn at(&self, x: usize, y: usize) -> u16 {
|
||||
self.data[y * self.stride + x]
|
||||
}
|
||||
#[inline(always)]
|
||||
pub fn row(&self, y: usize) -> &[u16] {
|
||||
&self.data[y * self.stride..(y + 1) * self.stride]
|
||||
}
|
||||
}
|
||||
|
||||
/// The intra modes this encoder uses (values are the spec's y / uv mode numbers).
|
||||
pub(crate) const INTRA_MODES_USED: [usize; 7] = [DC_PRED, V_PRED, H_PRED, PAETH_PRED, SMOOTH_PRED, SMOOTH_V_PRED, SMOOTH_H_PRED];
|
||||
|
||||
/// Inputs of the intra prediction process.
|
||||
pub(crate) struct IntraEdge {
|
||||
pub x: usize,
|
||||
pub y: usize,
|
||||
pub have_left: bool,
|
||||
pub have_above: bool,
|
||||
pub have_above_right: bool,
|
||||
pub have_below_left: bool,
|
||||
pub log2: u32,
|
||||
/// ((MiCols * MI_SIZE) >> subX) - 1 and the same for rows.
|
||||
pub max_x: i32,
|
||||
pub max_y: i32,
|
||||
pub bit_depth: u32,
|
||||
}
|
||||
|
||||
/// Edge samples: `above[0..2n]`, `left[0..2n]` and the top-left corner.
|
||||
pub(crate) struct Edges {
|
||||
above: [i32; 64],
|
||||
left: [i32; 64],
|
||||
corner: i32,
|
||||
}
|
||||
|
||||
pub(crate) fn edges(pl: &Plane, e: &IntraEdge) -> Edges {
|
||||
let w = 1usize << e.log2;
|
||||
let h = w;
|
||||
let bd = e.bit_depth;
|
||||
let (x, y) = (e.x, e.y);
|
||||
let n = w + h;
|
||||
let mut above = [0i32; 64];
|
||||
let mut left = [0i32; 64];
|
||||
if !e.have_above && e.have_left {
|
||||
let v = pl.at(x - 1, y) as i32;
|
||||
above[..n].iter_mut().for_each(|a| *a = v);
|
||||
} else if !e.have_above && !e.have_left {
|
||||
let v = (1 << (bd - 1)) - 1;
|
||||
above[..n].iter_mut().for_each(|a| *a = v);
|
||||
} else {
|
||||
let above_limit = e.max_x.min(x as i32 + if e.have_above_right { 2 * w as i32 } else { w as i32 } - 1);
|
||||
let row = pl.row(y - 1);
|
||||
for (i, a) in above[..n].iter_mut().enumerate() {
|
||||
*a = row[above_limit.min(x as i32 + i as i32) as usize] as i32;
|
||||
}
|
||||
}
|
||||
if !e.have_left && e.have_above {
|
||||
let v = pl.at(x, y - 1) as i32;
|
||||
left[..n].iter_mut().for_each(|a| *a = v);
|
||||
} else if !e.have_left && !e.have_above {
|
||||
let v = (1 << (bd - 1)) + 1;
|
||||
left[..n].iter_mut().for_each(|a| *a = v);
|
||||
} else {
|
||||
let left_limit = e.max_y.min(y as i32 + if e.have_below_left { 2 * h as i32 } else { h as i32 } - 1);
|
||||
for (i, l) in left[..n].iter_mut().enumerate() {
|
||||
*l = pl.at(x - 1, left_limit.min(y as i32 + i as i32) as usize) as i32;
|
||||
}
|
||||
}
|
||||
let corner = if e.have_above && e.have_left {
|
||||
pl.at(x - 1, y - 1) as i32
|
||||
} else if e.have_above {
|
||||
pl.at(x, y - 1) as i32
|
||||
} else if e.have_left {
|
||||
pl.at(x - 1, y) as i32
|
||||
} else {
|
||||
1 << (bd - 1)
|
||||
};
|
||||
Edges { above, left, corner }
|
||||
}
|
||||
|
||||
fn sm_weights(log2: u32) -> &'static [u8] {
|
||||
match log2 {
|
||||
2 => &SM_WEIGHTS_TX_4X4,
|
||||
3 => &SM_WEIGHTS_TX_8X8,
|
||||
4 => &SM_WEIGHTS_TX_16X16,
|
||||
_ => &SM_WEIGHTS_TX_32X32,
|
||||
}
|
||||
}
|
||||
|
||||
/// Predicts a square `2^log2` block with `mode` into `pred` (row stride = block width).
|
||||
pub(crate) fn predict_intra(ed: &Edges, e: &IntraEdge, mode: usize, pred: &mut [u16]) {
|
||||
let log2 = e.log2;
|
||||
let w = 1usize << log2;
|
||||
let h = w;
|
||||
let bd = e.bit_depth;
|
||||
let (above, left) = (&ed.above, &ed.left);
|
||||
match mode {
|
||||
DC_PRED => {
|
||||
let v = if e.have_left && e.have_above {
|
||||
let sum: i32 = left[..h].iter().sum::<i32>() + above[..w].iter().sum::<i32>() + ((w + h) >> 1) as i32;
|
||||
sum / (w + h) as i32
|
||||
} else if e.have_left {
|
||||
(left[..h].iter().sum::<i32>() + (h >> 1) as i32) >> log2
|
||||
} else if e.have_above {
|
||||
(above[..w].iter().sum::<i32>() + (w >> 1) as i32) >> log2
|
||||
} else {
|
||||
1 << (bd - 1)
|
||||
};
|
||||
let v = v.clamp(0, (1 << bd) - 1) as u16;
|
||||
pred[..w * h].iter_mut().for_each(|p| *p = v);
|
||||
}
|
||||
V_PRED => {
|
||||
for i in 0..h {
|
||||
for j in 0..w {
|
||||
pred[i * w + j] = above[j] as u16;
|
||||
}
|
||||
}
|
||||
}
|
||||
H_PRED => {
|
||||
for i in 0..h {
|
||||
pred[i * w..i * w + w].iter_mut().for_each(|p| *p = left[i] as u16);
|
||||
}
|
||||
}
|
||||
PAETH_PRED => {
|
||||
let tl = ed.corner;
|
||||
for i in 0..h {
|
||||
for j in 0..w {
|
||||
let a = above[j];
|
||||
let l = left[i];
|
||||
let base = a + l - tl;
|
||||
let p_left = (base - l).abs();
|
||||
let p_top = (base - a).abs();
|
||||
let p_tl = (base - tl).abs();
|
||||
pred[i * w + j] = if p_left <= p_top && p_left <= p_tl {
|
||||
l as u16
|
||||
} else if p_top <= p_tl {
|
||||
a as u16
|
||||
} else {
|
||||
tl as u16
|
||||
};
|
||||
}
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
let wx = sm_weights(log2);
|
||||
let wy = wx;
|
||||
for i in 0..h {
|
||||
for j in 0..w {
|
||||
let v = if mode == SMOOTH_PRED {
|
||||
let s = wy[i] as i32 * above[j] + (256 - wy[i] as i32) * left[h - 1] + wx[j] as i32 * left[i] + (256 - wx[j] as i32) * above[w - 1];
|
||||
round2(s, 9)
|
||||
} else if mode == SMOOTH_V_PRED {
|
||||
round2(wy[i] as i32 * above[j] + (256 - wy[i] as i32) * left[h - 1], 8)
|
||||
} else {
|
||||
round2(wx[j] as i32 * left[i] + (256 - wx[j] as i32) * above[w - 1], 8)
|
||||
};
|
||||
pred[i * w + j] = v as u16;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Block inter prediction (7.11.3.3 motion vector scaling without scaling, 7.11.3.4 block
|
||||
/// inter prediction with the EIGHTTAP filter, single reference so InterRound1 = 11 and the
|
||||
/// final rounding is the identity) of the `w x h` block at plane position (x, y) displaced by
|
||||
/// `mv` (1/8 luma sample units, [row, col]). `last_x` / `last_y` are the largest sample
|
||||
/// coordinates of the reference plane; the result goes to `pred` (row stride `w`).
|
||||
pub(crate) fn predict_inter(
|
||||
refp: &Plane,
|
||||
last_x: i32,
|
||||
last_y: i32,
|
||||
x: usize,
|
||||
y: usize,
|
||||
w: usize,
|
||||
h: usize,
|
||||
mv: [i32; 2],
|
||||
ss: usize,
|
||||
bit_depth: u32,
|
||||
pred: &mut [u16],
|
||||
) {
|
||||
// positions in 1/16 sample units ((startX >> 6) & 15 is the filter phase)
|
||||
let px = ((x as i32) << 4) + ((2 * mv[1]) >> ss);
|
||||
let py = ((y as i32) << 4) + ((2 * mv[0]) >> ss);
|
||||
let (fx, fy) = ((px & 15) as usize, (py & 15) as usize);
|
||||
let (ix, iy) = (px >> 4, py >> 4);
|
||||
let max = (1i32 << bit_depth) - 1;
|
||||
if fx == 0 && fy == 0 {
|
||||
// the filters' phase-0 kernel is a unit impulse: a clamped copy
|
||||
for r in 0..h {
|
||||
let ry = (iy + r as i32).clamp(0, last_y) as usize;
|
||||
let row = refp.row(ry);
|
||||
let out = &mut pred[r * w..r * w + w];
|
||||
let x0 = ix;
|
||||
if x0 >= 0 && x0 + w as i32 - 1 <= last_x {
|
||||
out.copy_from_slice(&row[x0 as usize..x0 as usize + w]);
|
||||
} else {
|
||||
for (c, o) in out.iter_mut().enumerate() {
|
||||
*o = row[(x0 + c as i32).clamp(0, last_x) as usize];
|
||||
}
|
||||
}
|
||||
}
|
||||
return;
|
||||
}
|
||||
let filt_h = if w <= 4 { 4 } else { EIGHTTAP };
|
||||
let filt_v = if h <= 4 { 4 } else { EIGHTTAP };
|
||||
let kh = &SUBPEL_FILTERS[filt_h][fx];
|
||||
let kv = &SUBPEL_FILTERS[filt_v][fy];
|
||||
let inter_h = h + 7;
|
||||
let mut inter = vec![0i32; inter_h * w];
|
||||
for r in 0..inter_h {
|
||||
let ry = (iy + r as i32 - 3).clamp(0, last_y) as usize;
|
||||
let row = refp.row(ry);
|
||||
for c in 0..w {
|
||||
let bx = ix + c as i32 - 3;
|
||||
let mut s = 0i32;
|
||||
for t in 0..8 {
|
||||
s += kh[t] as i32 * row[(bx + t as i32).clamp(0, last_x) as usize] as i32;
|
||||
}
|
||||
inter[r * w + c] = round2(s, 3);
|
||||
}
|
||||
}
|
||||
for r in 0..h {
|
||||
for c in 0..w {
|
||||
let mut s = 0i32;
|
||||
for t in 0..8 {
|
||||
s += kv[t] as i32 * inter[(r + t) * w + c];
|
||||
}
|
||||
pred[r * w + c] = round2(s, 11).clamp(0, max) as u16;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn phase_zero_filters_are_impulses() {
|
||||
for f in SUBPEL_FILTERS.iter() {
|
||||
assert_eq!(f[0], [0, 0, 0, 128, 0, 0, 0, 0]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,267 @@
|
|||
//! Multi-symbol arithmetic encoder: the exact inverse of the symbol decoder of spec 8.2.
|
||||
//!
|
||||
//! The decoder keeps `SymbolValue`, the bitwise complement of a 15-bit window into the tile data
|
||||
//! minus the lower bounds of the symbols decoded so far, and selects symbol `k` when
|
||||
//! `cur(k) <= SymbolValue < cur(k - 1)` with
|
||||
//! `cur(k) = ((SymbolRange >> 8) * ((32768 - cdf[k]) >> 6) >> 1) + 4 * (N - k - 1)` and
|
||||
//! `cur(-1) = SymbolRange`. The encoder therefore runs an ordinary carry-propagating range
|
||||
//! coder over the complemented bit stream: coding symbol `k` adds `cur(k)` to `low` and sets the
|
||||
//! range to `cur(k - 1) - cur(k)`; both are renormalised exactly like the decoder's range.
|
||||
//!
|
||||
//! At the end (spec 8.2.4 exit process) the decoder has consumed `S` renormalisation bits plus a
|
||||
//! 15-bit window, and the bits from position `S` on must be a one followed by zeros. In the
|
||||
//! complemented domain that window is `0b011_1111_1111_1111`, so the encoder picks the `S`-bit
|
||||
//! prefix `D` with `D * 2^15 + 2^14 - 1` inside `[low, low + range)`, i.e.
|
||||
//! `D = (low + 2^14) >> 15`, and writes `!D`, a one bit and zero padding.
|
||||
|
||||
pub(crate) struct SymbolWriter {
|
||||
/// Low end of the interval, holding `15 + cnt` bits not yet moved to `out` (plus carries).
|
||||
low: u64,
|
||||
rng: u32,
|
||||
cnt: u32,
|
||||
out: Vec<u8>,
|
||||
pub disable_update: bool,
|
||||
}
|
||||
|
||||
impl SymbolWriter {
|
||||
pub fn new(disable_update: bool) -> Self {
|
||||
SymbolWriter { low: 0, rng: 1 << 15, cnt: 0, out: Vec::with_capacity(1 << 16), disable_update }
|
||||
}
|
||||
|
||||
/// Adds `k` to the bytes already output (carry propagation).
|
||||
fn carry(&mut self, mut k: u64) {
|
||||
let mut i = self.out.len();
|
||||
while k != 0 {
|
||||
i -= 1;
|
||||
let t = self.out[i] as u64 + k;
|
||||
self.out[i] = t as u8;
|
||||
k = t >> 8;
|
||||
}
|
||||
}
|
||||
|
||||
#[inline(always)]
|
||||
fn encode(&mut self, cdf: &[u16], n: usize, s: usize) {
|
||||
let r8 = self.rng >> 8;
|
||||
let cur = |k: usize| -> u32 { ((r8 * ((32768 - cdf[k] as u32) >> 6)) >> 1) + 4 * (n - k - 1) as u32 };
|
||||
let hi = if s == 0 { self.rng } else { cur(s - 1) };
|
||||
let lo = cur(s);
|
||||
debug_assert!(hi > lo, "symbol {s} has an empty interval");
|
||||
self.low += lo as u64;
|
||||
self.rng = hi - lo;
|
||||
let d = 15 - (31 - self.rng.leading_zeros());
|
||||
if d == 0 {
|
||||
return;
|
||||
}
|
||||
self.rng <<= d;
|
||||
self.low <<= d;
|
||||
self.cnt += d;
|
||||
while self.cnt >= 8 {
|
||||
let sh = 15 + self.cnt - 8;
|
||||
let c = self.low >> sh;
|
||||
self.low &= (1u64 << sh) - 1;
|
||||
self.cnt -= 8;
|
||||
self.carry(c >> 8);
|
||||
self.out.push((c & 0xff) as u8);
|
||||
}
|
||||
}
|
||||
|
||||
/// Codes `s` with the adaptive CDF `cdf` (N values plus the adaptation counter).
|
||||
#[inline]
|
||||
pub fn symbol(&mut self, s: usize, cdf: &mut [u16]) {
|
||||
let n = cdf.len() - 1;
|
||||
debug_assert!(s < n);
|
||||
self.encode(cdf, n, s);
|
||||
if !self.disable_update {
|
||||
update_cdf(cdf, n, s);
|
||||
}
|
||||
}
|
||||
|
||||
/// Codes `s` with a CDF that is not adapted.
|
||||
pub fn symbol_fixed(&mut self, s: usize, cdf: &[u16]) {
|
||||
let n = cdf.len() - 1;
|
||||
self.encode(cdf, n, s);
|
||||
}
|
||||
|
||||
/// `read_bool()` counterpart.
|
||||
pub fn bool(&mut self, b: bool) {
|
||||
self.encode(&[1 << 14, 1 << 15, 0], 2, b as usize);
|
||||
}
|
||||
|
||||
/// `L(n)` counterpart (most significant bit first).
|
||||
#[cfg(test)]
|
||||
pub fn literal(&mut self, n: u32, v: u32) {
|
||||
for i in (0..n).rev() {
|
||||
self.bool((v >> i) & 1 != 0);
|
||||
}
|
||||
}
|
||||
|
||||
/// Finishes the tile: returns its bytes (with the trailing one bit and zero padding).
|
||||
pub fn finish(mut self) -> Vec<u8> {
|
||||
let low = self.low + (1 << 14);
|
||||
let mut tail = low >> 15;
|
||||
self.carry(tail >> self.cnt);
|
||||
tail &= (1u64 << self.cnt) - 1;
|
||||
let mut bytes: Vec<u8> = self.out.iter().map(|b| !b).collect();
|
||||
let n = self.cnt + 1; // 1..=8 bits: !tail then the trailing one bit
|
||||
let v = (((!tail) & ((1u64 << self.cnt) - 1)) << 1) | 1;
|
||||
bytes.push((v << (8 - n)) as u8);
|
||||
bytes
|
||||
}
|
||||
}
|
||||
|
||||
/// CDF adaptation (spec 8.2.6, as in the decoder).
|
||||
#[inline(always)]
|
||||
pub(crate) fn update_cdf(cdf: &mut [u16], n: usize, symbol: usize) {
|
||||
let count = cdf[n];
|
||||
let rate = 3 + (count > 15) as u32 + (count > 31) as u32 + (31 - (n as u32).leading_zeros()).min(2);
|
||||
let mut tmp = 0u32;
|
||||
for i in 0..n - 1 {
|
||||
if i == symbol {
|
||||
tmp = 1 << 15;
|
||||
}
|
||||
let c = cdf[i] as u32;
|
||||
if tmp < c {
|
||||
cdf[i] = (c - ((c - tmp) >> rate)) as u16;
|
||||
} else {
|
||||
cdf[i] = (c + ((tmp - c) >> rate)) as u16;
|
||||
}
|
||||
}
|
||||
cdf[n] += (count < 32) as u16;
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
/// The spec's symbol decoder (8.2.2 - 8.2.4), for checking the encoder.
|
||||
struct Dec<'a> {
|
||||
data: &'a [u8],
|
||||
pos: usize,
|
||||
value: u32,
|
||||
range: u32,
|
||||
max_bits: i64,
|
||||
}
|
||||
|
||||
impl<'a> Dec<'a> {
|
||||
fn new(data: &'a [u8]) -> Self {
|
||||
let mut d = Dec { data, pos: 0, value: 0, range: 1 << 15, max_bits: 8 * data.len() as i64 - 15 };
|
||||
let nb = (data.len() * 8).min(15) as u32;
|
||||
let buf = d.bits(nb);
|
||||
d.value = ((1 << 15) - 1) ^ (buf << (15 - nb));
|
||||
d
|
||||
}
|
||||
fn bits(&mut self, n: u32) -> u32 {
|
||||
let mut x = 0;
|
||||
for _ in 0..n {
|
||||
let b = if self.pos < self.data.len() * 8 { (self.data[self.pos >> 3] >> (7 - (self.pos & 7))) & 1 } else { 0 };
|
||||
self.pos += 1;
|
||||
x = (x << 1) | b as u32;
|
||||
}
|
||||
x
|
||||
}
|
||||
fn symbol(&mut self, cdf: &mut [u16], adapt: bool) -> usize {
|
||||
let n = cdf.len() - 1;
|
||||
let mut cur = self.range;
|
||||
let mut s = 0;
|
||||
let mut prev;
|
||||
loop {
|
||||
prev = cur;
|
||||
let f = (1u32 << 15) - cdf[s] as u32;
|
||||
cur = (((self.range >> 8) * (f >> 6)) >> 1) + 4 * (n - s - 1) as u32;
|
||||
if self.value >= cur {
|
||||
break;
|
||||
}
|
||||
s += 1;
|
||||
}
|
||||
self.range = prev - cur;
|
||||
self.value -= cur;
|
||||
let b = 15 - (31 - self.range.leading_zeros());
|
||||
if b > 0 {
|
||||
self.range <<= b;
|
||||
let nb = (b as i64).min(self.max_bits.max(0)) as u32;
|
||||
let nd = self.bits(nb) << (b - nb);
|
||||
self.value = nd ^ (((self.value + 1) << b) - 1);
|
||||
self.max_bits -= b as i64;
|
||||
}
|
||||
if adapt {
|
||||
update_cdf(cdf, n, s);
|
||||
}
|
||||
s
|
||||
}
|
||||
/// Position of the trailing one bit (8.2.4).
|
||||
fn trailing_position(&self) -> i64 {
|
||||
let consumed = (self.pos as i64).min(self.data.len() as i64 * 8);
|
||||
consumed - (15i64).min(self.max_bits + 15)
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bools_only() {
|
||||
for len in 1..200 {
|
||||
let bits: Vec<bool> = (0..len).map(|i| (i * 7 + len) % 3 == 0).collect();
|
||||
let mut w = SymbolWriter::new(false);
|
||||
for &b in &bits {
|
||||
w.bool(b);
|
||||
}
|
||||
let data = w.finish();
|
||||
let mut d = Dec::new(&data);
|
||||
for (i, &b) in bits.iter().enumerate() {
|
||||
assert_eq!(d.symbol(&mut [1 << 14, 1 << 15, 0], false) == 1, b, "len {len} bit {i} data {data:?}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn round_trip_random_symbols() {
|
||||
let mut seed = 0x1234_5678u32;
|
||||
let mut rnd = || {
|
||||
seed ^= seed << 13;
|
||||
seed ^= seed >> 17;
|
||||
seed ^= seed << 5;
|
||||
seed
|
||||
};
|
||||
for trial in 0..300 {
|
||||
let count = 1 + (rnd() % 3000) as usize;
|
||||
let mut cdfs_e = vec![[4096u16, 11264, 19328, 32768, 0]; 4];
|
||||
let mut cdfs_d = cdfs_e.clone();
|
||||
let mut syms = Vec::new();
|
||||
let mut w = SymbolWriter::new(false);
|
||||
for _ in 0..count {
|
||||
let kind = rnd() % 3;
|
||||
let ctx = (rnd() % 4) as usize;
|
||||
// skewed symbols so that the range also gets small
|
||||
let s = if rnd() % 4 == 0 { (rnd() % 4) as usize } else { 0 };
|
||||
match kind {
|
||||
0 => w.symbol(s, &mut cdfs_e[ctx]),
|
||||
1 => w.bool(s & 1 == 1),
|
||||
_ => w.literal(5, s as u32 * 7),
|
||||
}
|
||||
syms.push((kind, ctx, s));
|
||||
}
|
||||
let data = w.finish();
|
||||
let mut d = Dec::new(&data);
|
||||
for (idx, &(kind, ctx, s)) in syms.iter().enumerate() {
|
||||
let got = match kind {
|
||||
0 => d.symbol(&mut cdfs_d[ctx], true),
|
||||
1 => d.symbol(&mut [1 << 14, 1 << 15, 0], false),
|
||||
_ => {
|
||||
let mut v = 0;
|
||||
for _ in 0..5 {
|
||||
v = (v << 1) | d.symbol(&mut [1 << 14, 1 << 15, 0], false);
|
||||
}
|
||||
v / 7
|
||||
}
|
||||
};
|
||||
let want = if kind == 1 { s & 1 } else { s };
|
||||
assert_eq!(got, want, "trial {trial} idx {idx} kind {kind} of {}", syms.len());
|
||||
}
|
||||
let tp = d.trailing_position();
|
||||
assert!(tp >= 0 && (tp as usize) < data.len() * 8, "trial {trial}: trailing bit position {tp}");
|
||||
let tp = tp as usize;
|
||||
assert_eq!((data[tp >> 3] >> (7 - (tp & 7))) & 1, 1, "trial {trial}: trailing one bit");
|
||||
for p in tp + 1..data.len() * 8 {
|
||||
assert_eq!((data[p >> 3] >> (7 - (p & 7))) & 1, 0, "trial {trial}: padding");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
|
|
@ -0,0 +1,454 @@
|
|||
//! Transforms. The inverse transforms are the normative ones of spec 7.13 (mirrored from
|
||||
//! filmcraft-av1, so that the encoder reconstructs exactly like a decoder); the forward
|
||||
//! transforms are floating-point matrices scaled to match them (not normative).
|
||||
|
||||
use std::sync::OnceLock;
|
||||
|
||||
use crate::tables::*;
|
||||
|
||||
#[inline(always)]
|
||||
fn round2_64(x: i64, n: u32) -> i32 {
|
||||
if n == 0 { x as i32 } else { ((x + (1i64 << (n - 1))) >> n) as i32 }
|
||||
}
|
||||
|
||||
#[inline(always)]
|
||||
pub(crate) fn round2(x: i32, n: u32) -> i32 {
|
||||
if n == 0 { x } else { (x + (1 << (n - 1))) >> n }
|
||||
}
|
||||
|
||||
#[inline(always)]
|
||||
fn cos128(angle: i32) -> i64 {
|
||||
let a = (angle & 255) as usize;
|
||||
(match a {
|
||||
0..=64 => COS128_LOOKUP[a] as i32,
|
||||
65..=128 => -(COS128_LOOKUP[128 - a] as i32),
|
||||
129..=192 => -(COS128_LOOKUP[a - 128] as i32),
|
||||
_ => COS128_LOOKUP[256 - a] as i32,
|
||||
}) as i64
|
||||
}
|
||||
|
||||
#[inline(always)]
|
||||
fn sin128(angle: i32) -> i64 {
|
||||
cos128(angle - 64)
|
||||
}
|
||||
|
||||
fn brev(num_bits: u32, x: usize) -> usize {
|
||||
let mut t = 0;
|
||||
for i in 0..num_bits {
|
||||
let bit = (x >> i) & 1;
|
||||
t += bit << (num_bits - 1 - i);
|
||||
}
|
||||
t
|
||||
}
|
||||
|
||||
#[inline(always)]
|
||||
fn bf(t: &mut [i32], a: usize, b: usize, angle: i32, flip: bool) {
|
||||
let x = t[a] as i64 * cos128(angle) - t[b] as i64 * sin128(angle);
|
||||
let y = t[a] as i64 * sin128(angle) + t[b] as i64 * cos128(angle);
|
||||
t[a] = round2_64(x, 12);
|
||||
t[b] = round2_64(y, 12);
|
||||
if flip {
|
||||
t.swap(a, b);
|
||||
}
|
||||
}
|
||||
|
||||
#[inline(always)]
|
||||
fn hd(t: &mut [i32], a: usize, b: usize, flip: bool, r: u32) {
|
||||
let (a, b) = if flip { (b, a) } else { (a, b) };
|
||||
let x = t[a];
|
||||
let y = t[b];
|
||||
let lo = -(1i32 << (r - 1));
|
||||
let hi = (1i32 << (r - 1)) - 1;
|
||||
t[a] = (x + y).clamp(lo, hi);
|
||||
t[b] = (x - y).clamp(lo, hi);
|
||||
}
|
||||
|
||||
/// Inverse DCT of length 2^n, n = 2..=5 (7.13.2.3), in place.
|
||||
fn inverse_dct(t: &mut [i32], n: u32, r: u32) {
|
||||
let len = 1usize << n;
|
||||
let mut copy = [0i32; 32];
|
||||
copy[..len].copy_from_slice(&t[..len]);
|
||||
for i in 0..len {
|
||||
t[i] = copy[brev(n, i)];
|
||||
}
|
||||
if n >= 5 {
|
||||
for i in 0..8 {
|
||||
bf(t, 16 + i, 31 - i, 6 + ((brev(3, 7 - i) as i32) << 3), false);
|
||||
}
|
||||
}
|
||||
if n >= 4 {
|
||||
for i in 0..4 {
|
||||
bf(t, 8 + i, 15 - i, 12 + ((brev(2, 3 - i) as i32) << 4), false);
|
||||
}
|
||||
}
|
||||
if n >= 5 {
|
||||
for i in 0..8 {
|
||||
hd(t, 16 + 2 * i, 17 + 2 * i, i & 1 == 1, r);
|
||||
}
|
||||
}
|
||||
if n >= 3 {
|
||||
for i in 0..2 {
|
||||
bf(t, 4 + i, 7 - i, 56 - 32 * i as i32, false);
|
||||
}
|
||||
}
|
||||
if n >= 4 {
|
||||
for i in 0..4 {
|
||||
hd(t, 8 + 2 * i, 9 + 2 * i, i & 1 == 1, r);
|
||||
}
|
||||
}
|
||||
if n >= 5 {
|
||||
for i in 0..2 {
|
||||
for j in 0..2 {
|
||||
bf(t, 30 - 4 * i - j, 17 + 4 * i + j, 24 + ((j as i32) << 6) + (((1 - i) as i32) << 5), true);
|
||||
}
|
||||
}
|
||||
}
|
||||
for i in 0..2 {
|
||||
bf(t, 2 * i, 2 * i + 1, 32 + 16 * i as i32, i == 0);
|
||||
}
|
||||
if n >= 3 {
|
||||
for i in 0..2 {
|
||||
hd(t, 4 + 2 * i, 5 + 2 * i, i == 1, r);
|
||||
}
|
||||
}
|
||||
if n >= 4 {
|
||||
for i in 0..2 {
|
||||
bf(t, 14 - i, 9 + i, 48 + 64 * i as i32, true);
|
||||
}
|
||||
}
|
||||
if n >= 5 {
|
||||
for i in 0..4 {
|
||||
for j in 0..2 {
|
||||
hd(t, 16 + 4 * i + j, 19 + 4 * i - j, i & 1 == 1, r);
|
||||
}
|
||||
}
|
||||
}
|
||||
for i in 0..2 {
|
||||
hd(t, i, 3 - i, false, r);
|
||||
}
|
||||
if n >= 3 {
|
||||
bf(t, 6, 5, 32, true);
|
||||
}
|
||||
if n >= 4 {
|
||||
for i in 0..2 {
|
||||
for j in 0..2 {
|
||||
hd(t, 8 + 4 * i + j, 11 + 4 * i - j, i == 1, r);
|
||||
}
|
||||
}
|
||||
}
|
||||
if n >= 5 {
|
||||
for i in 0..4 {
|
||||
bf(t, 29 - i, 18 + i, 48 + (i as i32 >> 1) * 64, true);
|
||||
}
|
||||
}
|
||||
if n >= 3 {
|
||||
for i in 0..4 {
|
||||
hd(t, i, 7 - i, false, r);
|
||||
}
|
||||
}
|
||||
if n >= 4 {
|
||||
for i in 0..2 {
|
||||
bf(t, 13 - i, 10 + i, 32, true);
|
||||
}
|
||||
}
|
||||
if n >= 5 {
|
||||
for i in 0..2 {
|
||||
for j in 0..4 {
|
||||
hd(t, 16 + i * 8 + j, 23 + i * 8 - j, i == 1, r);
|
||||
}
|
||||
}
|
||||
}
|
||||
if n >= 4 {
|
||||
for i in 0..8 {
|
||||
hd(t, i, 15 - i, false, r);
|
||||
}
|
||||
}
|
||||
if n >= 5 {
|
||||
for i in 0..4 {
|
||||
bf(t, 27 - i, 20 + i, 32, true);
|
||||
}
|
||||
}
|
||||
if n >= 5 {
|
||||
for i in 0..16 {
|
||||
hd(t, i, 31 - i, false, r);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const SINPI_1_9: i64 = 1321;
|
||||
const SINPI_2_9: i64 = 2482;
|
||||
const SINPI_3_9: i64 = 3344;
|
||||
const SINPI_4_9: i64 = 3803;
|
||||
|
||||
fn inverse_adst4(t: &mut [i32]) {
|
||||
let (t0, t1, t2, t3) = (t[0] as i64, t[1] as i64, t[2] as i64, t[3] as i64);
|
||||
let mut s = [SINPI_1_9 * t0, SINPI_2_9 * t0, SINPI_3_9 * t1, SINPI_4_9 * t2, SINPI_1_9 * t2, SINPI_2_9 * t3, SINPI_4_9 * t3];
|
||||
let a7 = t0 - t2;
|
||||
let b7 = a7 + t3;
|
||||
s[0] += s[3];
|
||||
s[1] -= s[4];
|
||||
s[3] = s[2];
|
||||
s[2] = SINPI_3_9 * b7;
|
||||
s[0] += s[5];
|
||||
s[1] -= s[6];
|
||||
let x = [s[0] + s[3], s[1] + s[3], s[2], s[0] + s[1] - s[3]];
|
||||
for i in 0..4 {
|
||||
t[i] = round2_64(x[i], 12);
|
||||
}
|
||||
}
|
||||
|
||||
fn adst_in_permute(t: &mut [i32], n: u32) {
|
||||
let n0 = 1usize << n;
|
||||
let mut copy = [0i32; 16];
|
||||
copy[..n0].copy_from_slice(&t[..n0]);
|
||||
for i in 0..n0 {
|
||||
let idx = if i & 1 == 1 { i - 1 } else { n0 - i - 1 };
|
||||
t[i] = copy[idx];
|
||||
}
|
||||
}
|
||||
|
||||
fn adst_out_permute(t: &mut [i32], n: u32) {
|
||||
let n0 = 1usize << n;
|
||||
let mut copy = [0i32; 16];
|
||||
copy[..n0].copy_from_slice(&t[..n0]);
|
||||
for i in 0..n0 {
|
||||
let a = (i >> 3) & 1;
|
||||
let b = ((i >> 2) & 1) ^ ((i >> 3) & 1);
|
||||
let c = ((i >> 1) & 1) ^ ((i >> 2) & 1);
|
||||
let d = (i & 1) ^ ((i >> 1) & 1);
|
||||
let idx = ((d << 3) | (c << 2) | (b << 1) | a) >> (4 - n);
|
||||
t[i] = if i & 1 == 1 { -copy[idx] } else { copy[idx] };
|
||||
}
|
||||
}
|
||||
|
||||
fn inverse_adst8(t: &mut [i32], r: u32) {
|
||||
adst_in_permute(t, 3);
|
||||
for i in 0..4 {
|
||||
bf(t, 2 * i, 2 * i + 1, 60 - 16 * i as i32, true);
|
||||
}
|
||||
for i in 0..4 {
|
||||
hd(t, i, 4 + i, false, r);
|
||||
}
|
||||
for i in 0..2 {
|
||||
bf(t, 4 + 3 * i, 5 + i, 48 - 32 * i as i32, true);
|
||||
}
|
||||
for i in 0..2 {
|
||||
for j in 0..2 {
|
||||
hd(t, 4 * j + i, 2 + 4 * j + i, false, r);
|
||||
}
|
||||
}
|
||||
for i in 0..2 {
|
||||
bf(t, 2 + 4 * i, 3 + 4 * i, 32, true);
|
||||
}
|
||||
adst_out_permute(t, 3);
|
||||
}
|
||||
|
||||
fn inverse_adst16(t: &mut [i32], r: u32) {
|
||||
adst_in_permute(t, 4);
|
||||
for i in 0..8 {
|
||||
bf(t, 2 * i, 2 * i + 1, 62 - 8 * i as i32, true);
|
||||
}
|
||||
for i in 0..8 {
|
||||
hd(t, i, 8 + i, false, r);
|
||||
}
|
||||
for i in 0..2 {
|
||||
bf(t, 8 + 2 * i, 9 + 2 * i, 56 - 32 * i as i32, true);
|
||||
bf(t, 13 + 2 * i, 12 + 2 * i, 8 + 32 * i as i32, true);
|
||||
}
|
||||
for i in 0..4 {
|
||||
for j in 0..2 {
|
||||
hd(t, 8 * j + i, 4 + 8 * j + i, false, r);
|
||||
}
|
||||
}
|
||||
for i in 0..2 {
|
||||
for j in 0..2 {
|
||||
bf(t, 4 + 8 * j + 3 * i, 5 + 8 * j + i, 48 - 32 * i as i32, true);
|
||||
}
|
||||
}
|
||||
for i in 0..2 {
|
||||
for j in 0..4 {
|
||||
hd(t, 4 * j + i, 2 + 4 * j + i, false, r);
|
||||
}
|
||||
}
|
||||
for i in 0..4 {
|
||||
bf(t, 2 + 4 * i, 3 + 4 * i, 32, true);
|
||||
}
|
||||
adst_out_permute(t, 4);
|
||||
}
|
||||
|
||||
/// 1-D transform kind.
|
||||
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
|
||||
pub(crate) enum Kind {
|
||||
Dct,
|
||||
Adst,
|
||||
}
|
||||
|
||||
/// (vertical / column kind, horizontal / row kind) of a transform type.
|
||||
pub(crate) fn kinds(tx_type: usize) -> (Kind, Kind) {
|
||||
match tx_type {
|
||||
DCT_DCT => (Kind::Dct, Kind::Dct),
|
||||
ADST_DCT => (Kind::Adst, Kind::Dct),
|
||||
DCT_ADST => (Kind::Dct, Kind::Adst),
|
||||
ADST_ADST => (Kind::Adst, Kind::Adst),
|
||||
_ => unreachable!("unsupported transform type {tx_type}"),
|
||||
}
|
||||
}
|
||||
|
||||
fn inverse_1d(t: &mut [i32], kind: Kind, n: u32, r: u32) {
|
||||
match kind {
|
||||
Kind::Dct => inverse_dct(t, n, r),
|
||||
Kind::Adst => match n {
|
||||
2 => inverse_adst4(t),
|
||||
3 => inverse_adst8(t, r),
|
||||
_ => inverse_adst16(t, r),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
/// 2-D inverse transform (7.13.3) of a square `n x n` block (n = 4..=32) of dequantised
|
||||
/// coefficients (row stride `n`) into `residual` (row stride `n`).
|
||||
pub(crate) fn inverse_2d(dequant: &[i32], residual: &mut [i32], tx_sz: usize, tx_type: usize, bit_depth: u32) {
|
||||
let log2 = TX_WIDTH_LOG2[tx_sz] as u32;
|
||||
let n = 1usize << log2;
|
||||
let row_shift = TRANSFORM_ROW_SHIFT[tx_sz] as u32;
|
||||
let row_clamp = bit_depth + 8;
|
||||
let col_clamp = (bit_depth + 6).max(16);
|
||||
let (ck, rk) = kinds(tx_type);
|
||||
let lo = -(1i32 << (col_clamp - 1));
|
||||
let hi = (1i32 << (col_clamp - 1)) - 1;
|
||||
let mut t = [0i32; 32];
|
||||
for i in 0..n {
|
||||
let row = &dequant[i * n..i * n + n];
|
||||
if row.iter().all(|&v| v == 0) {
|
||||
residual[i * n..i * n + n].iter_mut().for_each(|v| *v = 0);
|
||||
continue;
|
||||
}
|
||||
t[..n].copy_from_slice(row);
|
||||
inverse_1d(&mut t, rk, log2, row_clamp);
|
||||
for j in 0..n {
|
||||
residual[i * n + j] = round2(t[j], row_shift).clamp(lo, hi);
|
||||
}
|
||||
}
|
||||
for j in 0..n {
|
||||
for i in 0..n {
|
||||
t[i] = residual[i * n + j];
|
||||
}
|
||||
inverse_1d(&mut t, ck, log2, col_clamp);
|
||||
for i in 0..n {
|
||||
residual[i * n + j] = round2(t[i], 4);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Orthonormal forward basis matrices `[k][i]` (row k = basis function k), sizes 4..=32.
|
||||
struct Bases {
|
||||
dct: [Vec<f32>; 4],
|
||||
adst: [Vec<f32>; 3],
|
||||
}
|
||||
|
||||
fn bases() -> &'static Bases {
|
||||
static B: OnceLock<Bases> = OnceLock::new();
|
||||
B.get_or_init(|| {
|
||||
let pi = std::f64::consts::PI;
|
||||
let dct = std::array::from_fn(|l| {
|
||||
let n = 4usize << l;
|
||||
let mut m = vec![0f32; n * n];
|
||||
for k in 0..n {
|
||||
let c = if k == 0 { std::f64::consts::FRAC_1_SQRT_2 } else { 1.0 };
|
||||
for i in 0..n {
|
||||
m[k * n + i] = (c * (2.0 / n as f64).sqrt() * (pi * (2 * i + 1) as f64 * k as f64 / (2 * n) as f64).cos()) as f32;
|
||||
}
|
||||
}
|
||||
m
|
||||
});
|
||||
let adst = std::array::from_fn(|l| {
|
||||
let n = 4usize << l;
|
||||
let mut m = vec![0f32; n * n];
|
||||
for k in 0..n {
|
||||
for i in 0..n {
|
||||
m[k * n + i] = if n == 4 {
|
||||
// the AV1 4-point ADST (sin(pi (i + 1)(2k + 1) / 9) basis)
|
||||
(2.0 / 3.0 * (pi * ((i + 1) * (2 * k + 1)) as f64 / 9.0).sin()) as f32
|
||||
} else {
|
||||
((2.0 / n as f64).sqrt() * (pi * ((2 * i + 1) * (2 * k + 1)) as f64 / (4 * n) as f64).sin()) as f32
|
||||
};
|
||||
}
|
||||
}
|
||||
m
|
||||
});
|
||||
Bases { dct, adst }
|
||||
})
|
||||
}
|
||||
|
||||
fn basis(kind: Kind, n: usize) -> &'static [f32] {
|
||||
let l = n.trailing_zeros() as usize - 2;
|
||||
match kind {
|
||||
Kind::Dct => &bases().dct[l],
|
||||
Kind::Adst => &bases().adst[l],
|
||||
}
|
||||
}
|
||||
|
||||
/// Forward 2-D transform of the `n x n` residual (row stride `n`) into orthonormal-domain
|
||||
/// coefficients scaled by 8 (so that `level = out / q` with the spec's dequantiser).
|
||||
pub(crate) fn forward_2d(residual: &[i32], out: &mut [f32], n: usize, tx_type: usize) {
|
||||
let (ck, rk) = kinds(tx_type);
|
||||
let bc = basis(ck, n);
|
||||
let br = basis(rk, n);
|
||||
// tmp[i][k] = sum_j residual[i][j] * br[k][j] (horizontal)
|
||||
let mut tmp = vec![0f32; n * n];
|
||||
for i in 0..n {
|
||||
let row = &residual[i * n..i * n + n];
|
||||
for k in 0..n {
|
||||
let b = &br[k * n..k * n + n];
|
||||
let mut s = 0f32;
|
||||
for j in 0..n {
|
||||
s += row[j] as f32 * b[j];
|
||||
}
|
||||
tmp[i * n + k] = s;
|
||||
}
|
||||
}
|
||||
// out[k][j] = 8 * sum_i bc[k][i] * tmp[i][j] (vertical)
|
||||
for k in 0..n {
|
||||
let b = &bc[k * n..k * n + n];
|
||||
for j in 0..n {
|
||||
let mut s = 0f32;
|
||||
for i in 0..n {
|
||||
s += b[i] * tmp[i * n + j];
|
||||
}
|
||||
out[k * n + j] = 8.0 * s;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
/// forward (float) then the normative inverse reproduces the residual: the forward basis
|
||||
/// and scale match the decoder's inverse transforms for every supported size and type.
|
||||
#[test]
|
||||
fn forward_matches_inverse() {
|
||||
let mut seed = 99u32;
|
||||
let mut rnd = || {
|
||||
seed = seed.wrapping_mul(1_103_515_245).wrapping_add(12345);
|
||||
((seed >> 16) % 201) as i32 - 100
|
||||
};
|
||||
for (tx, n) in [(TX_4X4, 4), (TX_8X8, 8), (TX_16X16, 16), (TX_32X32, 32)] {
|
||||
let denom = if tx == TX_32X32 { 2.0 } else { 1.0 };
|
||||
for tx_type in [DCT_DCT, ADST_DCT, DCT_ADST, ADST_ADST] {
|
||||
if n == 32 && tx_type != DCT_DCT {
|
||||
continue;
|
||||
}
|
||||
let res: Vec<i32> = (0..n * n).map(|_| rnd()).collect();
|
||||
let mut coef = vec![0f32; n * n];
|
||||
forward_2d(&res, &mut coef, n, tx_type);
|
||||
// q = 8 / denom keeps full precision: dequant = level * q / denom = coef
|
||||
let deq: Vec<i32> = coef.iter().map(|&c| (c / denom).round() as i32).collect();
|
||||
let mut out = vec![0i32; n * n];
|
||||
inverse_2d(&deq, &mut out, tx, tx_type, 8);
|
||||
let err: i32 = res.iter().zip(&out).map(|(a, b)| (a - b).abs()).max().unwrap();
|
||||
assert!(err <= 1, "tx {tx} type {tx_type}: max error {err}");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,981 @@
|
|||
//! Tile syntax writer (spec 5.11 with the 8.3 CDF selection). It keeps the same per-4x4 mode
|
||||
//! info, coefficient contexts and adaptive CDFs as a decoder and writes the decisions taken by
|
||||
//! `decide.rs` as symbols. The context derivations and the motion vector prediction process
|
||||
//! (7.10.2) mirror the decoder exactly, since symbols and coded motion vector differences
|
||||
//! depend on them.
|
||||
|
||||
use crate::cdf::Cdfs;
|
||||
use crate::decide::{FrameGeom, Leaf, Node, intra_uv_tx_type, scan_for, tx_in_set, tx_set, tx_size_for};
|
||||
use crate::symbol::SymbolWriter;
|
||||
use crate::tables::*;
|
||||
|
||||
type MvP = [i32; 2];
|
||||
|
||||
/// Mode info per 4x4 (MiSizes, YModes, RefFrames, Mvs, ...).
|
||||
struct Mi {
|
||||
cols: usize,
|
||||
mi_size: Vec<u8>,
|
||||
y_mode: Vec<u8>,
|
||||
ref_frame: Vec<[i8; 2]>,
|
||||
mv: Vec<MvP>,
|
||||
skip: Vec<bool>,
|
||||
is_inter: Vec<bool>,
|
||||
tx_type: Vec<u8>,
|
||||
written: Vec<bool>,
|
||||
}
|
||||
|
||||
impl Mi {
|
||||
#[inline(always)]
|
||||
fn idx(&self, r: usize, c: usize) -> usize {
|
||||
r * self.cols + c
|
||||
}
|
||||
}
|
||||
|
||||
/// Output of find_mv_stack.
|
||||
#[derive(Default)]
|
||||
struct MvStack {
|
||||
ref_stack_mv: [MvP; MAX_REF_MV_STACK_SIZE + 1],
|
||||
weight_stack: [u32; MAX_REF_MV_STACK_SIZE + 1],
|
||||
num_mv_found: usize,
|
||||
new_mv_count: usize,
|
||||
found_match: bool,
|
||||
close_matches: usize,
|
||||
total_matches: usize,
|
||||
new_mv_context: usize,
|
||||
ref_mv_context: usize,
|
||||
zero_mv_context: usize,
|
||||
drl_ctx_stack: [usize; MAX_REF_MV_STACK_SIZE + 1],
|
||||
}
|
||||
|
||||
pub(crate) struct TileWriter<'a> {
|
||||
g: &'a FrameGeom,
|
||||
intra_only: bool,
|
||||
qidx: u8,
|
||||
sw: SymbolWriter,
|
||||
cdf: Box<Cdfs>,
|
||||
kf_y_mode: [[[u16; 14]; 5]; 5],
|
||||
mi: Mi,
|
||||
above_level: [Vec<u8>; 3],
|
||||
above_dc: [Vec<u8>; 3],
|
||||
left_level: [Vec<u8>; 3],
|
||||
left_dc: [Vec<u8>; 3],
|
||||
// current block
|
||||
r: usize,
|
||||
c: usize,
|
||||
bsize: usize,
|
||||
avail_u: bool,
|
||||
avail_l: bool,
|
||||
mvs: MvStack,
|
||||
}
|
||||
|
||||
fn block_size_px(bs: usize) -> usize {
|
||||
4 * NUM_4X4_BLOCKS_WIDE[bs] as usize
|
||||
}
|
||||
|
||||
fn is_directional(mode: usize) -> bool {
|
||||
(V_PRED..=D67_PRED).contains(&mode)
|
||||
}
|
||||
|
||||
impl<'a> TileWriter<'a> {
|
||||
pub fn new(g: &'a FrameGeom, intra_only: bool, qidx: u8) -> Self {
|
||||
let n = g.mi_cols * g.mi_rows;
|
||||
TileWriter {
|
||||
g,
|
||||
intra_only,
|
||||
qidx,
|
||||
sw: SymbolWriter::new(false),
|
||||
cdf: Cdfs::new(qidx as u32),
|
||||
kf_y_mode: DEFAULT_INTRA_FRAME_Y_MODE_CDF,
|
||||
mi: Mi {
|
||||
cols: g.mi_cols,
|
||||
mi_size: vec![0; n],
|
||||
y_mode: vec![0; n],
|
||||
ref_frame: vec![[0, -1]; n],
|
||||
mv: vec![[0, 0]; n],
|
||||
skip: vec![false; n],
|
||||
is_inter: vec![false; n],
|
||||
tx_type: vec![0; n],
|
||||
written: vec![false; n],
|
||||
},
|
||||
above_level: std::array::from_fn(|_| vec![0; g.mi_cols + 64]),
|
||||
above_dc: std::array::from_fn(|_| vec![0; g.mi_cols + 64]),
|
||||
left_level: std::array::from_fn(|_| vec![0; g.mi_rows + 64]),
|
||||
left_dc: std::array::from_fn(|_| vec![0; g.mi_rows + 64]),
|
||||
r: 0,
|
||||
c: 0,
|
||||
bsize: 0,
|
||||
avail_u: false,
|
||||
avail_l: false,
|
||||
mvs: MvStack::default(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn finish(self) -> Vec<u8> {
|
||||
self.sw.finish()
|
||||
}
|
||||
|
||||
/// Start of a superblock row (clear_left_context).
|
||||
pub fn start_sb_row(&mut self) {
|
||||
for p in 0..3 {
|
||||
self.left_level[p].iter_mut().for_each(|v| *v = 0);
|
||||
self.left_dc[p].iter_mut().for_each(|v| *v = 0);
|
||||
}
|
||||
}
|
||||
|
||||
pub fn superblock(&mut self, r: usize, c: usize, node: &Node) {
|
||||
self.partition(r, c, BLOCK_64X64, node);
|
||||
}
|
||||
|
||||
fn partition(&mut self, r: usize, c: usize, bsize: usize, node: &Node) {
|
||||
let (mi_rows, mi_cols) = (self.g.mi_rows, self.g.mi_cols);
|
||||
if r >= mi_rows || c >= mi_cols {
|
||||
return;
|
||||
}
|
||||
let avail_u = r > 0;
|
||||
let avail_l = c > 0;
|
||||
let n4 = NUM_4X4_BLOCKS_WIDE[bsize] as usize;
|
||||
let half = n4 >> 1;
|
||||
let has_rows = r + half < mi_rows;
|
||||
let has_cols = c + half < mi_cols;
|
||||
let split = matches!(node, Node::Split(_));
|
||||
let bsl = MI_WIDTH_LOG2[bsize] as usize;
|
||||
let above = avail_u && (MI_WIDTH_LOG2[self.mi.mi_size[self.mi.idx(r - 1, c)] as usize] as usize) < bsl;
|
||||
let left = avail_l && (MI_HEIGHT_LOG2[self.mi.mi_size[self.mi.idx(r, c - 1)] as usize] as usize) < bsl;
|
||||
let ctx = left as usize * 2 + above as usize;
|
||||
let partition = if split { PARTITION_SPLIT } else { PARTITION_NONE };
|
||||
if has_rows && has_cols {
|
||||
let cdf: &mut [u16] = match bsl {
|
||||
1 => &mut self.cdf.partition_w8[ctx],
|
||||
2 => &mut self.cdf.partition_w16[ctx],
|
||||
3 => &mut self.cdf.partition_w32[ctx],
|
||||
_ => &mut self.cdf.partition_w64[ctx],
|
||||
};
|
||||
self.sw.symbol(partition, cdf);
|
||||
} else if has_cols || has_rows {
|
||||
assert!(split, "blocks crossing the frame edge are split");
|
||||
let pcdf: Vec<u16> = match bsl {
|
||||
2 => self.cdf.partition_w16[ctx].to_vec(),
|
||||
3 => self.cdf.partition_w32[ctx].to_vec(),
|
||||
_ => self.cdf.partition_w64[ctx].to_vec(),
|
||||
};
|
||||
let pr = |k: usize| pcdf[k] as i32 - if k == 0 { 0 } else { pcdf[k - 1] as i32 };
|
||||
let psum = if has_cols {
|
||||
// split_or_horz
|
||||
pr(PARTITION_VERT) + pr(PARTITION_SPLIT) + pr(PARTITION_HORZ_A) + pr(PARTITION_VERT_A) + pr(PARTITION_VERT_B) + pr(PARTITION_VERT_4)
|
||||
} else {
|
||||
// split_or_vert
|
||||
pr(PARTITION_HORZ) + pr(PARTITION_SPLIT) + pr(PARTITION_HORZ_A) + pr(PARTITION_HORZ_B) + pr(PARTITION_VERT_A) + pr(PARTITION_HORZ_4)
|
||||
};
|
||||
let cdf = [((1i32 << 15) - psum) as u16, 1 << 15, 0];
|
||||
self.sw.symbol_fixed(1, &cdf);
|
||||
} else {
|
||||
assert!(split);
|
||||
}
|
||||
match node {
|
||||
Node::Split(ch) => {
|
||||
let sub = PARTITION_SUBSIZE[PARTITION_SPLIT][bsize] as usize;
|
||||
self.partition(r, c, sub, &ch[0]);
|
||||
self.partition(r, c + half, sub, &ch[1]);
|
||||
self.partition(r + half, c, sub, &ch[2]);
|
||||
self.partition(r + half, c + half, sub, &ch[3]);
|
||||
}
|
||||
Node::Leaf(leaf) => {
|
||||
debug_assert_eq!(leaf.bsize, bsize);
|
||||
self.block(r, c, leaf);
|
||||
}
|
||||
Node::Out => unreachable!("decided block inside the frame"),
|
||||
}
|
||||
}
|
||||
|
||||
fn block(&mut self, r: usize, c: usize, b: &Leaf) {
|
||||
self.r = r;
|
||||
self.c = c;
|
||||
self.bsize = b.bsize;
|
||||
self.avail_u = r > 0;
|
||||
self.avail_l = c > 0;
|
||||
let bsize = b.bsize;
|
||||
let n4 = NUM_4X4_BLOCKS_WIDE[bsize] as usize;
|
||||
// skip
|
||||
let mut ctx = 0;
|
||||
if self.avail_u {
|
||||
ctx += self.mi.skip[self.mi.idx(r - 1, c)] as usize;
|
||||
}
|
||||
if self.avail_l {
|
||||
ctx += self.mi.skip[self.mi.idx(r, c - 1)] as usize;
|
||||
}
|
||||
let mut y_mode = b.y_mode;
|
||||
if self.intra_only {
|
||||
self.sw.symbol(b.skip as usize, &mut self.cdf.skip[ctx]);
|
||||
let above = if self.avail_u { self.mi.y_mode[self.mi.idx(r - 1, c)] as usize } else { DC_PRED };
|
||||
let left = if self.avail_l { self.mi.y_mode[self.mi.idx(r, c - 1)] as usize } else { DC_PRED };
|
||||
let actx = INTRA_MODE_CONTEXT[above] as usize;
|
||||
let lctx = INTRA_MODE_CONTEXT[left] as usize;
|
||||
self.sw.symbol(b.y_mode, &mut self.kf_y_mode[actx][lctx]);
|
||||
self.intra_angle_and_uv(b);
|
||||
} else {
|
||||
self.sw.symbol(b.skip as usize, &mut self.cdf.skip[ctx]);
|
||||
// is_inter
|
||||
let left = if self.avail_l { self.mi.ref_frame[self.mi.idx(r, c - 1)] } else { [INTRA_FRAME as i8, -1] };
|
||||
let above = if self.avail_u { self.mi.ref_frame[self.mi.idx(r - 1, c)] } else { [INTRA_FRAME as i8, -1] };
|
||||
let left_intra = left[0] <= INTRA_FRAME as i8;
|
||||
let above_intra = above[0] <= INTRA_FRAME as i8;
|
||||
let ctx = if self.avail_u && self.avail_l {
|
||||
if left_intra && above_intra { 3 } else { (left_intra || above_intra) as usize }
|
||||
} else if self.avail_u || self.avail_l {
|
||||
2 * (if self.avail_u { above_intra } else { left_intra }) as usize
|
||||
} else {
|
||||
0
|
||||
};
|
||||
self.sw.symbol(b.is_inter as usize, &mut self.cdf.is_inter[ctx]);
|
||||
if b.is_inter {
|
||||
y_mode = self.inter_block_mode_info(b, left, above);
|
||||
} else {
|
||||
let ctx = SIZE_GROUP[bsize] as usize;
|
||||
self.sw.symbol(b.y_mode, &mut self.cdf.y_mode[ctx]);
|
||||
self.intra_angle_and_uv(b);
|
||||
}
|
||||
}
|
||||
// mode info stored before the residual
|
||||
let (rows, cols) = (n4.min(self.g.mi_rows - r), n4.min(self.g.mi_cols - c));
|
||||
for y in 0..rows {
|
||||
for x in 0..cols {
|
||||
let i = self.mi.idx(r + y, c + x);
|
||||
self.mi.y_mode[i] = y_mode as u8;
|
||||
self.mi.ref_frame[i] = if b.is_inter { [LAST_FRAME as i8, -1] } else { [INTRA_FRAME as i8, -1] };
|
||||
self.mi.written[i] = true;
|
||||
if b.is_inter {
|
||||
self.mi.mv[i] = b.mv;
|
||||
}
|
||||
}
|
||||
}
|
||||
if b.skip {
|
||||
// reset_block_context( )
|
||||
for plane in 0..3 {
|
||||
let s = (plane > 0) as usize;
|
||||
for i in (c >> s)..((c + n4) >> s) {
|
||||
self.above_level[plane][i] = 0;
|
||||
self.above_dc[plane][i] = 0;
|
||||
}
|
||||
for i in (r >> s)..((r + n4) >> s) {
|
||||
self.left_level[plane][i] = 0;
|
||||
self.left_dc[plane][i] = 0;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
let log2 = (n4 * 4).trailing_zeros();
|
||||
assert!(log2 <= 5, "64x64 blocks are coded as skip blocks");
|
||||
for plane in 0..3 {
|
||||
let s = (plane > 0) as u32;
|
||||
self.coeffs(plane, (c * 4) >> s, (r * 4) >> s, tx_size_for(log2 - s), b, &b.coefs[plane]);
|
||||
}
|
||||
}
|
||||
for y in 0..rows {
|
||||
for x in 0..cols {
|
||||
let i = self.mi.idx(r + y, c + x);
|
||||
self.mi.is_inter[i] = b.is_inter;
|
||||
self.mi.skip[i] = b.skip;
|
||||
self.mi.mi_size[i] = bsize as u8;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn intra_angle_and_uv(&mut self, b: &Leaf) {
|
||||
if self.bsize >= BLOCK_8X8 && is_directional(b.y_mode) {
|
||||
self.sw.symbol(MAX_ANGLE_DELTA, &mut self.cdf.angle_delta[b.y_mode - V_PRED]);
|
||||
}
|
||||
let bs = block_size_px(self.bsize);
|
||||
if bs <= 32 {
|
||||
self.sw.symbol(b.uv_mode, &mut self.cdf.uv_mode_cfl_allowed[b.y_mode]);
|
||||
} else {
|
||||
self.sw.symbol(b.uv_mode, &mut self.cdf.uv_mode_cfl_not_allowed[b.y_mode]);
|
||||
}
|
||||
if self.bsize >= BLOCK_8X8 && is_directional(b.uv_mode) {
|
||||
self.sw.symbol(MAX_ANGLE_DELTA, &mut self.cdf.angle_delta[b.uv_mode - V_PRED]);
|
||||
}
|
||||
}
|
||||
|
||||
/// read_ref_frames (LAST_FRAME), find_mv_stack, the inter mode, DRL and motion vector.
|
||||
/// Returns the inter mode (YMode).
|
||||
fn inter_block_mode_info(&mut self, b: &Leaf, left: [i8; 2], above: [i8; 2]) -> usize {
|
||||
let count = |ft: i8| -> usize {
|
||||
let mut n = 0;
|
||||
if self.avail_u {
|
||||
n += (above[0] == ft) as usize + (above[1] == ft) as usize;
|
||||
}
|
||||
if self.avail_l {
|
||||
n += (left[0] == ft) as usize + (left[1] == ft) as usize;
|
||||
}
|
||||
n
|
||||
};
|
||||
let cmp = |a: usize, b: usize| match a.cmp(&b) {
|
||||
std::cmp::Ordering::Less => 0,
|
||||
std::cmp::Ordering::Equal => 1,
|
||||
std::cmp::Ordering::Greater => 2,
|
||||
};
|
||||
let (last, last2, last3, gold) = (count(LAST_FRAME as i8), count(LAST2_FRAME as i8), count(LAST3_FRAME as i8), count(GOLDEN_FRAME as i8));
|
||||
let (bwd, alt2, alt) = (count(BWDREF_FRAME as i8), count(ALTREF2_FRAME as i8), count(ALTREF_FRAME as i8));
|
||||
let ctx_p1 = cmp(last + last2 + last3 + gold, bwd + alt2 + alt);
|
||||
let ctx_comp_ref = cmp(last + last2, last3 + gold);
|
||||
let ctx_comp_ref_p1 = cmp(last, last2);
|
||||
self.sw.symbol(0, &mut self.cdf.single_ref[ctx_p1][0]);
|
||||
self.sw.symbol(0, &mut self.cdf.single_ref[ctx_comp_ref][2]);
|
||||
self.sw.symbol(0, &mut self.cdf.single_ref[ctx_comp_ref_p1][3]);
|
||||
self.find_mv_stack();
|
||||
let mv = b.mv;
|
||||
let s = &self.mvs;
|
||||
let mode = if mv == [0, 0] {
|
||||
GLOBALMV
|
||||
} else if mv == s.ref_stack_mv[0] {
|
||||
NEARESTMV
|
||||
} else if mv == s.ref_stack_mv[1] {
|
||||
NEARMV
|
||||
} else {
|
||||
NEWMV
|
||||
};
|
||||
let (nctx, zctx, rctx) = (s.new_mv_context, s.zero_mv_context, s.ref_mv_context);
|
||||
self.sw.symbol((mode != NEWMV) as usize, &mut self.cdf.new_mv[nctx]);
|
||||
if mode != NEWMV {
|
||||
self.sw.symbol((mode != GLOBALMV) as usize, &mut self.cdf.zero_mv[zctx]);
|
||||
if mode != GLOBALMV {
|
||||
self.sw.symbol((mode == NEARMV) as usize, &mut self.cdf.ref_mv[rctx]);
|
||||
}
|
||||
}
|
||||
// DRL: ref_mv_idx 0 for NEWMV, 1 for NEARMV
|
||||
if mode == NEWMV {
|
||||
if self.mvs.num_mv_found > 1 {
|
||||
let ctx = self.mvs.drl_ctx_stack[0];
|
||||
self.sw.symbol(0, &mut self.cdf.drl_mode[ctx]);
|
||||
}
|
||||
} else if mode == NEARMV && self.mvs.num_mv_found > 2 {
|
||||
let ctx = self.mvs.drl_ctx_stack[1];
|
||||
self.sw.symbol(0, &mut self.cdf.drl_mode[ctx]);
|
||||
}
|
||||
if mode == NEWMV {
|
||||
let pred = self.mvs.ref_stack_mv[0];
|
||||
self.write_mv([mv[0] - pred[0], mv[1] - pred[1]]);
|
||||
}
|
||||
mode
|
||||
}
|
||||
|
||||
fn write_mv(&mut self, diff: MvP) {
|
||||
let joint = ((diff[0] != 0) as usize) << 1 | (diff[1] != 0) as usize;
|
||||
// MV_JOINT_ZERO, MV_JOINT_HNZVZ (col only), MV_JOINT_HZVNZ (row only), MV_JOINT_HNZVNZ
|
||||
self.sw.symbol(joint, &mut self.cdf.mv_joint);
|
||||
for (comp, &d) in diff.iter().enumerate() {
|
||||
if d != 0 {
|
||||
self.write_mv_component(comp, d);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn write_mv_component(&mut self, comp: usize, v: i32) {
|
||||
debug_assert!(v % 2 == 0, "quarter-sample precision");
|
||||
self.sw.symbol((v < 0) as usize, &mut self.cdf.mv_sign[comp]);
|
||||
let m = v.unsigned_abs() - 1; // odd: the implied hp bit is 1
|
||||
let fr = ((m >> 1) & 3) as usize;
|
||||
let int = m >> 3;
|
||||
if int < CLASS0_SIZE as u32 {
|
||||
self.sw.symbol(MV_CLASS_0, &mut self.cdf.mv_class[comp]);
|
||||
self.sw.symbol(int as usize, &mut self.cdf.mv_class0_bit[comp]);
|
||||
self.sw.symbol(fr, &mut self.cdf.mv_class0_fr[comp][int as usize]);
|
||||
} else {
|
||||
let class = 31 - int.leading_zeros();
|
||||
let d = int - (1 << class);
|
||||
self.sw.symbol(class as usize, &mut self.cdf.mv_class[comp]);
|
||||
for i in 0..class as usize {
|
||||
self.sw.symbol(((d >> i) & 1) as usize, &mut self.cdf.mv_bit[comp][i]);
|
||||
}
|
||||
self.sw.symbol(fr, &mut self.cdf.mv_fr[comp]);
|
||||
}
|
||||
}
|
||||
|
||||
// ---- coefficients ----------------------------------------------------------------------
|
||||
|
||||
fn coeffs(&mut self, plane: usize, start_x: usize, start_y: usize, tx_sz: usize, b: &Leaf, levels: &[i32]) {
|
||||
let x4 = start_x >> 2;
|
||||
let y4 = start_y >> 2;
|
||||
let w4 = TX_WIDTH[tx_sz] as usize >> 2;
|
||||
let h4 = TX_HEIGHT[tx_sz] as usize >> 2;
|
||||
let tx_sz_ctx = (TX_SIZE_SQR[tx_sz] as usize + TX_SIZE_SQR_UP[tx_sz] as usize + 1) >> 1;
|
||||
let ptype = (plane > 0) as usize;
|
||||
let scan = scan_for(tx_sz);
|
||||
let mut eob = 0;
|
||||
for (i, &pos) in scan.iter().enumerate() {
|
||||
if levels[pos as usize] != 0 {
|
||||
eob = i + 1;
|
||||
}
|
||||
}
|
||||
let ctx = self.all_zero_ctx(plane, tx_sz, x4, y4, w4, h4);
|
||||
self.sw.symbol((eob == 0) as usize, &mut self.cdf.txb_skip[tx_sz_ctx][ctx]);
|
||||
let mut cul_level = 0u32;
|
||||
let mut dc_category = 0u8;
|
||||
if eob == 0 {
|
||||
if plane == 0 {
|
||||
self.set_tx_types(x4, y4, tx_sz, DCT_DCT);
|
||||
}
|
||||
} else {
|
||||
if plane == 0 {
|
||||
// transform_type( )
|
||||
let set = tx_set(tx_sz, b.is_inter);
|
||||
let t = b.tx_type_y;
|
||||
if set > 0 && self.qidx > 0 {
|
||||
let sqr = TX_SIZE_SQR[tx_sz] as usize;
|
||||
if b.is_inter {
|
||||
let s = TX_TYPE_INTER_INV_SET3.iter().position(|&v| v as usize == t).expect("inter tx type in set");
|
||||
self.sw.symbol(s, &mut self.cdf.inter_tx_type_set3[sqr]);
|
||||
} else {
|
||||
let s = TX_TYPE_INTRA_INV_SET2.iter().position(|&v| v as usize == t).expect("intra tx type in set");
|
||||
self.sw.symbol(s, &mut self.cdf.intra_tx_type_set2[sqr][b.y_mode]);
|
||||
}
|
||||
} else {
|
||||
assert_eq!(t, DCT_DCT);
|
||||
}
|
||||
self.set_tx_types(x4, y4, tx_sz, t);
|
||||
}
|
||||
let plane_tx_type = self.compute_tx_type(plane, tx_sz, x4, y4, b);
|
||||
debug_assert!(plane_tx_type == DCT_DCT || matches!(plane_tx_type, ADST_DCT | DCT_ADST | ADST_ADST));
|
||||
let eob_multisize = (TX_WIDTH_LOG2[tx_sz] as usize).min(5) + (TX_HEIGHT_LOG2[tx_sz] as usize).min(5) - 4;
|
||||
let eob_pt = if eob < 2 { eob } else { (32 - ((eob - 1) as u32).leading_zeros()) as usize + 1 };
|
||||
let ectx = 0; // TX_CLASS_2D
|
||||
let s = eob_pt - 1;
|
||||
match eob_multisize {
|
||||
0 => self.sw.symbol(s, &mut self.cdf.eob_pt_16[ptype][ectx]),
|
||||
1 => self.sw.symbol(s, &mut self.cdf.eob_pt_32[ptype][ectx]),
|
||||
2 => self.sw.symbol(s, &mut self.cdf.eob_pt_64[ptype][ectx]),
|
||||
3 => self.sw.symbol(s, &mut self.cdf.eob_pt_128[ptype][ectx]),
|
||||
4 => self.sw.symbol(s, &mut self.cdf.eob_pt_256[ptype][ectx]),
|
||||
5 => self.sw.symbol(s, &mut self.cdf.eob_pt_512[ptype]),
|
||||
_ => self.sw.symbol(s, &mut self.cdf.eob_pt_1024[ptype]),
|
||||
}
|
||||
if eob_pt >= 3 {
|
||||
let eob_shift = eob_pt - 3;
|
||||
let offset = eob - ((1 << (eob_pt - 2)) + 1);
|
||||
self.sw.symbol((offset >> eob_shift) & 1, &mut self.cdf.eob_extra[tx_sz_ctx][ptype][eob_pt - 3]);
|
||||
for i in 1..eob_pt - 2 {
|
||||
let sh = eob_pt - 2 - 1 - i;
|
||||
self.sw.bool((offset >> sh) & 1 == 1);
|
||||
}
|
||||
}
|
||||
let bwl = TX_WIDTH_LOG2[tx_sz] as usize;
|
||||
let height = TX_HEIGHT[tx_sz] as usize;
|
||||
let width = 1usize << bwl;
|
||||
// Quant[] as the decoder sees it while parsing: the levels coded so far (capped)
|
||||
let mut quant = vec![0i32; width * height];
|
||||
for c in (0..eob).rev() {
|
||||
let pos = scan[c] as usize;
|
||||
let a = levels[pos].unsigned_abs().min(15) as i32;
|
||||
if c == eob - 1 {
|
||||
let ctx = coeff_base_eob_ctx(c, bwl, height);
|
||||
self.sw.symbol((a.min(3) - 1) as usize, &mut self.cdf.coeff_base_eob[tx_sz_ctx][ptype][ctx]);
|
||||
} else {
|
||||
let ctx = coeff_base_ctx(&quant, tx_sz, bwl, width, height, pos);
|
||||
self.sw.symbol(a.min(3) as usize, &mut self.cdf.coeff_base[tx_sz_ctx][ptype][ctx]);
|
||||
}
|
||||
if a > NUM_BASE_LEVELS as i32 {
|
||||
let ctx = coeff_br_ctx(&quant, bwl, height, pos);
|
||||
let mut rem = a - 3;
|
||||
for _ in 0..(COEFF_BASE_RANGE / (BR_CDF_SIZE - 1)) {
|
||||
let br = rem.min(3);
|
||||
self.sw.symbol(br as usize, &mut self.cdf.coeff_br[tx_sz_ctx.min(TX_32X32)][ptype][ctx]);
|
||||
rem -= br;
|
||||
if br < 3 {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
quant[pos] = a;
|
||||
}
|
||||
for c in 0..eob {
|
||||
let pos = scan[c] as usize;
|
||||
let v = levels[pos];
|
||||
if v != 0 {
|
||||
if c == 0 {
|
||||
let ctx = self.dc_sign_ctx(plane, x4, y4, w4, h4);
|
||||
self.sw.symbol((v < 0) as usize, &mut self.cdf.dc_sign[ptype][ctx]);
|
||||
} else {
|
||||
self.sw.bool(v < 0);
|
||||
}
|
||||
}
|
||||
let a = v.unsigned_abs();
|
||||
if a > (NUM_BASE_LEVELS + COEFF_BASE_RANGE) as u32 {
|
||||
// Golomb code of a - 15
|
||||
let x = a - 15 + 1;
|
||||
let length = 32 - x.leading_zeros();
|
||||
for _ in 0..length - 1 {
|
||||
self.sw.bool(false);
|
||||
}
|
||||
self.sw.bool(true);
|
||||
for i in (0..length - 1).rev() {
|
||||
self.sw.bool((x >> i) & 1 == 1);
|
||||
}
|
||||
}
|
||||
if pos == 0 && a > 0 {
|
||||
dc_category = if v < 0 { 1 } else { 2 };
|
||||
}
|
||||
cul_level += a & 0xFFFFF;
|
||||
}
|
||||
cul_level = cul_level.min(63);
|
||||
}
|
||||
for i in 0..w4 {
|
||||
if x4 + i < self.above_level[plane].len() {
|
||||
self.above_level[plane][x4 + i] = cul_level as u8;
|
||||
self.above_dc[plane][x4 + i] = dc_category;
|
||||
}
|
||||
}
|
||||
for i in 0..h4 {
|
||||
if y4 + i < self.left_level[plane].len() {
|
||||
self.left_level[plane][y4 + i] = cul_level as u8;
|
||||
self.left_dc[plane][y4 + i] = dc_category;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn set_tx_types(&mut self, x4: usize, y4: usize, tx_sz: usize, t: usize) {
|
||||
for i in 0..(TX_WIDTH[tx_sz] as usize >> 2) {
|
||||
for j in 0..(TX_HEIGHT[tx_sz] as usize >> 2) {
|
||||
if y4 + j < self.g.mi_rows && x4 + i < self.g.mi_cols {
|
||||
let k = self.mi.idx(y4 + j, x4 + i);
|
||||
self.mi.tx_type[k] = t as u8;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn compute_tx_type(&self, plane: usize, tx_sz: usize, x4: usize, y4: usize, b: &Leaf) -> usize {
|
||||
if TX_SIZE_SQR_UP[tx_sz] as usize > TX_32X32 {
|
||||
return DCT_DCT;
|
||||
}
|
||||
if plane == 0 {
|
||||
return self.mi.tx_type[self.mi.idx(y4, x4)] as usize;
|
||||
}
|
||||
if b.is_inter {
|
||||
let set = tx_set(tx_sz, true);
|
||||
let xx = self.c.max(x4 << 1);
|
||||
let yy = self.r.max(y4 << 1);
|
||||
let t = self.mi.tx_type[self.mi.idx(yy, xx)] as usize;
|
||||
return if tx_in_set(set, true, t) { t } else { DCT_DCT };
|
||||
}
|
||||
intra_uv_tx_type(b.uv_mode, tx_sz)
|
||||
}
|
||||
|
||||
fn all_zero_ctx(&self, plane: usize, tx_sz: usize, x4: usize, y4: usize, w4: usize, h4: usize) -> usize {
|
||||
let s = (plane > 0) as usize;
|
||||
let max_x4 = self.g.mi_cols >> s;
|
||||
let max_y4 = self.g.mi_rows >> s;
|
||||
let w = TX_WIDTH[tx_sz] as usize;
|
||||
let h = TX_HEIGHT[tx_sz] as usize;
|
||||
let bsize = SUBSAMPLED_SIZE[self.bsize][s][s] as usize;
|
||||
let bw = block_size_px(bsize);
|
||||
let bh = 4 * NUM_4X4_BLOCKS_HIGH[bsize] as usize;
|
||||
let al = &self.above_level[plane];
|
||||
let ll = &self.left_level[plane];
|
||||
if plane == 0 {
|
||||
let mut top = 0u32;
|
||||
let mut left = 0u32;
|
||||
for k in 0..w4 {
|
||||
if x4 + k < max_x4 {
|
||||
top = top.max(al[x4 + k] as u32);
|
||||
}
|
||||
}
|
||||
for k in 0..h4 {
|
||||
if y4 + k < max_y4 {
|
||||
left = left.max(ll[y4 + k] as u32);
|
||||
}
|
||||
}
|
||||
top = top.min(255);
|
||||
left = left.min(255);
|
||||
if bw == w && bh == h {
|
||||
0
|
||||
} else if top == 0 && left == 0 {
|
||||
1
|
||||
} else if top == 0 || left == 0 {
|
||||
2 + (top.max(left) > 3) as usize
|
||||
} else if top.max(left) <= 3 {
|
||||
4
|
||||
} else if top.min(left) <= 3 {
|
||||
5
|
||||
} else {
|
||||
6
|
||||
}
|
||||
} else {
|
||||
let ad = &self.above_dc[plane];
|
||||
let ld = &self.left_dc[plane];
|
||||
let mut above = 0u8;
|
||||
let mut left = 0u8;
|
||||
for i in 0..w4 {
|
||||
if x4 + i < max_x4 {
|
||||
above |= al[x4 + i] | ad[x4 + i];
|
||||
}
|
||||
}
|
||||
for i in 0..h4 {
|
||||
if y4 + i < max_y4 {
|
||||
left |= ll[y4 + i] | ld[y4 + i];
|
||||
}
|
||||
}
|
||||
let mut ctx = (above != 0) as usize + (left != 0) as usize + 7;
|
||||
if bw * bh > w * h {
|
||||
ctx += 3;
|
||||
}
|
||||
ctx
|
||||
}
|
||||
}
|
||||
|
||||
fn dc_sign_ctx(&self, plane: usize, x4: usize, y4: usize, w4: usize, h4: usize) -> usize {
|
||||
let s = (plane > 0) as usize;
|
||||
let max_x4 = self.g.mi_cols >> s;
|
||||
let max_y4 = self.g.mi_rows >> s;
|
||||
let mut dc_sign = 0i32;
|
||||
for k in 0..w4 {
|
||||
if x4 + k < max_x4 {
|
||||
match self.above_dc[plane][x4 + k] {
|
||||
1 => dc_sign -= 1,
|
||||
2 => dc_sign += 1,
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
for k in 0..h4 {
|
||||
if y4 + k < max_y4 {
|
||||
match self.left_dc[plane][y4 + k] {
|
||||
1 => dc_sign -= 1,
|
||||
2 => dc_sign += 1,
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
match dc_sign.signum() {
|
||||
-1 => 1,
|
||||
1 => 2,
|
||||
_ => 0,
|
||||
}
|
||||
}
|
||||
|
||||
// ---- motion vector prediction (7.10.2), single reference LAST_FRAME --------------------
|
||||
|
||||
fn inside(&self, r: isize, c: isize) -> bool {
|
||||
c >= 0 && c < self.g.mi_cols as isize && r >= 0 && r < self.g.mi_rows as isize
|
||||
}
|
||||
|
||||
fn find_mv_stack(&mut self) {
|
||||
let n4 = NUM_4X4_BLOCKS_WIDE[self.bsize] as isize;
|
||||
let (bw4, bh4) = (n4, n4);
|
||||
self.mvs = MvStack::default();
|
||||
self.scan_row(-1);
|
||||
let mut found_above = self.mvs.found_match;
|
||||
self.mvs.found_match = false;
|
||||
self.scan_col(-1);
|
||||
let mut found_left = self.mvs.found_match;
|
||||
self.mvs.found_match = false;
|
||||
if bw4.max(bh4) <= 16 {
|
||||
self.scan_point(-1, bw4);
|
||||
}
|
||||
if self.mvs.found_match {
|
||||
found_above = true;
|
||||
}
|
||||
self.mvs.close_matches = found_above as usize + found_left as usize;
|
||||
let num_nearest = self.mvs.num_mv_found;
|
||||
let num_new = self.mvs.new_mv_count;
|
||||
for idx in 0..num_nearest {
|
||||
self.mvs.weight_stack[idx] += REF_CAT_LEVEL as u32;
|
||||
}
|
||||
self.mvs.zero_mv_context = 0;
|
||||
// use_ref_frame_mvs = 0: no temporal candidates
|
||||
self.scan_point(-1, -1);
|
||||
if self.mvs.found_match {
|
||||
found_above = true;
|
||||
}
|
||||
self.mvs.found_match = false;
|
||||
self.scan_row(-3);
|
||||
if self.mvs.found_match {
|
||||
found_above = true;
|
||||
}
|
||||
self.mvs.found_match = false;
|
||||
self.scan_col(-3);
|
||||
if self.mvs.found_match {
|
||||
found_left = true;
|
||||
}
|
||||
self.mvs.found_match = false;
|
||||
if bh4 > 1 {
|
||||
self.scan_row(-5);
|
||||
}
|
||||
if self.mvs.found_match {
|
||||
found_above = true;
|
||||
}
|
||||
self.mvs.found_match = false;
|
||||
if bw4 > 1 {
|
||||
self.scan_col(-5);
|
||||
}
|
||||
if self.mvs.found_match {
|
||||
found_left = true;
|
||||
}
|
||||
self.mvs.total_matches = found_above as usize + found_left as usize;
|
||||
let n = self.mvs.num_mv_found;
|
||||
self.sort_stack(0, num_nearest);
|
||||
self.sort_stack(num_nearest, n);
|
||||
if self.mvs.num_mv_found < 2 {
|
||||
self.extra_search();
|
||||
}
|
||||
self.context_and_clamping(num_new);
|
||||
}
|
||||
|
||||
fn scan_row(&mut self, delta_row: isize) {
|
||||
let bw4 = NUM_4X4_BLOCKS_WIDE[self.bsize] as isize;
|
||||
let (mi_row, mi_col) = (self.r as isize, self.c as isize);
|
||||
let end4 = bw4.min(self.g.mi_cols as isize - mi_col).min(16);
|
||||
let mut delta_col = 0;
|
||||
let use_step16 = bw4 >= 16;
|
||||
let mut delta_row = delta_row;
|
||||
if delta_row.abs() > 1 {
|
||||
delta_row += mi_row & 1;
|
||||
delta_col = 1 - (mi_col & 1);
|
||||
}
|
||||
let mut i = 0;
|
||||
while i < end4 {
|
||||
let mv_row = mi_row + delta_row;
|
||||
let mv_col = mi_col + delta_col + i;
|
||||
if !self.inside(mv_row, mv_col) {
|
||||
break;
|
||||
}
|
||||
let k = self.mi.idx(mv_row as usize, mv_col as usize);
|
||||
let mut len = bw4.min(NUM_4X4_BLOCKS_WIDE[self.mi.mi_size[k] as usize] as isize);
|
||||
if delta_row.abs() > 1 {
|
||||
len = len.max(2);
|
||||
}
|
||||
if use_step16 {
|
||||
len = len.max(4);
|
||||
}
|
||||
self.add_ref_mv_candidate(mv_row as usize, mv_col as usize, len as u32 * 2);
|
||||
i += len;
|
||||
}
|
||||
}
|
||||
|
||||
fn scan_col(&mut self, delta_col: isize) {
|
||||
let bh4 = NUM_4X4_BLOCKS_HIGH[self.bsize] as isize;
|
||||
let (mi_row, mi_col) = (self.r as isize, self.c as isize);
|
||||
let end4 = bh4.min(self.g.mi_rows as isize - mi_row).min(16);
|
||||
let mut delta_row = 0;
|
||||
let use_step16 = bh4 >= 16;
|
||||
let mut delta_col = delta_col;
|
||||
if delta_col.abs() > 1 {
|
||||
delta_row = 1 - (mi_row & 1);
|
||||
delta_col += mi_col & 1;
|
||||
}
|
||||
let mut i = 0;
|
||||
while i < end4 {
|
||||
let mv_row = mi_row + delta_row + i;
|
||||
let mv_col = mi_col + delta_col;
|
||||
if !self.inside(mv_row, mv_col) {
|
||||
break;
|
||||
}
|
||||
let k = self.mi.idx(mv_row as usize, mv_col as usize);
|
||||
let mut len = bh4.min(NUM_4X4_BLOCKS_HIGH[self.mi.mi_size[k] as usize] as isize);
|
||||
if delta_col.abs() > 1 {
|
||||
len = len.max(2);
|
||||
}
|
||||
if use_step16 {
|
||||
len = len.max(4);
|
||||
}
|
||||
self.add_ref_mv_candidate(mv_row as usize, mv_col as usize, len as u32 * 2);
|
||||
i += len;
|
||||
}
|
||||
}
|
||||
|
||||
fn scan_point(&mut self, delta_row: isize, delta_col: isize) {
|
||||
let mv_row = self.r as isize + delta_row;
|
||||
let mv_col = self.c as isize + delta_col;
|
||||
if self.inside(mv_row, mv_col) && self.mi.written[self.mi.idx(mv_row as usize, mv_col as usize)] {
|
||||
self.add_ref_mv_candidate(mv_row as usize, mv_col as usize, 4);
|
||||
}
|
||||
}
|
||||
|
||||
fn add_ref_mv_candidate(&mut self, mv_row: usize, mv_col: usize, weight: u32) {
|
||||
let i = self.mi.idx(mv_row, mv_col);
|
||||
if !self.mi.is_inter[i] {
|
||||
return;
|
||||
}
|
||||
let rf = self.mi.ref_frame[i];
|
||||
for cand_list in 0..2 {
|
||||
if rf[cand_list] == LAST_FRAME as i8 {
|
||||
self.search_stack(i, cand_list, weight);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn search_stack(&mut self, i: usize, cand_list: usize, weight: u32) {
|
||||
debug_assert_eq!(cand_list, 0);
|
||||
let cand_mode = self.mi.y_mode[i] as usize;
|
||||
// gm_type is IDENTITY: the stored motion vector is used
|
||||
let mut cand = self.mi.mv[i];
|
||||
lower_mv_precision(&mut cand);
|
||||
if cand_mode == NEWMV {
|
||||
self.mvs.new_mv_count += 1;
|
||||
}
|
||||
self.mvs.found_match = true;
|
||||
let n = self.mvs.num_mv_found;
|
||||
if let Some(idx) = (0..n).find(|&k| self.mvs.ref_stack_mv[k] == cand) {
|
||||
self.mvs.weight_stack[idx] += weight;
|
||||
} else if n < MAX_REF_MV_STACK_SIZE {
|
||||
self.mvs.ref_stack_mv[n] = cand;
|
||||
self.mvs.weight_stack[n] = weight;
|
||||
self.mvs.num_mv_found += 1;
|
||||
}
|
||||
}
|
||||
|
||||
fn sort_stack(&mut self, start: usize, end: usize) {
|
||||
let mut end = end;
|
||||
let s = &mut self.mvs;
|
||||
while end > start {
|
||||
let mut new_end = start;
|
||||
for idx in start + 1..end {
|
||||
if s.weight_stack[idx - 1] < s.weight_stack[idx] {
|
||||
s.weight_stack.swap(idx - 1, idx);
|
||||
s.ref_stack_mv.swap(idx - 1, idx);
|
||||
new_end = idx;
|
||||
}
|
||||
}
|
||||
end = new_end;
|
||||
}
|
||||
}
|
||||
|
||||
fn extra_search(&mut self) {
|
||||
let (mi_row, mi_col) = (self.r as isize, self.c as isize);
|
||||
let n4 = NUM_4X4_BLOCKS_WIDE[self.bsize] as isize;
|
||||
let w4 = n4.min(16).min(self.g.mi_cols as isize - mi_col);
|
||||
let h4 = n4.min(16).min(self.g.mi_rows as isize - mi_row);
|
||||
let num4x4 = w4.min(h4);
|
||||
for pass in 0..2 {
|
||||
let mut idx = 0;
|
||||
while idx < num4x4 && self.mvs.num_mv_found < 2 {
|
||||
let (mv_row, mv_col) = if pass == 0 { (mi_row - 1, mi_col + idx) } else { (mi_row + idx, mi_col - 1) };
|
||||
if !self.inside(mv_row, mv_col) {
|
||||
break;
|
||||
}
|
||||
let k = self.mi.idx(mv_row as usize, mv_col as usize);
|
||||
// add_extra_mv_candidate( ) (no sign bias: a single forward reference)
|
||||
let rfs = self.mi.ref_frame[k];
|
||||
for &cand_ref in rfs.iter().take(2) {
|
||||
if cand_ref > INTRA_FRAME as i8 {
|
||||
let cand = self.mi.mv[k];
|
||||
let n = self.mvs.num_mv_found;
|
||||
if !(0..n).any(|j| self.mvs.ref_stack_mv[j] == cand) {
|
||||
self.mvs.ref_stack_mv[n] = cand;
|
||||
self.mvs.weight_stack[n] = 2;
|
||||
self.mvs.num_mv_found += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
let sz = self.mi.mi_size[k] as usize;
|
||||
idx += if pass == 0 { NUM_4X4_BLOCKS_WIDE[sz] } else { NUM_4X4_BLOCKS_HIGH[sz] } as isize;
|
||||
}
|
||||
}
|
||||
for idx in self.mvs.num_mv_found..2 {
|
||||
self.mvs.ref_stack_mv[idx] = [0, 0]; // GlobalMvs[ 0 ]
|
||||
}
|
||||
}
|
||||
|
||||
fn context_and_clamping(&mut self, num_new: usize) {
|
||||
let bs = block_size_px(self.bsize) as i32;
|
||||
let n = self.mvs.num_mv_found;
|
||||
for idx in 0..n {
|
||||
let mut z = 0;
|
||||
if idx + 1 < n {
|
||||
let w0 = self.mvs.weight_stack[idx];
|
||||
let w1 = self.mvs.weight_stack[idx + 1];
|
||||
if w0 >= REF_CAT_LEVEL as u32 {
|
||||
if w1 < REF_CAT_LEVEL as u32 {
|
||||
z = 1;
|
||||
}
|
||||
} else {
|
||||
z = 2;
|
||||
}
|
||||
}
|
||||
self.mvs.drl_ctx_stack[idx] = z;
|
||||
}
|
||||
let n4 = NUM_4X4_BLOCKS_WIDE[self.bsize] as i32;
|
||||
for idx in 0..n {
|
||||
let mut mv = self.mvs.ref_stack_mv[idx];
|
||||
let border_r = MV_BORDER as i32 + bs * 8;
|
||||
let to_top = -((self.r as i32 * 4) * 8);
|
||||
let to_bottom = ((self.g.mi_rows as i32 - n4 - self.r as i32) * 4) * 8;
|
||||
mv[0] = mv[0].clamp(to_top - border_r, to_bottom + border_r);
|
||||
let to_left = -((self.c as i32 * 4) * 8);
|
||||
let to_right = ((self.g.mi_cols as i32 - n4 - self.c as i32) * 4) * 8;
|
||||
mv[1] = mv[1].clamp(to_left - border_r, to_right + border_r);
|
||||
self.mvs.ref_stack_mv[idx] = mv;
|
||||
}
|
||||
let (close, total) = (self.mvs.close_matches, self.mvs.total_matches);
|
||||
if close == 0 {
|
||||
self.mvs.new_mv_context = total.min(1);
|
||||
self.mvs.ref_mv_context = total;
|
||||
} else if close == 1 {
|
||||
self.mvs.new_mv_context = 3 - num_new.min(1);
|
||||
self.mvs.ref_mv_context = 2 + total;
|
||||
} else {
|
||||
self.mvs.new_mv_context = 5 - num_new.min(1);
|
||||
self.mvs.ref_mv_context = 5;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// lower_mv_precision( ) with allow_high_precision_mv = 0 and force_integer_mv = 0.
|
||||
fn lower_mv_precision(mv: &mut MvP) {
|
||||
for c in mv.iter_mut() {
|
||||
if *c & 1 != 0 {
|
||||
*c += if *c > 0 { -1 } else { 1 };
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn coeff_base_eob_ctx(c: usize, bwl: usize, height: usize) -> usize {
|
||||
if c == 0 {
|
||||
0
|
||||
} else if c <= (height << bwl) / 8 {
|
||||
1
|
||||
} else if c <= (height << bwl) / 4 {
|
||||
2
|
||||
} else {
|
||||
3
|
||||
}
|
||||
}
|
||||
|
||||
/// Context of coeff_base for TX_CLASS_2D (8.3.2).
|
||||
fn coeff_base_ctx(quant: &[i32], tx_sz: usize, bwl: usize, width: usize, height: usize, pos: usize) -> usize {
|
||||
let row = pos >> bwl;
|
||||
let col = pos - (row << bwl);
|
||||
let mut mag = 0i32;
|
||||
for off in SIG_REF_DIFF_OFFSET[TX_CLASS_2D].iter() {
|
||||
let ref_row = row + off[0] as usize;
|
||||
let ref_col = col + off[1] as usize;
|
||||
if ref_row < height && ref_col < width {
|
||||
mag += quant[(ref_row << bwl) + ref_col].abs().min(3);
|
||||
}
|
||||
}
|
||||
let ctx = ((mag + 1) >> 1).min(4) as usize;
|
||||
if row == 0 && col == 0 {
|
||||
return 0;
|
||||
}
|
||||
ctx + COEFF_BASE_CTX_OFFSET[tx_sz][row.min(4)][col.min(4)] as usize
|
||||
}
|
||||
|
||||
/// Context of coeff_br for TX_CLASS_2D.
|
||||
fn coeff_br_ctx(quant: &[i32], bwl: usize, txh: usize, pos: usize) -> usize {
|
||||
let txw = 1usize << bwl;
|
||||
let row = pos >> bwl;
|
||||
let col = pos - (row << bwl);
|
||||
let mut mag = 0i32;
|
||||
for off in MAG_REF_OFFSET_WITH_TX_CLASS[TX_CLASS_2D].iter() {
|
||||
let ref_row = row + off[0] as usize;
|
||||
let ref_col = col + off[1] as usize;
|
||||
if ref_row < txh && ref_col < txw {
|
||||
mag += quant[ref_row * txw + ref_col].min((COEFF_BASE_RANGE + NUM_BASE_LEVELS + 1) as i32);
|
||||
}
|
||||
}
|
||||
let mag = ((mag + 1) >> 1).min(6) as usize;
|
||||
if pos == 0 {
|
||||
mag
|
||||
} else if row < 2 && col < 2 {
|
||||
mag + 7
|
||||
} else {
|
||||
mag + 14
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,114 @@
|
|||
#![allow(dead_code)]
|
||||
|
||||
use effectcraft_av1enc::{Encoder, EncoderConfig, Frame, Packet};
|
||||
|
||||
/// A synthetic 4:2:0 picture.
|
||||
pub struct Pic {
|
||||
pub w: usize,
|
||||
pub h: usize,
|
||||
pub y: Vec<u16>,
|
||||
pub u: Vec<u16>,
|
||||
pub v: Vec<u16>,
|
||||
}
|
||||
|
||||
impl Pic {
|
||||
pub fn frame(&self) -> Frame<'_> {
|
||||
Frame { y: &self.y, u: &self.u, v: &self.v, y_stride: self.w, uv_stride: self.w.div_ceil(2) }
|
||||
}
|
||||
}
|
||||
|
||||
/// Smooth gradients with a soft blob, moved by (dx, dy) samples per frame index `t`, plus a
|
||||
/// textured rectangle that moves differently.
|
||||
pub fn synth(w: usize, h: usize, bd: u32, t: usize) -> Pic {
|
||||
let max = ((1u32 << bd) - 1) as f64;
|
||||
let (cw, ch) = (w.div_ceil(2), h.div_ceil(2));
|
||||
let sample = |x: f64, y: f64| -> f64 {
|
||||
let fx = x - 2.0 * t as f64;
|
||||
let fy = y - 1.0 * t as f64;
|
||||
let base = 0.25 + 0.5 * (fx / 97.0).sin() * (fy / 61.0).cos() * 0.5 + 0.2 * ((fx + fy) / (w + h) as f64);
|
||||
let d = ((fx - w as f64 / 2.0).powi(2) + (fy - h as f64 / 2.0).powi(2)).sqrt();
|
||||
let blob = 0.3 * (-d * d / (2.0 * (w.min(h) as f64 / 5.0).powi(2))).exp();
|
||||
// a checkered patch moving right by 4 per frame
|
||||
let px = x - 4.0 * t as f64 - w as f64 / 4.0;
|
||||
let py = y - h as f64 / 3.0;
|
||||
let patch =
|
||||
if (0.0..24.0).contains(&px) && (0.0..16.0).contains(&py) { if ((px as i32 / 4) + (py as i32 / 4)) % 2 == 0 { 0.25 } else { -0.25 } } else { 0.0 };
|
||||
(base + blob + patch).clamp(0.0, 1.0)
|
||||
};
|
||||
let mut y = vec![0u16; w * h];
|
||||
for j in 0..h {
|
||||
for i in 0..w {
|
||||
y[j * w + i] = (sample(i as f64, j as f64) * max).round() as u16;
|
||||
}
|
||||
}
|
||||
let mut u = vec![0u16; cw * ch];
|
||||
let mut v = vec![0u16; cw * ch];
|
||||
for j in 0..ch {
|
||||
for i in 0..cw {
|
||||
let s = sample(2.0 * i as f64, 2.0 * j as f64);
|
||||
u[j * cw + i] = ((0.5 + 0.3 * (s - 0.5) + 0.1 * (i as f64 / 13.0).sin()) * max).round() as u16;
|
||||
v[j * cw + i] = ((0.5 - 0.25 * (s - 0.5) + 0.1 * (j as f64 / 7.0).cos()) * max).round() as u16;
|
||||
}
|
||||
}
|
||||
Pic { w, h, y, u, v }
|
||||
}
|
||||
|
||||
/// Deterministic noise picture (hard case for prediction).
|
||||
pub fn noise(w: usize, h: usize, bd: u32, seed: u32) -> Pic {
|
||||
let mut s = seed.wrapping_mul(2_654_435_761).max(1);
|
||||
let mut rnd = || {
|
||||
s ^= s << 13;
|
||||
s ^= s >> 17;
|
||||
s ^= s << 5;
|
||||
s
|
||||
};
|
||||
let mask = (1u32 << bd) - 1;
|
||||
let (cw, ch) = (w.div_ceil(2), h.div_ceil(2));
|
||||
Pic {
|
||||
w,
|
||||
h,
|
||||
y: (0..w * h).map(|_| (rnd() & mask) as u16).collect(),
|
||||
u: (0..cw * ch).map(|_| (rnd() & mask) as u16).collect(),
|
||||
v: (0..cw * ch).map(|_| (rnd() & mask) as u16).collect(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn psnr(a: &[u16], b: &[u16], bd: u32) -> f64 {
|
||||
let max = ((1u32 << bd) - 1) as f64;
|
||||
let mse = a.iter().zip(b).map(|(&x, &y)| (x as f64 - y as f64).powi(2)).sum::<f64>() / a.len() as f64;
|
||||
if mse == 0.0 { 99.0 } else { 10.0 * (max * max / mse).log10() }
|
||||
}
|
||||
|
||||
pub struct Run {
|
||||
pub packets: Vec<Packet>,
|
||||
pub recons: Vec<[Vec<u16>; 3]>,
|
||||
pub psnr_y: Vec<f64>,
|
||||
}
|
||||
|
||||
/// Encodes `pics`, checks every packet decodes (filmcraft-av1) to exactly the encoder's
|
||||
/// reconstruction and returns packets, reconstructions and luma PSNRs.
|
||||
pub fn encode_and_check(cfg: EncoderConfig, pics: &[Pic]) -> Run {
|
||||
let bd = cfg.bit_depth as u32;
|
||||
let mut enc = Encoder::new(cfg).expect("config");
|
||||
let mut dec = filmcraft_av1::Decoder::new();
|
||||
let mut run = Run { packets: Vec::new(), recons: Vec::new(), psnr_y: Vec::new() };
|
||||
for (i, p) in pics.iter().enumerate() {
|
||||
let pkt = enc.encode(&p.frame());
|
||||
let rec = enc.last_reconstruction().expect("reconstruction");
|
||||
let out = dec.decode(&pkt.data).unwrap_or_else(|e| panic!("frame {i}: decode error {e:?}"));
|
||||
assert_eq!(out.len(), 1, "frame {i}: one picture per packet");
|
||||
let d = &out[0];
|
||||
assert_eq!((d.width as usize, d.height as usize), (p.w, p.h));
|
||||
for pl in 0..3 {
|
||||
if d.planes[pl] != rec[pl] {
|
||||
let first = d.planes[pl].iter().zip(&rec[pl]).position(|(a, b)| a != b).unwrap_or(usize::MAX);
|
||||
let pw = if pl == 0 { p.w } else { p.w.div_ceil(2) };
|
||||
panic!("frame {i} plane {pl}: decoder output differs from the reconstruction (first at x {} y {})", first % pw, first / pw);
|
||||
}
|
||||
}
|
||||
run.psnr_y.push(psnr(&p.y, &rec[0], bd));
|
||||
run.recons.push(rec);
|
||||
run.packets.push(pkt);
|
||||
}
|
||||
run
|
||||
}
|
||||
|
|
@ -0,0 +1,138 @@
|
|||
//! External oracle: streams decoded by ffmpeg's libdav1d must equal the encoder's reconstruction
|
||||
//! bit for bit. Skipped when `ffmpeg` (with libdav1d) is not installed. ffmpeg's native `av1`
|
||||
//! decoder only drives hardware decoders, so it is not used.
|
||||
|
||||
mod common;
|
||||
|
||||
use std::process::Command;
|
||||
|
||||
use common::*;
|
||||
use effectcraft_av1enc::{Encoder, EncoderConfig, Packet, RateControl};
|
||||
|
||||
fn have_libdav1d() -> bool {
|
||||
let Ok(out) = Command::new("ffmpeg").args(["-hide_banner", "-decoders"]).output() else {
|
||||
eprintln!("ffmpeg not found: skipping the external oracle");
|
||||
return false;
|
||||
};
|
||||
let ok = out.status.success() && String::from_utf8_lossy(&out.stdout).contains("libdav1d");
|
||||
if !ok {
|
||||
eprintln!("ffmpeg without libdav1d: skipping the external oracle");
|
||||
}
|
||||
ok
|
||||
}
|
||||
|
||||
/// IVF container with one temporal unit (temporal delimiter + packet) per frame.
|
||||
fn ivf(w: usize, h: usize, packets: &[Packet]) -> Vec<u8> {
|
||||
let mut v = Vec::new();
|
||||
v.extend_from_slice(b"DKIF");
|
||||
v.extend_from_slice(&0u16.to_le_bytes());
|
||||
v.extend_from_slice(&32u16.to_le_bytes());
|
||||
v.extend_from_slice(b"AV01");
|
||||
v.extend_from_slice(&(w as u16).to_le_bytes());
|
||||
v.extend_from_slice(&(h as u16).to_le_bytes());
|
||||
v.extend_from_slice(&30u32.to_le_bytes());
|
||||
v.extend_from_slice(&1u32.to_le_bytes());
|
||||
v.extend_from_slice(&(packets.len() as u32).to_le_bytes());
|
||||
v.extend_from_slice(&0u32.to_le_bytes());
|
||||
for (i, p) in packets.iter().enumerate() {
|
||||
v.extend_from_slice(&((p.data.len() + 2) as u32).to_le_bytes());
|
||||
v.extend_from_slice(&(i as u64).to_le_bytes());
|
||||
v.extend_from_slice(&[0x12, 0x00]); // OBU_TEMPORAL_DELIMITER
|
||||
v.extend_from_slice(&p.data);
|
||||
}
|
||||
v
|
||||
}
|
||||
|
||||
/// Encodes `pics`, decodes the IVF stream with `ffmpeg -c:v libdav1d` and compares every
|
||||
/// decoded frame with the encoder's reconstruction.
|
||||
fn check(name: &str, cfg: EncoderConfig, pics: &[Pic]) {
|
||||
let (w, h, bd) = (cfg.width as usize, cfg.height as usize, cfg.bit_depth);
|
||||
let mut enc = Encoder::new(cfg).expect("config");
|
||||
let mut packets = Vec::new();
|
||||
let mut recons = Vec::new();
|
||||
for p in pics {
|
||||
packets.push(enc.encode(&p.frame()));
|
||||
recons.push(enc.last_reconstruction().expect("reconstruction"));
|
||||
}
|
||||
let dir = std::env::temp_dir().join(format!("effectcraft-av1enc-{}-{name}", std::process::id()));
|
||||
std::fs::create_dir_all(&dir).expect("temp dir");
|
||||
let input = dir.join("in.ivf");
|
||||
let output = dir.join("out.yuv");
|
||||
std::fs::write(&input, ivf(w, h, &packets)).expect("write ivf");
|
||||
let pix_fmt = if bd == 8 { "yuv420p" } else { "yuv420p10le" };
|
||||
let res = Command::new("ffmpeg")
|
||||
.args(["-v", "error", "-y", "-c:v", "libdav1d", "-i"])
|
||||
.arg(&input)
|
||||
.args(["-f", "rawvideo", "-pix_fmt", pix_fmt])
|
||||
.arg(&output)
|
||||
.output()
|
||||
.expect("run ffmpeg");
|
||||
let stderr = String::from_utf8_lossy(&res.stderr);
|
||||
assert!(res.status.success() && stderr.trim().is_empty(), "{name}: ffmpeg -c:v libdav1d failed: {stderr}");
|
||||
let raw = std::fs::read(&output).expect("decoded yuv");
|
||||
let _ = std::fs::remove_dir_all(&dir);
|
||||
let samples: Vec<u16> =
|
||||
if bd == 8 { raw.iter().map(|&b| b as u16).collect() } else { raw.chunks_exact(2).map(|c| u16::from_le_bytes([c[0], c[1]])).collect() };
|
||||
let (cw, ch) = (w.div_ceil(2), h.div_ceil(2));
|
||||
let fsize = w * h + 2 * cw * ch;
|
||||
assert_eq!(samples.len(), fsize * pics.len(), "{name}: decoded frame count");
|
||||
for (i, rec) in recons.iter().enumerate() {
|
||||
let f = &samples[i * fsize..(i + 1) * fsize];
|
||||
let planes = [&f[..w * h], &f[w * h..w * h + cw * ch], &f[w * h + cw * ch..]];
|
||||
for p in 0..3 {
|
||||
if planes[p] != &rec[p][..] {
|
||||
let pw = if p == 0 { w } else { cw };
|
||||
let k = planes[p].iter().zip(&rec[p]).position(|(a, b)| a != b).unwrap_or(0);
|
||||
panic!("{name}: frame {i} plane {p} differs from the reconstruction at ({}, {})", k % pw, k / pw);
|
||||
}
|
||||
}
|
||||
}
|
||||
eprintln!("{name}: libdav1d bit-exact over {} frames ({} bytes)", pics.len(), packets.iter().map(|p| p.data.len()).sum::<usize>());
|
||||
}
|
||||
|
||||
fn cfg(w: usize, h: usize, bd: u8, rate: RateControl, keyint: u32) -> EncoderConfig {
|
||||
let mut c = EncoderConfig::new(w as u32, h as u32, 30, 1);
|
||||
c.bit_depth = bd;
|
||||
c.rate = rate;
|
||||
c.keyint = keyint;
|
||||
c
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn libdav1d_decodes_bit_exact() {
|
||||
if !have_libdav1d() {
|
||||
return;
|
||||
}
|
||||
// a moving pattern with inter frames
|
||||
let pics: Vec<Pic> = (0..8).map(|t| synth(176, 144, 8, t)).collect();
|
||||
check("inter8", cfg(176, 144, 8, RateControl::ConstantQ(90), 5), &pics);
|
||||
// all-intra at a fine quantizer (large coefficients, Golomb codes)
|
||||
let pics: Vec<Pic> = (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);
|
||||
// odd sizes, 8 and 10 bit
|
||||
let pics: Vec<Pic> = (0..4).map(|t| synth(33, 17, 8, t)).collect();
|
||||
check("odd8", cfg(33, 17, 8, RateControl::ConstantQ(60), 3), &pics);
|
||||
let pics: Vec<Pic> = (0..5).map(|t| synth(101, 75, 10, t)).collect();
|
||||
check("odd10", cfg(101, 75, 10, RateControl::ConstantQ(110), 4), &pics);
|
||||
// a noisy moving picture, bitrate mode
|
||||
let pics: Vec<Pic> = (0..6).map(|t| noisy(128, 80, t)).collect();
|
||||
check("bitrate", cfg(128, 80, 8, RateControl::Bitrate { kbps: 300 }, 10), &pics);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn libdav1d_decodes_1080p() {
|
||||
if !have_libdav1d() {
|
||||
return;
|
||||
}
|
||||
let pics: Vec<Pic> = (0..3).map(|t| noisy(1920, 1080, t)).collect();
|
||||
check("1080p", cfg(1920, 1080, 8, RateControl::ConstantQ(120), 30), &pics);
|
||||
}
|
||||
|
||||
fn noisy(w: usize, h: usize, t: usize) -> Pic {
|
||||
let mut p = synth(w, h, 8, t);
|
||||
let n = noise(w, h, 8, 7 + t as u32);
|
||||
for (a, b) in p.y.iter_mut().zip(&n.y) {
|
||||
*a = (*a as i32 + (*b as i32 - 128) / 16).clamp(0, 255) as u16;
|
||||
}
|
||||
p
|
||||
}
|
||||
|
|
@ -0,0 +1,92 @@
|
|||
//! Round trips through FilmCraft's spec-derived AV1 decoder: decoded frames must equal the
|
||||
//! encoder's own reconstruction bit for bit.
|
||||
|
||||
mod common;
|
||||
|
||||
use common::*;
|
||||
use effectcraft_av1enc::{EncoderConfig, RateControl};
|
||||
|
||||
fn cfg(w: u32, h: u32, bd: u8, q: u8, keyint: u32) -> EncoderConfig {
|
||||
let mut c = EncoderConfig::new(w, h, 30, 1);
|
||||
c.bit_depth = bd;
|
||||
c.rate = RateControl::ConstantQ(q);
|
||||
c.keyint = keyint;
|
||||
c
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn intra_8bit_smooth_high_quality() {
|
||||
let pics = vec![synth(128, 96, 8, 0)];
|
||||
let run = encode_and_check(cfg(128, 96, 8, 40, 1), &pics);
|
||||
eprintln!("intra 8-bit q40: {} bytes, PSNR-Y {:.2} dB", run.packets[0].data.len(), run.psnr_y[0]);
|
||||
assert!(run.packets[0].keyframe);
|
||||
assert!(run.psnr_y[0] > 35.0, "PSNR {:.2}", run.psnr_y[0]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn odd_sizes_all_intra() {
|
||||
for (w, h) in [(33, 17), (1, 1), (7, 9), (65, 63), (100, 30)] {
|
||||
let pics: Vec<Pic> = (0..2).map(|t| synth(w, h, 8, t)).collect();
|
||||
let run = encode_and_check(cfg(w as u32, h as u32, 8, 90, 1), &pics);
|
||||
assert!(run.packets.iter().all(|p| p.keyframe));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn noise_low_q_exact() {
|
||||
let pics = vec![noise(72, 40, 8, 1)];
|
||||
let run = encode_and_check(cfg(72, 40, 8, 1, 1), &pics);
|
||||
assert!(run.psnr_y[0] > 40.0, "PSNR {:.2}", run.psnr_y[0]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn inter_moving_pattern() {
|
||||
let pics: Vec<Pic> = (0..6).map(|t| synth(160, 112, 8, t)).collect();
|
||||
let run = encode_and_check(cfg(160, 112, 8, 80, 30), &pics);
|
||||
let sizes: Vec<usize> = run.packets.iter().map(|p| p.data.len()).collect();
|
||||
eprintln!("inter sizes {sizes:?} psnr {:?}", run.psnr_y);
|
||||
assert!(run.packets[0].keyframe && run.packets[1..].iter().all(|p| !p.keyframe));
|
||||
let inter_avg = sizes[1..].iter().sum::<usize>() as f64 / (sizes.len() - 1) as f64;
|
||||
assert!(inter_avg < sizes[0] as f64 * 0.5, "inter frames {inter_avg} vs key {}", sizes[0]);
|
||||
assert!(run.psnr_y.iter().all(|&p| p > 33.0));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ten_bit_inter() {
|
||||
let pics: Vec<Pic> = (0..4).map(|t| synth(96, 72, 10, t)).collect();
|
||||
let run = encode_and_check(cfg(96, 72, 10, 60, 3), &pics);
|
||||
eprintln!("10-bit psnr {:?}", run.psnr_y);
|
||||
assert!(run.psnr_y.iter().all(|&p| p > 35.0));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn odd_size_inter() {
|
||||
let pics: Vec<Pic> = (0..5).map(|t| synth(57, 43, 8, t)).collect();
|
||||
encode_and_check(cfg(57, 43, 8, 120, 10), &pics);
|
||||
let pics: Vec<Pic> = (0..3).map(|t| synth(33, 17, 10, t)).collect();
|
||||
encode_and_check(cfg(33, 17, 10, 120, 10), &pics);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bitrate_mode_tracks_target() {
|
||||
let n = 12;
|
||||
let pics: Vec<Pic> = (0..n).map(|t| synth(192, 128, 8, t)).collect();
|
||||
let mut c = cfg(192, 128, 8, 0, 12);
|
||||
c.rate = RateControl::Bitrate { kbps: 400 };
|
||||
let run = encode_and_check(c, &pics);
|
||||
let bits: usize = run.packets.iter().map(|p| p.data.len() * 8).sum();
|
||||
let kbps = bits as f64 / (n as f64 / 30.0) / 1000.0;
|
||||
eprintln!("bitrate mode: {kbps:.1} kbit/s for a 400 kbit/s target, psnr {:?}", run.psnr_y);
|
||||
assert!(kbps > 400.0 / 2.0 && kbps < 400.0 * 2.0, "{kbps} kbit/s");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn headers() {
|
||||
let enc = effectcraft_av1enc::Encoder::new(cfg(1920, 1080, 10, 100, 60)).unwrap();
|
||||
let c = enc.av1c();
|
||||
assert_eq!(c[0], 0x81);
|
||||
assert_eq!(c[1] & 0x1f, 8, "1080p30 is level 4.0");
|
||||
assert_eq!(c[2], 0b0100_1100);
|
||||
assert_eq!(&c[4..], &enc.sequence_header_obu()[..]);
|
||||
assert_eq!(enc.sequence_header_obu()[0], 0x0A);
|
||||
}
|
||||
Loading…
Reference in New Issue