mirror of
https://github.com/storytold/effectcraft.git
synced 2026-11-08 10:41:09 +00:00
Brandon ThomasandClaude Opus 5.5
406706e3d5
- effectcraft-vp9enc (L0): VP9 profile 0 key-frame encoder from the VP9 bitstream spec (8x8 intra blocks, DC/V/H/TM, 4x4 DCT/ADST, lossless WHT at quality 100, tile columns); round-trips exactly through FilmCraft's VP9 decoder and ffmpeg. - export: WebM muxer (EBML/Matroska spec) with VP9, VP9 alpha in BlockAdditional (AlphaMode 1), Opus (CodecDelay = pre-skip); WAV and AIFF outputs. - OutputFormat WebM/Wav/Aiff in the Render Queue, renderQueue.add formats, CLI --format. - Tests: FilmCraft decode-back (frames, size, pixels, audio RMS), ffprobe/ffmpeg oracles. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
858 lines
29 KiB
Rust
858 lines
29 KiB
Rust
//! A VP9 intra-frame encoder (profile 0, 8-bit 4:2:0), written from the *VP9 Bitstream &
|
||
//! Decoding Process Specification* v0.6 by mirroring its decoding process in reverse (see
|
||
//! README.md). Every frame is a key frame: one tile row (tile columns as the frame width
|
||
//! requires), 64×64 superblocks split to 8×8 blocks, 4×4 transforms (the Walsh–Hadamard
|
||
//! transform at quality 100, which is lossless), intra modes DC / V / H / TM chosen per block by
|
||
//! prediction error, default probabilities (no probability updates). The encoder reconstructs
|
||
//! exactly what a decoder reconstructs (the specification's integer inverse transforms and
|
||
//! intra edge rules), so prediction never drifts.
|
||
|
||
mod bool;
|
||
mod tables;
|
||
|
||
use bool::BoolEncoder;
|
||
use tables::*;
|
||
|
||
const DC_PRED: u8 = 0;
|
||
const V_PRED: u8 = 1;
|
||
const H_PRED: u8 = 2;
|
||
const TM_PRED: u8 = 9;
|
||
const MODES: [u8; 4] = [DC_PRED, V_PRED, H_PRED, TM_PRED];
|
||
|
||
const DCT_DCT: u8 = 0;
|
||
const ADST_DCT: u8 = 1;
|
||
const DCT_ADST: u8 = 2;
|
||
const ADST_ADST: u8 = 3;
|
||
|
||
/// mode2txfm_map for the modes used here.
|
||
fn mode_tx_type(mode: u8) -> u8 {
|
||
match mode {
|
||
V_PRED => ADST_DCT,
|
||
H_PRED => DCT_ADST,
|
||
TM_PRED => ADST_ADST,
|
||
_ => DCT_DCT,
|
||
}
|
||
}
|
||
|
||
const PARTITION_TREE: [i8; 6] = [0, 2, -1, 4, -2, -3];
|
||
const INTRA_MODE_TREE: [i8; 18] = [0, 2, -9, 4, -1, 6, 8, 12, -2, 10, -4, -5, -3, 14, -8, 16, -6, -7];
|
||
|
||
/// extra_bits[token] = (cat, numExtra, base).
|
||
const EXTRA_BITS: [(u8, u8, u16); 11] =
|
||
[(0, 0, 0), (0, 0, 1), (0, 0, 2), (0, 0, 3), (0, 0, 4), (1, 1, 5), (2, 2, 7), (3, 3, 11), (4, 4, 19), (5, 5, 35), (6, 14, 67)];
|
||
const CAT_PROBS: [&[u8]; 7] = [
|
||
&[],
|
||
&[159],
|
||
&[165, 145],
|
||
&[173, 148, 140],
|
||
&[176, 155, 140, 135],
|
||
&[180, 157, 141, 134, 130],
|
||
&[254, 254, 254, 252, 249, 243, 230, 196, 177, 153, 140, 133, 130, 129],
|
||
];
|
||
|
||
/// Encoder settings.
|
||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||
pub struct EncoderConfig {
|
||
pub width: u32,
|
||
pub height: u32,
|
||
/// 1..=100; 100 is lossless.
|
||
pub quality: u8,
|
||
/// Full-range (0–255) rather than limited (16–235) YUV; only signalled.
|
||
pub full_range: bool,
|
||
}
|
||
|
||
/// Probabilities of the full token tree (9.3.2) for model probability `p`.
|
||
fn pareto(p: u8) -> [u8; 8] {
|
||
let p = p.max(1) as usize;
|
||
let x = (p - 1) / 2;
|
||
let mut t = [0u8; 8];
|
||
for n in 0..8 {
|
||
t[n] = if p & 1 == 1 { PARETO_TABLE[x * 8 + n] } else { ((PARETO_TABLE[x * 8 + n] as u32 + PARETO_TABLE[(x + 1) * 8 + n] as u32) >> 1) as u8 };
|
||
}
|
||
t
|
||
}
|
||
|
||
fn scan(tx_type: u8) -> &'static [u16; 16] {
|
||
match tx_type {
|
||
ADST_DCT => &ROW_SCAN_4X4,
|
||
DCT_ADST => &COL_SCAN_4X4,
|
||
_ => &DEFAULT_SCAN_4X4,
|
||
}
|
||
}
|
||
|
||
/// Coefficient context neighbours (9.3.2) of each scan position.
|
||
fn neighbors(tx_type: u8) -> [(usize, usize); 16] {
|
||
let sc = scan(tx_type);
|
||
let mut nb = [(0, 0); 16];
|
||
for (c, &pos) in sc.iter().enumerate().skip(1) {
|
||
let pos = pos as usize;
|
||
let (i, j) = (pos / 4, pos % 4);
|
||
nb[c] = if i > 0 && j > 0 {
|
||
let a = (i - 1) * 4 + j;
|
||
let b = i * 4 + j - 1;
|
||
match tx_type {
|
||
DCT_ADST => (a, a),
|
||
ADST_DCT => (b, b),
|
||
_ => (a, b),
|
||
}
|
||
} else if i > 0 {
|
||
((i - 1) * 4 + j, (i - 1) * 4 + j)
|
||
} else {
|
||
(j - 1, j - 1)
|
||
};
|
||
}
|
||
nb
|
||
}
|
||
|
||
// ---------------------------------------------------------------- transforms (8.7)
|
||
|
||
fn r14(v: i32) -> i32 {
|
||
(v + (1 << 13)) >> 14
|
||
}
|
||
|
||
fn idct4(x: [i32; 4]) -> [i32; 4] {
|
||
let v0 = r14(x[0] * 11585 - x[2] * 11585);
|
||
let v1 = r14(x[0] * 11585 + x[2] * 11585);
|
||
let v2 = r14(x[1] * 6270 - x[3] * 15137);
|
||
let v3 = r14(x[1] * 15137 + x[3] * 6270);
|
||
[v1 + v3, v0 + v2, v0 - v2, v1 - v3]
|
||
}
|
||
|
||
fn iadst4(x: [i32; 4]) -> [i32; 4] {
|
||
let (s19, s29, s39, s49) = (5283i64, 9929i64, 13377i64, 15212i64);
|
||
let x: [i64; 4] = x.map(|v| v as i64);
|
||
let s0 = s19 * x[0];
|
||
let s1 = s29 * x[0];
|
||
let s2 = s39 * x[1];
|
||
let s3 = s49 * x[2];
|
||
let s4 = s19 * x[2];
|
||
let s5 = s29 * x[3];
|
||
let s6 = s49 * x[3];
|
||
let s7 = s39 * (x[0] - x[2] + x[3]);
|
||
let x0 = s0 + s3 + s5;
|
||
let x1 = s1 - s4 - s6;
|
||
let x2 = s7;
|
||
let x3 = s2;
|
||
let r = |v: i64| ((v + (1 << 13)) >> 14) as i32;
|
||
[r(x0 + x3), r(x1 + x3), r(x2), r(x0 + x1 - x3)]
|
||
}
|
||
|
||
/// The decoder's 4×4 inverse transform of dequantised coefficients (row-major) to a residual.
|
||
fn inverse4(coefs: &[i32; 16], tx_type: u8) -> [i32; 16] {
|
||
let row_adst = matches!(tx_type, DCT_ADST | ADST_ADST);
|
||
let col_adst = matches!(tx_type, ADST_DCT | ADST_ADST);
|
||
let mut t = [[0i32; 4]; 4];
|
||
for i in 0..4 {
|
||
let r = [coefs[i * 4], coefs[i * 4 + 1], coefs[i * 4 + 2], coefs[i * 4 + 3]];
|
||
t[i] = if row_adst { iadst4(r) } else { idct4(r) };
|
||
}
|
||
let mut out = [0i32; 16];
|
||
for j in 0..4 {
|
||
let c = [t[0][j], t[1][j], t[2][j], t[3][j]];
|
||
let c = if col_adst { iadst4(c) } else { idct4(c) };
|
||
for i in 0..4 {
|
||
out[i * 4 + j] = (c[i] + 8) >> 4;
|
||
}
|
||
}
|
||
out
|
||
}
|
||
|
||
/// Inverse Walsh–Hadamard (8.7.1.10) on one line.
|
||
fn iwht(x: [i32; 4], shift: u32) -> [i32; 4] {
|
||
let mut a = x[0] >> shift;
|
||
let mut c = x[1] >> shift;
|
||
let mut d = x[2] >> shift;
|
||
let mut b = x[3] >> shift;
|
||
a += c;
|
||
d -= b;
|
||
let e = (a - d) >> 1;
|
||
b = e - b;
|
||
c = e - c;
|
||
a -= b;
|
||
d += c;
|
||
[a, b, c, d]
|
||
}
|
||
|
||
/// The exact inverse of [`iwht`] (with shift 0): inputs that produce outputs `y`.
|
||
fn fwht(y: [i32; 4]) -> [i32; 4] {
|
||
let (aa, bb, cc, dd) = (y[0], y[1], y[2], y[3]);
|
||
let d1 = dd - cc;
|
||
let a1 = aa + bb;
|
||
let e = (a1 - d1) >> 1;
|
||
let b0 = e - bb;
|
||
let c0 = e - cc;
|
||
let a0 = a1 - c0;
|
||
let d0 = d1 + b0;
|
||
[a0, c0, d0, b0]
|
||
}
|
||
|
||
fn inverse_wht(coefs: &[i32; 16]) -> [i32; 16] {
|
||
let mut t = [[0i32; 4]; 4];
|
||
for i in 0..4 {
|
||
t[i] = iwht([coefs[i * 4], coefs[i * 4 + 1], coefs[i * 4 + 2], coefs[i * 4 + 3]], 2);
|
||
}
|
||
let mut out = [0i32; 16];
|
||
for j in 0..4 {
|
||
let r = iwht([t[0][j], t[1][j], t[2][j], t[3][j]], 0);
|
||
for i in 0..4 {
|
||
out[i * 4 + j] = r[i];
|
||
}
|
||
}
|
||
out
|
||
}
|
||
|
||
/// Lossless forward transform: levels (dequantised by 4) whose inverse is exactly `res`.
|
||
fn forward_wht(res: &[i32; 16]) -> [i32; 16] {
|
||
let mut t = [[0i32; 4]; 4];
|
||
for j in 0..4 {
|
||
let col = fwht([res[j], res[4 + j], res[8 + j], res[12 + j]]);
|
||
for i in 0..4 {
|
||
t[i][j] = col[i];
|
||
}
|
||
}
|
||
let mut lv = [0i32; 16];
|
||
for i in 0..4 {
|
||
let r = fwht(t[i]);
|
||
lv[i * 4..i * 4 + 4].copy_from_slice(&r);
|
||
}
|
||
lv
|
||
}
|
||
|
||
/// Forward transform matrices: the inverses of the decoder's (linear part of the) inverse
|
||
/// transforms, so `inverse4(forward(res)) ≈ res`.
|
||
fn forward_matrices() -> &'static [[[f64; 16]; 16]; 4] {
|
||
use std::sync::OnceLock;
|
||
static M: OnceLock<[[[f64; 16]; 16]; 4]> = OnceLock::new();
|
||
M.get_or_init(|| {
|
||
let mut out = [[[0.0; 16]; 16]; 4];
|
||
for (ty, o) in out.iter_mut().enumerate() {
|
||
// Columns of the inverse: response to a scaled unit coefficient.
|
||
let k = 1 << 16;
|
||
let mut inv = [[0.0f64; 16]; 16];
|
||
for c in 0..16 {
|
||
let mut e = [0i32; 16];
|
||
e[c] = k;
|
||
let r = inverse4(&e, ty as u8);
|
||
for p in 0..16 {
|
||
inv[p][c] = r[p] as f64 / k as f64;
|
||
}
|
||
}
|
||
*o = invert16(inv);
|
||
}
|
||
out
|
||
})
|
||
}
|
||
|
||
fn invert16(m: [[f64; 16]; 16]) -> [[f64; 16]; 16] {
|
||
let mut a = m;
|
||
let mut inv = [[0.0; 16]; 16];
|
||
for (i, r) in inv.iter_mut().enumerate() {
|
||
r[i] = 1.0;
|
||
}
|
||
for col in 0..16 {
|
||
let piv = (col..16).max_by(|&x, &y| a[x][col].abs().total_cmp(&a[y][col].abs())).unwrap_or(col);
|
||
a.swap(col, piv);
|
||
inv.swap(col, piv);
|
||
let d = a[col][col];
|
||
for j in 0..16 {
|
||
a[col][j] /= d;
|
||
inv[col][j] /= d;
|
||
}
|
||
for r in 0..16 {
|
||
if r != col {
|
||
let f = a[r][col];
|
||
if f != 0.0 {
|
||
for j in 0..16 {
|
||
a[r][j] -= f * a[col][j];
|
||
inv[r][j] -= f * inv[col][j];
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
inv
|
||
}
|
||
|
||
// ---------------------------------------------------------------- planes and prediction
|
||
|
||
struct Plane {
|
||
w: usize,
|
||
/// Last valid column / row for prediction edges (the 8-aligned mode-info grid).
|
||
max_x: usize,
|
||
max_y: usize,
|
||
src: Vec<u8>,
|
||
rec: Vec<u8>,
|
||
}
|
||
|
||
impl Plane {
|
||
fn new(data: &[u8], vw: usize, vh: usize, w: usize, h: usize, max_x: usize, max_y: usize) -> Plane {
|
||
let mut src = vec![0u8; w * h];
|
||
for y in 0..h {
|
||
let sy = y.min(vh.saturating_sub(1));
|
||
for x in 0..w {
|
||
let sx = x.min(vw.saturating_sub(1));
|
||
src[y * w + x] = data.get(sy * vw + sx).copied().unwrap_or(128);
|
||
}
|
||
}
|
||
Plane { w, max_x, max_y, src, rec: vec![0; w * h] }
|
||
}
|
||
}
|
||
|
||
/// Intra prediction (8.5.1) of a square block for DC / V / H / TM.
|
||
fn predict(p: &Plane, x: usize, y: usize, log2: usize, mode: u8, have_left: bool, have_above: bool) -> [i32; 64] {
|
||
let size = 1usize << log2;
|
||
let base = 128i32;
|
||
let buf = &p.rec;
|
||
let w = p.w;
|
||
let mut above = [0i32; 8];
|
||
let mut left = [0i32; 8];
|
||
if have_above {
|
||
for i in 0..size {
|
||
above[i] = buf[(y - 1) * w + (x + i).min(p.max_x)] as i32;
|
||
}
|
||
} else {
|
||
above[..size].fill(base - 1);
|
||
}
|
||
let top_left = if have_above && have_left {
|
||
buf[(y - 1) * w + (x - 1).min(p.max_x)] as i32
|
||
} else if have_above {
|
||
base + 1
|
||
} else {
|
||
base - 1
|
||
};
|
||
if have_left {
|
||
for i in 0..size {
|
||
left[i] = buf[(y + i).min(p.max_y) * w + x - 1] as i32;
|
||
}
|
||
} else {
|
||
left[..size].fill(base + 1);
|
||
}
|
||
let mut out = [0i32; 64];
|
||
for i in 0..size {
|
||
for j in 0..size {
|
||
out[i * size + j] = match mode {
|
||
V_PRED => above[j],
|
||
H_PRED => left[i],
|
||
TM_PRED => (above[j] + left[i] - top_left).clamp(0, 255),
|
||
_ => 0,
|
||
};
|
||
}
|
||
}
|
||
if mode == DC_PRED {
|
||
let v = match (have_left, have_above) {
|
||
(true, true) => (left[..size].iter().sum::<i32>() + above[..size].iter().sum::<i32>() + size as i32) >> (log2 + 1),
|
||
(true, false) => (left[..size].iter().sum::<i32>() + (1 << (log2 - 1))) >> log2,
|
||
(false, true) => (above[..size].iter().sum::<i32>() + (1 << (log2 - 1))) >> log2,
|
||
(false, false) => base,
|
||
};
|
||
out[..size * size].fill(v);
|
||
}
|
||
out
|
||
}
|
||
|
||
// ---------------------------------------------------------------- the encoder
|
||
|
||
pub struct Vp9Encoder {
|
||
cfg: EncoderConfig,
|
||
qidx: u8,
|
||
lossless: bool,
|
||
}
|
||
|
||
/// Per-tile coding state.
|
||
struct Tile<'a> {
|
||
enc: &'a Vp9Encoder,
|
||
bc: BoolEncoder,
|
||
planes: &'a mut [Plane; 3],
|
||
mi_cols: usize,
|
||
mi_rows: usize,
|
||
col_start: usize,
|
||
above_nz: [Vec<u8>; 3],
|
||
left_nz: [Vec<u8>; 3],
|
||
above_part: Vec<u8>,
|
||
left_part: Vec<u8>,
|
||
/// y mode of every 8×8 block (for mode contexts).
|
||
modes: &'a mut [u8],
|
||
}
|
||
|
||
impl Vp9Encoder {
|
||
pub fn new(cfg: EncoderConfig) -> Self {
|
||
let q = cfg.quality.clamp(1, 100);
|
||
let lossless = q >= 100;
|
||
// Quality → base_q_idx: 99 ≈ 3, 90 ≈ 30, 50 ≈ 150, 1 ≈ 255.
|
||
let qidx = if lossless { 0 } else { ((100 - q as i32) * 255 / 85).clamp(1, 255) as u8 };
|
||
Vp9Encoder { cfg, qidx, lossless }
|
||
}
|
||
|
||
pub fn config(&self) -> &EncoderConfig {
|
||
&self.cfg
|
||
}
|
||
|
||
/// Encode one key frame from planar 8-bit 4:2:0 (`y`: width×height, `u`/`v`:
|
||
/// ((w+1)/2)×((h+1)/2), tightly packed).
|
||
pub fn encode_yuv420(&mut self, y: &[u8], u: &[u8], v: &[u8]) -> Vec<u8> {
|
||
let (w, h) = (self.cfg.width.max(1) as usize, self.cfg.height.max(1) as usize);
|
||
let mi_cols = w.div_ceil(8);
|
||
let mi_rows = h.div_ceil(8);
|
||
let sb_cols = mi_cols.div_ceil(8);
|
||
let sb_rows = mi_rows.div_ceil(8);
|
||
let (pw, ph) = (sb_cols * 64, sb_rows * 64);
|
||
let (cw, ch) = (w.div_ceil(2), h.div_ceil(2));
|
||
let mut planes = [
|
||
Plane::new(y, w, h, pw, ph, mi_cols * 8 - 1, mi_rows * 8 - 1),
|
||
Plane::new(u, cw, ch, pw / 2, ph / 2, mi_cols * 4 - 1, mi_rows * 4 - 1),
|
||
Plane::new(v, cw, ch, pw / 2, ph / 2, mi_cols * 4 - 1, mi_rows * 4 - 1),
|
||
];
|
||
// Tile columns: as few as the width allows (tile_info, 6.2.13).
|
||
let mut min_log2 = 0;
|
||
while (64usize << min_log2) < sb_cols {
|
||
min_log2 += 1;
|
||
}
|
||
let mut max_log2 = 1;
|
||
while (sb_cols >> max_log2) >= 4 {
|
||
max_log2 += 1;
|
||
}
|
||
max_log2 -= 1;
|
||
let tile_cols = 1usize << min_log2;
|
||
let mut modes = vec![0u8; mi_cols * mi_rows];
|
||
let mut tiles = vec![];
|
||
for t in 0..tile_cols {
|
||
let start = (((t * sb_cols) >> min_log2) * 8).min(mi_cols);
|
||
let end = ((((t + 1) * sb_cols) >> min_log2) * 8).min(mi_cols);
|
||
let mut tile = Tile {
|
||
enc: self,
|
||
bc: BoolEncoder::new(),
|
||
planes: &mut planes,
|
||
mi_cols,
|
||
mi_rows,
|
||
col_start: start,
|
||
above_nz: [vec![0; mi_cols * 2 + 16], vec![0; mi_cols * 2 + 16], vec![0; mi_cols * 2 + 16]],
|
||
left_nz: [vec![0; 16], vec![0; 16], vec![0; 16]],
|
||
above_part: vec![0; mi_cols + 8],
|
||
left_part: vec![0; 8],
|
||
modes: &mut modes,
|
||
};
|
||
tile.bc.write(false, 128);
|
||
for r in (0..mi_rows).step_by(8) {
|
||
for p in 0..3 {
|
||
tile.left_nz[p].fill(0);
|
||
}
|
||
tile.left_part.fill(0);
|
||
for c in (start..end).step_by(8) {
|
||
tile.partition(r, c, 12);
|
||
}
|
||
}
|
||
tiles.push(tile.bc.finish());
|
||
}
|
||
// Compressed header.
|
||
let mut ch_enc = BoolEncoder::new();
|
||
ch_enc.write(false, 128);
|
||
if !self.lossless {
|
||
ch_enc.literal(0, 2); // tx_mode ONLY_4X4
|
||
}
|
||
ch_enc.literal(0, 1); // no coefficient probability updates (TX_4X4)
|
||
for _ in 0..3 {
|
||
ch_enc.write(false, 252); // skip probabilities unchanged
|
||
}
|
||
let compressed = ch_enc.finish();
|
||
// Uncompressed header.
|
||
let mut bw = BitWriter::default();
|
||
bw.put(2, 2); // frame_marker
|
||
bw.put(0, 1); // profile low bit
|
||
bw.put(0, 1); // profile high bit
|
||
bw.put(0, 1); // show_existing_frame
|
||
bw.put(0, 1); // frame_type = KEY_FRAME
|
||
bw.put(1, 1); // show_frame
|
||
bw.put(0, 1); // error_resilient_mode
|
||
bw.put(0x49, 8);
|
||
bw.put(0x83, 8);
|
||
bw.put(0x42, 8);
|
||
bw.put(2, 3); // color_space CS_BT_709
|
||
bw.put(self.cfg.full_range as u32, 1);
|
||
bw.put(w as u32 - 1, 16);
|
||
bw.put(h as u32 - 1, 16);
|
||
bw.put(0, 1); // render_and_frame_size_different
|
||
bw.put(0, 1); // refresh_frame_context
|
||
bw.put(1, 1); // frame_parallel_decoding_mode
|
||
bw.put(0, 2); // frame_context_idx
|
||
bw.put(0, 6); // loop_filter_level
|
||
bw.put(0, 3); // sharpness
|
||
bw.put(0, 1); // mode_ref_delta_enabled
|
||
bw.put(self.qidx as u32, 8);
|
||
bw.put(0, 1); // delta_q_y_dc
|
||
bw.put(0, 1); // delta_q_uv_dc
|
||
bw.put(0, 1); // delta_q_uv_ac
|
||
bw.put(0, 1); // segmentation_enabled
|
||
if min_log2 < max_log2 {
|
||
bw.put(0, 1); // no extra tile columns
|
||
}
|
||
bw.put(0, 1); // tile_rows_log2 = 0
|
||
bw.put(compressed.len() as u32, 16);
|
||
let mut out = bw.finish();
|
||
out.extend_from_slice(&compressed);
|
||
let n = tiles.len();
|
||
for (i, t) in tiles.into_iter().enumerate() {
|
||
if i + 1 < n {
|
||
out.extend_from_slice(&(t.len() as u32).to_be_bytes());
|
||
}
|
||
out.extend_from_slice(&t);
|
||
}
|
||
out
|
||
}
|
||
}
|
||
|
||
#[derive(Default)]
|
||
struct BitWriter {
|
||
out: Vec<u8>,
|
||
acc: u32,
|
||
n: u32,
|
||
}
|
||
|
||
impl BitWriter {
|
||
fn put(&mut self, v: u32, bits: u32) {
|
||
for i in (0..bits).rev() {
|
||
self.acc = (self.acc << 1) | ((v >> i) & 1);
|
||
self.n += 1;
|
||
if self.n == 8 {
|
||
self.out.push(self.acc as u8);
|
||
self.acc = 0;
|
||
self.n = 0;
|
||
}
|
||
}
|
||
}
|
||
fn finish(mut self) -> Vec<u8> {
|
||
if self.n > 0 {
|
||
self.out.push((self.acc << (8 - self.n)) as u8);
|
||
}
|
||
self.out
|
||
}
|
||
}
|
||
|
||
const BLOCK_8X8: u8 = 3;
|
||
const NUM_8X8: [usize; 13] = [1, 1, 1, 1, 1, 2, 2, 2, 4, 4, 4, 8, 8];
|
||
const MI_WIDTH_LOG2: [usize; 13] = [0, 0, 0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3];
|
||
|
||
impl Tile<'_> {
|
||
fn partition(&mut self, r: usize, c: usize, bsize: u8) {
|
||
if r >= self.mi_rows || c >= self.mi_cols {
|
||
return;
|
||
}
|
||
let num8 = NUM_8X8[bsize as usize];
|
||
let half = num8 >> 1;
|
||
let has_rows = (r + half) < self.mi_rows;
|
||
let has_cols = (c + half) < self.mi_cols;
|
||
let bsl = MI_WIDTH_LOG2[bsize as usize];
|
||
let boffset = 3 - bsl;
|
||
let mut above = 0u8;
|
||
let mut left = 0u8;
|
||
for i in 0..num8 {
|
||
above |= self.above_part[c + i];
|
||
left |= self.left_part[(r & 7) + i];
|
||
}
|
||
let above = (above >> boffset) & 1;
|
||
let left = (left >> boffset) & 1;
|
||
let ctx = bsl * 4 + left as usize * 2 + above as usize;
|
||
let probs = &KF_PARTITION_PROBS[ctx * 3..ctx * 3 + 3];
|
||
if bsize == BLOCK_8X8 {
|
||
// PARTITION_NONE.
|
||
self.bc.write(false, probs[0]);
|
||
self.block(r, c);
|
||
self.above_part[c] = 15 >> 1;
|
||
self.left_part[r & 7] = 15 >> 1;
|
||
return;
|
||
}
|
||
// PARTITION_SPLIT.
|
||
if has_rows && has_cols {
|
||
self.bc.tree(&PARTITION_TREE, probs, 3);
|
||
} else if has_cols {
|
||
self.bc.write(true, probs[1]);
|
||
} else if has_rows {
|
||
self.bc.write(true, probs[2]);
|
||
}
|
||
let sub = bsize - 3;
|
||
self.partition(r, c, sub);
|
||
self.partition(r, c + half, sub);
|
||
self.partition(r + half, c, sub);
|
||
self.partition(r + half, c + half, sub);
|
||
}
|
||
|
||
fn block(&mut self, r: usize, c: usize) {
|
||
let avail_u = r > 0;
|
||
let avail_l = c > self.col_start;
|
||
// Mode decision on the 8×8 luma block and the 4×4 chroma blocks.
|
||
let y_mode = self.choose(0, c * 8, r * 8, 3, avail_l, avail_u);
|
||
let uv_mode = {
|
||
let mut best = (i64::MAX, DC_PRED);
|
||
for m in MODES {
|
||
let cost: i64 = (1..3).map(|p| self.cost(p, c * 4, r * 4, 2, m, avail_l, avail_u)).sum();
|
||
if cost < best.0 {
|
||
best = (cost, m);
|
||
}
|
||
}
|
||
best.1
|
||
};
|
||
// skip = 0 (context: neighbours' skip, always 0 here); tx_size implied (4×4).
|
||
self.bc.write(false, DEFAULT_SKIP_PROB[0]);
|
||
let above_mode = if avail_u { self.modes[(r - 1) * self.mi_cols + c] } else { DC_PRED };
|
||
let left_mode = if avail_l { self.modes[r * self.mi_cols + c - 1] } else { DC_PRED };
|
||
let o = (above_mode as usize * 10 + left_mode as usize) * 9;
|
||
self.bc.tree(&INTRA_MODE_TREE, &KF_Y_MODE_PROBS[o..o + 9], y_mode);
|
||
let o = y_mode as usize * 9;
|
||
self.bc.tree(&INTRA_MODE_TREE, &KF_UV_MODE_PROBS[o..o + 9], uv_mode);
|
||
self.modes[r * self.mi_cols + c] = y_mode;
|
||
// Residual: luma 2×2 transform blocks, then one 4×4 block per chroma plane.
|
||
for (dy, dx) in [(0, 0), (0, 1), (1, 0), (1, 1)] {
|
||
let tx = if self.enc.lossless { DCT_DCT } else { mode_tx_type(y_mode) };
|
||
self.tx_block(0, c * 8 + dx * 4, r * 8 + dy * 4, y_mode, tx, avail_l || dx > 0, avail_u || dy > 0);
|
||
}
|
||
for p in 1..3 {
|
||
self.tx_block(p, c * 4, r * 4, uv_mode, DCT_DCT, avail_l, avail_u);
|
||
}
|
||
}
|
||
|
||
fn cost(&self, plane: usize, x: usize, y: usize, log2: usize, mode: u8, have_left: bool, have_above: bool) -> i64 {
|
||
let p = &self.planes[plane];
|
||
let pred = predict(p, x, y, log2, mode, have_left, have_above);
|
||
let size = 1 << log2;
|
||
let mut s = 0i64;
|
||
for i in 0..size {
|
||
for j in 0..size {
|
||
s += (p.src[(y + i) * p.w + x + j] as i32 - pred[i * size + j]).abs() as i64;
|
||
}
|
||
}
|
||
s
|
||
}
|
||
|
||
fn choose(&self, plane: usize, x: usize, y: usize, log2: usize, have_left: bool, have_above: bool) -> u8 {
|
||
let mut best = (i64::MAX, DC_PRED);
|
||
for m in MODES {
|
||
let cost = self.cost(plane, x, y, log2, m, have_left, have_above);
|
||
if cost < best.0 {
|
||
best = (cost, m);
|
||
}
|
||
}
|
||
best.1
|
||
}
|
||
|
||
#[allow(clippy::too_many_arguments)]
|
||
fn tx_block(&mut self, plane: usize, x: usize, y: usize, mode: u8, tx_type: u8, have_left: bool, have_above: bool) {
|
||
let ss = (plane > 0) as usize;
|
||
let max_x = (self.mi_cols * 8) >> ss;
|
||
let max_y = (self.mi_rows * 8) >> ss;
|
||
let (x4, y4) = (x >> 2, y >> 2);
|
||
let ly = y4 & (15 >> ss);
|
||
if x >= max_x || y >= max_y {
|
||
self.above_nz[plane][x4] = 0;
|
||
self.left_nz[plane][ly] = 0;
|
||
return;
|
||
}
|
||
let p = &self.planes[plane];
|
||
let pred = predict(p, x, y, 2, mode, have_left, have_above);
|
||
let mut res = [0i32; 16];
|
||
for i in 0..4 {
|
||
for j in 0..4 {
|
||
res[i * 4 + j] = p.src[(y + i) * p.w + x + j] as i32 - pred[i * 4 + j];
|
||
}
|
||
}
|
||
let (dcq, acq) = (DC_QLOOKUP[self.enc.qidx as usize], AC_QLOOKUP[self.enc.qidx as usize]);
|
||
// Quantised levels (raster order).
|
||
let levels: [i32; 16] = if self.enc.lossless {
|
||
forward_wht(&res)
|
||
} else {
|
||
let f = &forward_matrices()[tx_type as usize];
|
||
let mut lv = [0i32; 16];
|
||
for k in 0..16 {
|
||
let coef: f64 = (0..16).map(|q| f[k][q] * res[q] as f64).sum();
|
||
let qs = if k == 0 { dcq } else { acq } as f64;
|
||
let a = coef.abs() / qs;
|
||
// Rounding with a small dead zone for AC coefficients.
|
||
let l = if k == 0 { (a + 0.5).floor() } else { (a + 0.38).floor() };
|
||
lv[k] = if coef < 0.0 { -(l as i32) } else { l as i32 };
|
||
}
|
||
lv
|
||
};
|
||
let sc = scan(tx_type);
|
||
let eob = (0..16).rev().find(|&i| levels[sc[i] as usize] != 0).map_or(0, |i| i + 1);
|
||
// Tokens.
|
||
let ctx0 = (self.above_nz[plane][x4] + self.left_nz[plane][ly]) as usize;
|
||
self.tokens(plane, tx_type, &levels, eob, ctx0);
|
||
// Reconstruct exactly as the decoder does.
|
||
let mut deq = [0i32; 16];
|
||
for k in 0..16 {
|
||
deq[k] = levels[k] * if k == 0 { dcq } else { acq };
|
||
}
|
||
let resid = if eob == 0 {
|
||
[0i32; 16]
|
||
} else if self.enc.lossless {
|
||
inverse_wht(&deq)
|
||
} else {
|
||
inverse4(&deq, tx_type)
|
||
};
|
||
let p = &mut self.planes[plane];
|
||
for i in 0..4 {
|
||
for j in 0..4 {
|
||
p.rec[(y + i) * p.w + x + j] = (pred[i * 4 + j] + resid[i * 4 + j]).clamp(0, 255) as u8;
|
||
}
|
||
}
|
||
let nz = (eob > 0) as u8;
|
||
self.above_nz[plane][x4] = nz;
|
||
self.left_nz[plane][ly] = nz;
|
||
}
|
||
|
||
fn tokens(&mut self, plane: usize, tx_type: u8, levels: &[i32; 16], eob: usize, ctx0: usize) {
|
||
let sc = scan(tx_type);
|
||
let nb = neighbors(tx_type);
|
||
let ptype = (plane > 0) as usize;
|
||
// coef_probs[TX_4X4][ptype][intra]: 6 bands × 6 contexts × 3.
|
||
let base = (ptype * 2) * 6 * 6 * 3;
|
||
let mut tc = [0u8; 16];
|
||
let mut check_eob = true;
|
||
let mut c = 0;
|
||
while c < 16 {
|
||
let pos = sc[c] as usize;
|
||
let band = COEFBAND_4X4[c] as usize;
|
||
let ctx = if c == 0 { ctx0 } else { ((1 + tc[nb[c].0] + tc[nb[c].1]) >> 1) as usize };
|
||
let o = base + (band * 6 + ctx) * 3;
|
||
let p = [DEFAULT_COEF_PROBS[o], DEFAULT_COEF_PROBS[o + 1], DEFAULT_COEF_PROBS[o + 2]];
|
||
if check_eob {
|
||
let more = c < eob;
|
||
self.bc.write(more, p[0]);
|
||
if !more {
|
||
break;
|
||
}
|
||
}
|
||
let v = levels[pos].unsigned_abs();
|
||
if v == 0 {
|
||
self.bc.write(false, p[1]);
|
||
tc[pos] = 0;
|
||
check_eob = false;
|
||
c += 1;
|
||
continue;
|
||
}
|
||
self.bc.write(true, p[1]);
|
||
check_eob = true;
|
||
let token = match v {
|
||
1 => 1,
|
||
2 => 2,
|
||
3 => 3,
|
||
4 => 4,
|
||
5..=6 => 5,
|
||
7..=10 => 6,
|
||
11..=18 => 7,
|
||
19..=34 => 8,
|
||
35..=66 => 9,
|
||
_ => 10,
|
||
};
|
||
if token == 1 {
|
||
self.bc.write(false, p[2]);
|
||
} else {
|
||
self.bc.write(true, p[2]);
|
||
let pp = pareto(p[2]);
|
||
let path: &[(usize, bool)] = match token {
|
||
2 => &[(0, false), (1, false)],
|
||
3 => &[(0, false), (1, true), (2, false)],
|
||
4 => &[(0, false), (1, true), (2, true)],
|
||
5 => &[(0, true), (3, false), (4, false)],
|
||
6 => &[(0, true), (3, false), (4, true)],
|
||
7 => &[(0, true), (3, true), (5, false), (6, false)],
|
||
8 => &[(0, true), (3, true), (5, false), (6, true)],
|
||
9 => &[(0, true), (3, true), (5, true), (7, false)],
|
||
_ => &[(0, true), (3, true), (5, true), (7, true)],
|
||
};
|
||
for &(n, b) in path {
|
||
self.bc.write(b, pp[n]);
|
||
}
|
||
let (cat, num_extra, base_v) = EXTRA_BITS[token];
|
||
let extra = (v - base_v as u32).min((1 << num_extra) - 1);
|
||
let cp = CAT_PROBS[cat as usize];
|
||
for e in 0..num_extra as usize {
|
||
self.bc.write((extra >> (num_extra as usize - 1 - e)) & 1 != 0, cp[e]);
|
||
}
|
||
}
|
||
tc[pos] = ENERGY_CLASS[token];
|
||
self.bc.write(levels[pos] < 0, 128);
|
||
c += 1;
|
||
}
|
||
}
|
||
}
|
||
|
||
// ---------------------------------------------------------------- colour conversion
|
||
|
||
/// BT.709 RGB(A) 8-bit → planar YUV 4:2:0 (limited or full range).
|
||
pub fn rgba_to_yuv420(rgba: &[u8], w: u32, h: u32, full_range: bool) -> (Vec<u8>, Vec<u8>, Vec<u8>) {
|
||
let (w, h) = (w as usize, h as usize);
|
||
let (cw, ch) = (w.div_ceil(2), h.div_ceil(2));
|
||
let mut y = vec![0u8; w * h];
|
||
let mut u = vec![0u8; cw * ch];
|
||
let mut v = vec![0u8; cw * ch];
|
||
let (ys, yo, cs) = if full_range { (255.0, 0.0, 255.0) } else { (219.0, 16.0, 224.0) };
|
||
let px = |x: usize, yy: usize| -> [f32; 3] {
|
||
let i = (yy.min(h - 1) * w + x.min(w - 1)) * 4;
|
||
[rgba[i] as f32 / 255.0, rgba[i + 1] as f32 / 255.0, rgba[i + 2] as f32 / 255.0]
|
||
};
|
||
let luma = |c: [f32; 3]| 0.2126 * c[0] + 0.7152 * c[1] + 0.0722 * c[2];
|
||
for yy in 0..h {
|
||
for x in 0..w {
|
||
y[yy * w + x] = (luma(px(x, yy)) * ys + yo).round().clamp(0.0, 255.0) as u8;
|
||
}
|
||
}
|
||
for cy in 0..ch {
|
||
for cx in 0..cw {
|
||
let mut acc = [0.0f32; 3];
|
||
for (dx, dy) in [(0, 0), (1, 0), (0, 1), (1, 1)] {
|
||
let c = px(cx * 2 + dx, cy * 2 + dy);
|
||
for k in 0..3 {
|
||
acc[k] += c[k] / 4.0;
|
||
}
|
||
}
|
||
let l = luma(acc);
|
||
let cb = (acc[2] - l) / 1.8556;
|
||
let cr = (acc[0] - l) / 1.5748;
|
||
u[cy * cw + cx] = (cb * cs + 128.0).round().clamp(0.0, 255.0) as u8;
|
||
v[cy * cw + cx] = (cr * cs + 128.0).round().clamp(0.0, 255.0) as u8;
|
||
}
|
||
}
|
||
(y, u, v)
|
||
}
|
||
|
||
/// The alpha channel as a luma plane (full range) with neutral chroma — the second VP9 stream of
|
||
/// WebM alpha.
|
||
pub fn alpha_to_yuv420(rgba: &[u8], w: u32, h: u32) -> (Vec<u8>, Vec<u8>, Vec<u8>) {
|
||
let (w, h) = (w as usize, h as usize);
|
||
let (cw, ch) = (w.div_ceil(2), h.div_ceil(2));
|
||
let y: Vec<u8> = rgba.chunks_exact(4).take(w * h).map(|p| p[3]).collect();
|
||
(y, vec![128; cw * ch], vec![128; cw * ch])
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
#[test]
|
||
fn wht_is_exactly_invertible() {
|
||
let mut s = 7u32;
|
||
for _ in 0..2000 {
|
||
let mut r = [0i32; 16];
|
||
for v in &mut r {
|
||
s = s.wrapping_mul(1664525).wrapping_add(1013904223);
|
||
*v = (s >> 16) as i32 % 511 - 255;
|
||
}
|
||
let lv = forward_wht(&r);
|
||
let deq = lv.map(|l| l * 4);
|
||
assert_eq!(inverse_wht(&deq), r);
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn forward_transforms_invert_the_decoder() {
|
||
for ty in 0..4u8 {
|
||
let f = &forward_matrices()[ty as usize];
|
||
let r: [i32; 16] = std::array::from_fn(|i| (i as i32 * 37 % 61) - 30);
|
||
let c: [i32; 16] = std::array::from_fn(|k| (0..16).map(|q| f[k][q] * r[q] as f64).sum::<f64>().round() as i32);
|
||
let back = inverse4(&c, ty);
|
||
for i in 0..16 {
|
||
assert!((back[i] - r[i]).abs() <= 1, "type {ty}: {back:?} vs {r:?}");
|
||
}
|
||
}
|
||
}
|
||
}
|