shifathosn
/
filmcraft
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

M6.1: filmcraft-h264enc — clean-room H.264 encoder (High/Main/Baseline, CABAC/CAVLC, intra/inter mode decision, deblocking, RC) 

- SPS/PPS/AUD/SEI with VUI (SAR, colour description, timing, reorder), avcC, Annex-B or length-prefixed packets
- Intra 16x16/4x4/8x8 + chroma, hexagon ME with half/quarter-sample refinement, 16x16/16x8/8x16/8x8, P_Skip
- B-frames (1-3, spatial direct/B_Skip, L0/L1/Bi), variance AQ, scene-cut IDR
- 4x4/8x8 transforms, deadzone quant, decimation; normative deblocking in the loop
- CABAC (init tables checked against the spec text) and CAVLC
- Rate control: QP, CRF-like, 1-pass ABR + VBV, closed-loop 2-pass VBR
- Slice-parallel encoding with rayon, concurrent B-frames, parallel lookahead/hpel
- ffmpeg-oracle tests: zero decode errors, exact frame counts, decoder output bit-identical to reconstruction

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Brandon Thomas 2026-09-30 09:32:20 -04:00
parent 8194dac6a8
commit 7103b58df9
37 changed files with 7899 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
Cargo.lock generated
View File

@ -1242,6 +1242,15 @@ dependencies = [
"serde",
]
[[package]]
name = "filmcraft-h264enc"
version = "0.1.0"
dependencies = [
"filmcraft-bitstream",
"rayon",
"thiserror 2.0.21",
]
[[package]]
name = "filmcraft-isobmff"
version = "0.1.0"

1
Cargo.toml
View File

@ -17,6 +17,7 @@ filmcraft-geom = { path = "crates/geom" }
filmcraft-color = { path = "crates/color" }
filmcraft-bitstream = { path = "crates/bitstream" }
filmcraft-isobmff = { path = "crates/isobmff" }
filmcraft-h264enc = { path = "crates/h264enc" }
filmcraft-frame = { path = "crates/frame" }
filmcraft-media = { path = "crates/media" }
filmcraft-codecs = { path = "crates/codecs" }

19
crates/h264enc/Cargo.toml Normal file
View File

@ -0,0 +1,19 @@
[package]
name = "filmcraft-h264enc"
description = "Clean-room pure-Rust H.264/AVC encoder (High/Main/Baseline, CABAC/CAVLC, B-frames, rate control)"
version.workspace = true
edition.workspace = true
license.workspace = true
rust-version.workspace = true
[lints]
workspace = true
[features]
default = ["threads"]
threads = ["dep:rayon"]
[dependencies]
filmcraft-bitstream = { workspace = true }
thiserror = { workspace = true }
rayon = { workspace = true, optional = true }

83
crates/h264enc/README.md Normal file
View File

@ -0,0 +1,83 @@
# filmcraft-h264enc
A clean-room, pure-Rust H.264/AVC encoder for FilmCraft's export path, written from the public
ITU-T H.264 specification (04/2013 edition) and general encoder literature. Layer L0: it depends only on
`filmcraft-bitstream`, `thiserror` and (optionally, feature `threads`, on by default) `rayon`. No `unsafe`,
builds for `wasm32-unknown-unknown` with or without `threads`.
## Usage
```rust
use filmcraft_h264enc::{Encoder, EncoderConfig, PacketFormat, Preset, Profile, RateControl, YuvFrame};
let mut cfg = EncoderConfig::new(1920, 1080, 30000, 1001);
cfg.profile = Profile::High;
cfg.preset = Preset::Balanced; // Speed / Balanced / Quality
cfg.rate = RateControl::Crf(20.0); // Crf(f32) | Cbr { kbps } | Vbr { target_kbps, max_kbps } | Qp(u8)
cfg.bframes = 2; // 0..=3
cfg.keyint = 250;
cfg.format = PacketFormat::LengthPrefixed; // MP4 samples; or AnnexB (default) for .h264
let mut enc = Encoder::new(cfg)?;
let avcc = enc.avcc(); // AVCDecoderConfigurationRecord for the MP4 `avcC` box
for (i, pic) in pictures.iter().enumerate() {
let frame = YuvFrame { y: &pic.y, u: &pic.u, v: &pic.v, y_stride: w, uv_stride: w / 2 };
for p in enc.encode(&frame, i as i64) { /* p.data, p.pts, p.dts, p.keyframe, p.frame_type */ }
}
for p in enc.flush() { /* ... */ }
```
- Input: planar 4:2:0 8-bit (`YuvFrame` with strides); any even size, non-multiple-of-16 sizes are
cropped via the SPS frame cropping. `try_encode` returns an error instead of panicking on short planes.
- Output: one `Packet` per picture in decoding order with `pts` (as given), `dts` (reorder-delayed, always
`<= pts` and strictly increasing), `keyframe` (IDR) and the frame type/QP. Annex-B packets carry an AUD,
SPS/PPS on every IDR and an identification SEI on the first picture.
- Two-pass: encode once with `cfg.pass = Pass::First`, take `enc.pass_stats()` (serialisable with
`PassStats::to_text/from_text`), then encode again with `Pass::Second(stats)`.
- `set_recon_capture(true)` + `take_recon()` expose the encoder's reconstructed pictures (used by the tests to
prove encoder/decoder match).
## Features
| Area | Implemented |
|---|---|
| Profiles | Constrained Baseline (CAVLC, I/P), Main (CABAC, I/P/B), High (CABAC, 8x8 transform, I/P/B); 4:2:0 8-bit, frame coding |
| Parameter sets | SPS (level auto-selected from size/rate/DPB/bitrate, cropping, VUI: SAR, colour description + range, timing, bitstream restriction/reorder), PPS, AUD, user-data SEI, avcC |
| Intra | 16x16 (4 modes, SATD), 4x4 (9 modes, SATD + mode-cost, sequential reconstruction), 8x8 (9 modes with reference filtering, High), chroma (4 modes) |
| Inter P | hexagon integer search from predictors, half- then quarter-sample refinement (SATD) on precomputed 6-tap half-sample planes; 16x16/16x8/8x16/8x8 partitions (Balanced/Quality); P_Skip detection with the normative skip MV |
| Inter B | 1–3 non-reference B-frames, spatial direct / B_Skip, L0/L1/Bi 16x16 with separate searches; all B-frames of a mini-GOP are encoded concurrently |
| Residual | 4x4 / 8x8 integer transforms, Intra16x16 and chroma DC Hadamards, deadzone quantisation (intra 1/3, inter 1/6), adaptive 4x4/8x8 inter transform choice, coefficient decimation for inter blocks |
| Entropy | CABAC (all I and cabac_init_idc 0 context tables, verified against the spec text) and CAVLC (all coeff_token/total_zeros/run_before tables, level escapes) |
| Loop filter | normative deblocking (bS 0–4 incl. bi-predicted rules, 8x8-transform edges, chroma QP mapping) in the reconstruction loop |
| Mode decision | SATD + λ·bits for Speed/Balanced; Quality adds RD-lite (SSD + λ·CABAC-estimated bits) between the best intra/inter/skip candidates |
| Rate control | `Qp`, `Crf` (constant quality with I/P/B offsets), 1-pass ABR with complexity model + VBV (`Cbr`, 1-pass `Vbr`), closed-loop 2-pass VBR; variance adaptive quantisation (`aq_strength`), lookahead scene-cut IDR insertion |
| Threading | slice-parallel encoding (MB-row slices, `disable_deblocking_filter_idc = 2` with >1 slice so each slice is also deblocked in parallel), parallel half-sample interpolation and lookahead |
## Verification
`cargo test --release -p filmcraft-h264enc` (ffmpeg/ffprobe used only as external oracles; tests skip when
absent). Every stream is decoded with `ffmpeg -ec 0 -f rawvideo` and must report no errors, produce exactly the
number of frames encoded, and be **bit-identical to the encoder's reconstruction** — this covers Baseline/Main/High,
all presets, 0–3 B-frames, 1–4 slices, odd sizes, `testsrc2` and `mandelbrot` sources, and length-prefixed output
re-wrapped with the avcC parameter sets. B-frame display order is checked with ffprobe.
Results on the synthetic content (moving gradient, grain noise, text-like edges, sub-pixel panning):
| Test | Result |
|---|---|
| 1080p, High, Speed, CRF 20, 2 B | mean luma PSNR 43.0 dB (IDR 39.1, P/B 43.2–44.5) |
| 640x360 CBR 1500 kbit/s | 1518 kbit/s (+1.2%) |
| 640x360 1-pass VBR 1000 kbit/s | 981 kbit/s (−1.9%) |
| 640x360 2-pass VBR 800 kbit/s | 840 kbit/s (+5.0%) |
Development tool: `cargo run --release -p filmcraft-h264enc --example h264enc_synth -- W H FRAMES
[profile] [preset] [bframes] [qpN|crfN|cbrN|vbrN] [slices] [synth|testsrc2|mandelbrot]` encodes, decodes with
ffmpeg and reports speed, bitrate, PSNR and any encoder/decoder mismatch (`MAP=1` prints a per-MB mismatch map,
`PERF=k` cycles k source pictures for speed runs).
## Limitations / not yet implemented
- Single reference frame per list; no B-pyramid, no weighted prediction, no interlaced/MBAFF, no 4:2:2/4:4:4 or
high bit depth, no custom scaling matrices, no sub-8x8 partitions, no B 16x8/8x16 partitions.
- CAVLC is used only for Baseline (High/Main always use CABAC); 8x8 transform is therefore CABAC-only.
- No trellis quantisation, no psy-RD; no HRD/buffering-period SEI (VBV is enforced internally only).
- Threading is slice-based only (no frame-level pipelining); with many slices, slice edges are not deblocked.

154
crates/h264enc/examples/h264enc_synth.rs Normal file
View File

@ -0,0 +1,154 @@
//! Development tool: encode synthetic (or lavfi) content and verify against ffmpeg.
//!
//! `cargo run --release -p filmcraft-h264enc --example h264enc_synth -- W H FRAMES [profile] [preset] [bframes] [rate] [slices] [src]`
//! profile: baseline|main|high, preset: speed|balanced|quality, rate: qpN | crfN | cbrN | vbrN
#[path = "../tests/common/mod.rs"]
mod common;
use filmcraft_h264enc::*;
use std::time::Instant;
fn main() {
let a: Vec<String> = std::env::args().collect();
let arg = |i: usize, d: &str| a.get(i).cloned().unwrap_or_else(|| d.to_string());
let w: usize = arg(1, "320").parse().unwrap();
let h: usize = arg(2, "240").parse().unwrap();
let n: usize = arg(3, "10").parse().unwrap();
let mut cfg = EncoderConfig::new(w as u32, h as u32, 30, 1);
cfg.profile = match arg(4, "high").as_str() {
"baseline" => Profile::Baseline,
"main" => Profile::Main,
_ => Profile::High,
};
cfg.preset = match arg(5, "balanced").as_str() {
"speed" => Preset::Speed,
"quality" => Preset::Quality,
_ => Preset::Balanced,
};
cfg.bframes = arg(6, "0").parse().unwrap();
let rate = arg(7, "qp26");
cfg.rate = if let Some(q) = rate.strip_prefix("qp") {
RateControl::Qp(q.parse().unwrap())
} else if let Some(q) = rate.strip_prefix("crf") {
RateControl::Crf(q.parse().unwrap())
} else if let Some(q) = rate.strip_prefix("cbr") {
RateControl::Cbr { kbps: q.parse().unwrap() }
} else {
let k: u32 = rate.strip_prefix("vbr").unwrap().parse().unwrap();
RateControl::Vbr { target_kbps: k, max_kbps: k * 2 }
};
cfg.slices = arg(8, "0").parse().unwrap();
let src_kind = arg(9, "synth");
if let Ok(t) = std::env::var("THREADS") {
cfg.threads = t.parse().unwrap();
}
if let Ok(k) = std::env::var("KEYINT") {
cfg.keyint = k.parse().unwrap();
}
if let Ok(k) = std::env::var("AQ") {
cfg.aq_strength = k.parse().unwrap();
}
// PERF=k: generate only k distinct pictures and cycle them (for long speed runs without verification).
let distinct = std::env::var("PERF").ok().map_or(n, |k| k.parse::<usize>().unwrap().min(n));
let base: Vec<common::Yuv> =
if src_kind == "synth" { (0..distinct).map(|t| common::synth(w, h, t)).collect() } else { common::ffmpeg_source(&src_kind, w, h, distinct) };
let perf = distinct < n;
let frames: Vec<&common::Yuv> = (0..n).map(|i| &base[i % distinct]).collect();
let dir = common::out_dir("dev");
let mut enc = Encoder::new(cfg).unwrap();
enc.set_recon_capture(true);
let mut out = Vec::new();
let mut pk = Vec::new();
let t0 = Instant::now();
if perf {
enc.set_recon_capture(false);
}
for (i, f) in frames.iter().enumerate() {
pk.extend(enc.encode(&f.frame(), i as i64));
}
pk.extend(enc.flush());
let el = t0.elapsed().as_secs_f64();
for p in &pk {
out.extend_from_slice(&p.data);
}
let file = dir.join("out.h264");
std::fs::write(&file, &out).unwrap();
let mut recon = enc.take_recon();
recon.sort_by_key(|r| r.pts);
let kbps = out.len() as f64 * 8.0 / (n as f64 / 30.0) / 1000.0;
println!("encoded {n} frames in {el:.3}s = {:.1} fps, {} bytes, {kbps:.1} kbps", n as f64 / el, out.len());
for p in pk.iter().take(12) {
print!("{:?}(pts{} dts{} qp{} {}B) ", p.frame_type, p.pts, p.dts, p.qp, p.data.len());
}
println!();
if common::ffmpeg().is_none() || perf {
return;
}
let (err, dec) = common::ffmpeg_decode(&file);
if !err.is_empty() {
println!("ffmpeg stderr: {err}");
}
let fs = w * h * 3 / 2;
println!("ffmpeg decoded {} frames (expected {n})", dec.len() / fs);
let mut src_all = Vec::new();
let mut psnr_sum = 0.0;
for (i, r) in recon.iter().enumerate() {
let mut rb = r.y.clone();
rb.extend_from_slice(&r.u);
rb.extend_from_slice(&r.v);
let s = frames[i].bytes();
src_all.extend_from_slice(&s);
let p = common::psnr(&frames[i].y, &r.y);
psnr_sum += p;
if dec.len() >= (i + 1) * fs {
let d = &dec[i * fs..(i + 1) * fs];
if d != rb.as_slice() {
let pos = d.iter().zip(&rb).position(|(a, b)| a != b).unwrap();
let (plane, off) = if pos < w * h {
("Y", pos)
} else if pos < w * h * 5 / 4 {
("U", pos - w * h)
} else {
("V", pos - w * h * 5 / 4)
};
let pw = if plane == "Y" { w } else { w / 2 };
let (x, y) = (off % pw, off / pw);
let mbsz = if plane == "Y" { 16 } else { 8 };
let nd = d.iter().zip(&rb).filter(|(a, b)| a != b).count();
println!(
"frame {i} (pts {}): MISMATCH first at {plane} ({x},{y}) mb ({},{}) dec={} rec={} ({nd} samples differ)",
r.pts,
x / mbsz,
y / mbsz,
d[pos],
rb[pos]
);
if std::env::var("MAP").is_ok() {
let (mbw, mbh) = (w.div_ceil(16), h.div_ceil(16));
for my in 0..mbh {
let mut line = String::new();
for mx in 0..mbw {
let mut n = 0;
for yy in my * 16..(my * 16 + 13).min(h) {
for xx in mx * 16..(mx * 16 + 13).min(w) {
n += (d[yy * w + xx] != rb[yy * w + xx]) as usize;
}
}
line.push(if n == 0 {
'.'
} else if n < 16 {
'e'
} else {
'X'
});
}
println!(" {line}");
}
break;
}
}
}
}
println!("mean luma PSNR vs source: {:.2} dB", psnr_sum / recon.len() as f64);
std::fs::write(dir.join("src.yuv"), &src_all).unwrap();
}

339
crates/h264enc/src/cabac.rs Normal file
View File

@ -0,0 +1,339 @@
//! CABAC arithmetic encoder (ITU-T H.264 §9.3.4.2) and a matching bit-cost estimator.
use crate::tables::{NUM_CTX, RANGE_LPS, TRANS_LPS};
/// A sink for CABAC bins. Implemented by the real arithmetic coder and by the rate estimator, so the same
/// syntax-element code drives both bitstream writing and RD cost estimation.
pub trait BinSink {
fn decision(&mut self, ctx: usize, bin: u32);
fn bypass(&mut self, bin: u32);
fn terminate(&mut self, bin: u32);
}
/// Context states: `pStateIdx << 1 | valMPS`.
#[derive(Clone)]
pub struct Contexts(pub [u8; NUM_CTX]);
impl Contexts {
pub fn init(table: &[(i8, i8); NUM_CTX], slice_qp: i32) -> Self {
let qp = slice_qp.clamp(0, 51);
let mut s = [0u8; NUM_CTX];
for (i, &(m, n)) in table.iter().enumerate() {
let pre = (((m as i32) * qp) >> 4) + n as i32;
let pre = pre.clamp(1, 126);
s[i] = if pre <= 63 { ((63 - pre) as u8) << 1 } else { (((pre - 64) as u8) << 1) | 1 };
}
Contexts(s)
}
}
/// The arithmetic encoder writing into a byte buffer.
#[derive(Clone)]
pub struct CabacEncoder {
pub ctx: Contexts,
low: u32,
range: u32,
outstanding: u32,
first_bit: bool,
buf: Vec<u8>,
acc: u64,
nbits: u32,
bins: u64,
}
impl CabacEncoder {
/// Start encoding after a slice header; `header` must be byte aligned (cabac_alignment_one_bits written).
pub fn new(ctx: Contexts, header: Vec<u8>) -> Self {
CabacEncoder { ctx, low: 0, range: 510, outstanding: 0, first_bit: true, buf: header, acc: 0, nbits: 0, bins: 0 }
}
#[inline]
fn write_bit_raw(&mut self, b: u32) {
self.acc = (self.acc << 1) | b as u64;
self.nbits += 1;
if self.nbits == 32 {
self.buf.extend_from_slice(&(self.acc as u32).to_be_bytes());
self.acc = 0;
self.nbits = 0;
}
}
#[inline]
fn put_bit(&mut self, b: u32) {
if self.first_bit {
self.first_bit = false;
} else {
self.write_bit_raw(b);
}
while self.outstanding > 0 {
self.write_bit_raw(1 - b);
self.outstanding -= 1;
}
}
#[inline]
fn renorm(&mut self) {
while self.range < 256 {
if self.low < 256 {
self.put_bit(0);
} else if self.low >= 512 {
self.low -= 512;
self.put_bit(1);
} else {
self.low -= 256;
self.outstanding += 1;
}
self.range <<= 1;
self.low <<= 1;
}
}
/// Terminate with end_of_slice_flag = 1 and flush; returns the RBSP bytes (with rbsp trailing bits).
pub fn finish(mut self) -> Vec<u8> {
self.terminate(1);
// EncodeFlush wrote the stop bit as the last bit; pad with zeros to byte alignment.
while !self.nbits.is_multiple_of(8) {
self.write_bit_raw(0);
}
let nbytes = self.nbits / 8;
for i in 0..nbytes {
self.buf.push((self.acc >> (8 * (nbytes - 1 - i))) as u8);
}
self.buf
}
}
impl BinSink for CabacEncoder {
#[inline]
fn decision(&mut self, ctx: usize, bin: u32) {
self.bins += 1;
let s = self.ctx.0[ctx];
let state = (s >> 1) as usize;
let mps = (s & 1) as u32;
let lps = RANGE_LPS[state][((self.range >> 6) & 3) as usize] as u32;
self.range -= lps;
if bin != mps {
self.low += self.range;
self.range = lps;
let new_mps = if state == 0 { 1 - mps } else { mps };
self.ctx.0[ctx] = (TRANS_LPS[state] << 1) | new_mps as u8;
} else if state < 62 {
self.ctx.0[ctx] = (((state + 1) as u8) << 1) | mps as u8;
}
self.renorm();
}
#[inline]
fn bypass(&mut self, bin: u32) {
self.bins += 1;
self.low <<= 1;
if bin != 0 {
self.low += self.range;
}
if self.low >= 1024 {
self.put_bit(1);
self.low -= 1024;
} else if self.low < 512 {
self.put_bit(0);
} else {
self.low -= 512;
self.outstanding += 1;
}
}
fn terminate(&mut self, bin: u32) {
self.bins += 1;
self.range -= 2;
if bin != 0 {
self.low += self.range;
// EncodeFlush
self.range = 2;
self.renorm();
self.put_bit((self.low >> 9) & 1);
let v = ((self.low >> 7) & 3) | 1;
self.write_bit_raw((v >> 1) & 1);
self.write_bit_raw(v & 1);
} else {
self.renorm();
}
}
}
/// Fractional-bit cost table: `COST[state][is_mps]` in 1/256 bit units.
pub struct CostTable {
pub cost: [[u32; 2]; 64],
}
impl CostTable {
pub fn new() -> Self {
let mut cost = [[0u32; 2]; 64];
// Probability of the LPS for each state, estimated from rangeTabLPS relative to a mid-range interval.
for (s, c) in cost.iter_mut().enumerate() {
let avg: f64 = RANGE_LPS[s].iter().map(|&v| v as f64).sum::<f64>() / 4.0;
let p_lps = (avg / 384.0).clamp(1e-4, 0.5);
c[0] = (-(p_lps.log2()) * 256.0).round() as u32;
c[1] = (-((1.0 - p_lps).log2()) * 256.0).round() as u32;
}
CostTable { cost }
}
}
impl Default for CostTable {
fn default() -> Self {
Self::new()
}
}
/// Rate estimator: tracks context states like the encoder and accumulates cost in 1/256 bits.
#[derive(Clone)]
pub struct CabacEstimator<'a> {
pub ctx: Contexts,
pub cost: u64,
table: &'a CostTable,
}
impl<'a> CabacEstimator<'a> {
pub fn new(ctx: Contexts, table: &'a CostTable) -> Self {
CabacEstimator { ctx, cost: 0, table }
}
}
impl BinSink for CabacEstimator<'_> {
#[inline]
fn decision(&mut self, ctx: usize, bin: u32) {
let s = self.ctx.0[ctx];
let state = (s >> 1) as usize;
let mps = (s & 1) as u32;
if bin != mps {
self.cost += self.table.cost[state][0] as u64;
let new_mps = if state == 0 { 1 - mps } else { mps };
self.ctx.0[ctx] = (TRANS_LPS[state] << 1) | new_mps as u8;
} else {
self.cost += self.table.cost[state][1] as u64;
if state < 62 {
self.ctx.0[ctx] = (((state + 1) as u8) << 1) | mps as u8;
}
}
}
#[inline]
fn bypass(&mut self, _bin: u32) {
self.cost += 256;
}
fn terminate(&mut self, bin: u32) {
self.cost += if bin != 0 { 7 * 256 } else { 2 };
}
}
#[cfg(test)]
mod tests {
use super::*;
/// A reference CABAC decoder (§9.3.1.2 / 9.3.3.2) used only to round-trip-test the encoder.
struct Dec<'a> {
data: &'a [u8],
pos: usize,
range: u32,
offset: u32,
ctx: Contexts,
}
impl<'a> Dec<'a> {
fn bit(&mut self) -> u32 {
let b = if self.pos / 8 < self.data.len() { (self.data[self.pos / 8] >> (7 - self.pos % 8)) & 1 } else { 0 };
self.pos += 1;
b as u32
}
fn new(data: &'a [u8], ctx: Contexts) -> Self {
let mut d = Dec { data, pos: 0, range: 510, offset: 0, ctx };
for _ in 0..9 {
d.offset = (d.offset << 1) | d.bit();
}
d
}
fn decision(&mut self, c: usize) -> u32 {
let s = self.ctx.0[c];
let state = (s >> 1) as usize;
let mps = (s & 1) as u32;
let lps = RANGE_LPS[state][((self.range >> 6) & 3) as usize] as u32;
self.range -= lps;
let bin;
if self.offset >= self.range {
bin = 1 - mps;
self.offset -= self.range;
self.range = lps;
let new_mps = if state == 0 { 1 - mps } else { mps };
self.ctx.0[c] = (TRANS_LPS[state] << 1) | new_mps as u8;
} else {
bin = mps;
if state < 62 {
self.ctx.0[c] = (((state + 1) as u8) << 1) | mps as u8;
}
}
while self.range < 256 {
self.range <<= 1;
self.offset = (self.offset << 1) | self.bit();
}
bin
}
fn bypass(&mut self) -> u32 {
self.offset = (self.offset << 1) | self.bit();
if self.offset >= self.range {
self.offset -= self.range;
1
} else {
0
}
}
fn terminate(&mut self) -> u32 {
self.range -= 2;
if self.offset >= self.range {
1
} else {
while self.range < 256 {
self.range <<= 1;
self.offset = (self.offset << 1) | self.bit();
}
0
}
}
}
#[test]
fn roundtrip_random_bins() {
let table = crate::tables::cabac_init_pb0();
let ctx = Contexts::init(&table, 26);
let mut enc = CabacEncoder::new(ctx.clone(), Vec::new());
let mut seed = 12345u32;
let mut rnd = || {
seed = seed.wrapping_mul(1664525).wrapping_add(1013904223);
seed >> 8
};
let mut ops = Vec::new();
for _ in 0..20000 {
let kind = rnd() % 10;
let c = (rnd() % 400) as usize;
// skewed bins so contexts adapt
let bin = if rnd() % 100 < 80 { (c % 2) as u32 } else { 1 - (c % 2) as u32 };
if kind < 7 {
enc.decision(c, bin);
ops.push((0, c, bin));
} else if kind < 9 {
enc.bypass(bin);
ops.push((1, 0, bin));
} else {
enc.terminate(0);
ops.push((2, 0, 0));
}
}
let data = enc.finish();
let mut dec = Dec::new(&data, ctx);
for (i, &(k, c, bin)) in ops.iter().enumerate() {
let got = match k {
0 => dec.decision(c),
1 => dec.bypass(),
_ => dec.terminate(),
};
assert_eq!(got, bin, "op {i}");
}
assert_eq!(dec.terminate(), 1);
}
}

411
crates/h264enc/src/cabac_mb.rs Normal file
View File

@ -0,0 +1,411 @@
//! CABAC binarisation and context selection for macroblock-layer syntax elements (§9.3.2, §9.3.3.1).
use crate::cabac::BinSink;
use crate::mbinfo::{MbInfo, MbKind, Part};
use crate::nal::SliceType;
use crate::syntax::{MbCode, NbCtx};
use crate::tables::{BLK4_XY, LAST8X8, SIG8X8_FRAME};
pub fn write_skip_flag<S: BinSink>(s: &mut S, nb: &NbCtx, skip: bool) {
let cond = |m: Option<&MbInfo>| m.is_some_and(|m| !m.kind.is_skip()) as usize;
let base = if nb.slice_type == SliceType::B { 24 } else { 11 };
s.decision(base + cond(nb.left) + cond(nb.top), skip as u32);
}
fn eg_bypass<S: BinSink>(s: &mut S, mut v: u32, mut k: u32) {
loop {
if v >= (1 << k) {
s.bypass(1);
v -= 1 << k;
k += 1;
} else {
s.bypass(0);
while k > 0 {
k -= 1;
s.bypass((v >> k) & 1);
}
break;
}
}
}
/// I16x16 / I_NxN mb_type bins. `off` = ctxIdxOffset of the (suffix) table; `i_slice` selects I-slice contexts.
fn write_intra_type<S: BinSink>(s: &mut S, nb: &NbCtx, c: &MbCode) {
let i_slice = nb.slice_type == SliceType::I;
let (off, b0ctx) = match nb.slice_type {
SliceType::I => {
let cond = |m: Option<&MbInfo>| m.is_some_and(|m| !m.kind.is_inxn()) as usize;
(3, 3 + cond(nb.left) + cond(nb.top))
}
SliceType::P => {
s.decision(14, 1); // prefix: intra
(17, 17)
}
SliceType::B => {
// prefix 111101
write_b_bins(s, nb, &[1, 1, 1, 1, 0, 1]);
(32, 32)
}
};
if c.kind != MbKind::I16x16 {
s.decision(b0ctx, 0);
return;
}
s.decision(b0ctx, 1);
s.terminate(0);
let luma = (c.cbp & 15 != 0) as u32;
let chroma = (c.cbp >> 4) as u32;
let pred = c.i16_mode as u32;
if i_slice {
s.decision(off + 3, luma);
s.decision(off + 4, (chroma != 0) as u32);
if chroma != 0 {
s.decision(off + 5, (chroma == 2) as u32);
s.decision(off + 6, pred >> 1);
s.decision(off + 7, pred & 1);
} else {
s.decision(off + 6, pred >> 1);
s.decision(off + 7, pred & 1);
}
} else {
s.decision(off + 1, luma);
s.decision(off + 2, (chroma != 0) as u32);
if chroma != 0 {
s.decision(off + 2, (chroma == 2) as u32);
}
s.decision(off + 3, pred >> 1);
s.decision(off + 3, pred & 1);
}
}
fn write_b_bins<S: BinSink>(s: &mut S, nb: &NbCtx, bins: &[u32]) {
for (i, &b) in bins.iter().enumerate() {
let ctx = match i {
0 => {
let cond = |m: Option<&MbInfo>| m.is_some_and(|m| !matches!(m.kind, MbKind::BSkip | MbKind::BDirect)) as usize;
27 + cond(nb.left) + cond(nb.top)
}
1 => 27 + 3,
2 => 27 + if bins[1] != 0 { 4 } else { 5 },
_ => 27 + 5,
};
s.decision(ctx, b);
}
}
const B_BINS: [&[u32]; 23] = [
&[0],
&[1, 0, 0],
&[1, 0, 1],
&[1, 1, 0, 0, 0, 0],
&[1, 1, 0, 0, 0, 1],
&[1, 1, 0, 0, 1, 0],
&[1, 1, 0, 0, 1, 1],
&[1, 1, 0, 1, 0, 0],
&[1, 1, 0, 1, 0, 1],
&[1, 1, 0, 1, 1, 0],
&[1, 1, 0, 1, 1, 1],
&[1, 1, 1, 1, 1, 0],
&[1, 1, 1, 0, 0, 0, 0],
&[1, 1, 1, 0, 0, 0, 1],
&[1, 1, 1, 0, 0, 1, 0],
&[1, 1, 1, 0, 0, 1, 1],
&[1, 1, 1, 0, 1, 0, 0],
&[1, 1, 1, 0, 1, 0, 1],
&[1, 1, 1, 0, 1, 1, 0],
&[1, 1, 1, 0, 1, 1, 1],
&[1, 1, 1, 1, 0, 0, 0],
&[1, 1, 1, 1, 0, 0, 1],
&[1, 1, 1, 1, 1, 1],
];
fn write_mvd<S: BinSink>(s: &mut S, comp: usize, v: i32, sum: u32) {
let base = if comp == 0 { 40 } else { 47 };
let a = v.unsigned_abs();
let inc = if sum < 3 {
0
} else if sum > 32 {
2
} else {
1
};
let u = a.min(9);
s.decision(base + inc, (u > 0) as u32);
if u > 0 {
for k in 1..9u32 {
let ctx = base + [3, 4, 5, 6, 6, 6, 6, 6][(k - 1) as usize];
if k < u {
s.decision(ctx, 1);
} else {
s.decision(ctx, 0);
break;
}
}
if a >= 9 {
eg_bypass(s, a - 9, 3);
}
s.bypass((v < 0) as u32);
}
}
/// Sum of |mvd| of neighbours A and B for the 4x4 block at raster `r` (current MB info already holds earlier partitions).
fn mvd_sum(nb: &NbCtx, r: usize, list: usize, comp: usize) -> u32 {
let (x, y) = (r & 3, r >> 2);
let a = if x > 0 { nb.cur.mvd[list][r - 1][comp] } else { nb.left.map_or(0, |m| m.mvd[list][y * 4 + 3][comp]) };
let b = if y > 0 { nb.cur.mvd[list][r - 4][comp] } else { nb.top.map_or(0, |m| m.mvd[list][12 + x][comp]) };
a as u32 + b as u32
}
fn write_cbp<S: BinSink>(s: &mut S, nb: &NbCtx, cbp: u8) {
let luma_bit = |m: &MbInfo, b8: usize| (m.cbp >> b8) & 1;
for b8 in 0..4usize {
// (A, B) neighbour 8x8 blocks: (Some(mb), idx) where None = current MB
let cond = |m: Option<&MbInfo>, idx: usize, in_cur: bool| -> usize {
if in_cur {
((cbp >> idx) & 1 == 0) as usize
} else {
match m {
None => 0,
Some(m) => (luma_bit(m, idx) == 0) as usize,
}
}
};
let (ca, cb) = match b8 {
0 => (cond(nb.left, 1, false), cond(nb.top, 2, false)),
1 => (cond(None, 0, true), cond(nb.top, 3, false)),
2 => (cond(nb.left, 3, false), cond(None, 0, true)),
_ => (cond(None, 2, true), cond(None, 1, true)),
};
s.decision(73 + ca + 2 * cb, ((cbp >> b8) & 1) as u32);
}
let chroma = cbp >> 4;
let cc = |m: Option<&MbInfo>, two: bool| -> usize {
m.is_some_and(|m| {
let c = m.cbp >> 4;
!m.kind.is_skip() && if two { c == 2 } else { c != 0 }
}) as usize
};
s.decision(77 + cc(nb.left, false) + 2 * cc(nb.top, false), (chroma != 0) as u32);
if chroma != 0 {
s.decision(77 + 4 + cc(nb.left, true) + 2 * cc(nb.top, true), (chroma == 2) as u32);
}
}
#[derive(Clone, Copy, PartialEq, Eq)]
pub enum Cat {
LumaDc = 0,
LumaAc = 1,
Luma4x4 = 2,
ChromaDc = 3,
ChromaAc = 4,
Luma8x8 = 5,
}
/// Encode one residual block; `coeffs` are the block's coefficients in scan order (length = maxNumCoeff).
pub fn write_residual<S: BinSink>(s: &mut S, cat: Cat, coeffs: &[i16], cbf_inc: u32) {
const CBF_OFF: [usize; 5] = [0, 4, 8, 12, 16];
const SIG_OFF: [usize; 5] = [0, 15, 29, 44, 47];
const ABS_OFF: [usize; 5] = [0, 10, 20, 30, 39];
let n = coeffs.len();
let last = coeffs.iter().rposition(|&c| c != 0);
let ci = cat as usize;
if cat != Cat::Luma8x8 {
s.decision(85 + CBF_OFF[ci] + cbf_inc as usize, last.is_some() as u32);
}
let Some(last) = last else { return };
let (sig_base, last_base, abs_base) = if cat == Cat::Luma8x8 { (402, 417, 426) } else { (105 + SIG_OFF[ci], 166 + SIG_OFF[ci], 227 + ABS_OFF[ci]) };
for i in 0..n - 1 {
let (si, li) = match cat {
Cat::ChromaDc => (i.min(2), i.min(2)),
Cat::Luma8x8 => (SIG8X8_FRAME[i] as usize, LAST8X8[i] as usize),
_ => (i, i),
};
let sig = coeffs[i] != 0;
s.decision(sig_base + si, sig as u32);
if sig {
let is_last = i == last;
s.decision(last_base + li, is_last as u32);
if is_last {
break;
}
}
}
let mut gt1 = 0usize;
let mut eq1 = 0usize;
let max_gt1 = if cat == Cat::ChromaDc { 3 } else { 4 };
for i in (0..=last).rev() {
let c = coeffs[i];
if c == 0 {
continue;
}
let a = c.unsigned_abs() as u32;
let v = a - 1;
let inc0 = if gt1 != 0 { 0 } else { (1 + eq1).min(4) };
s.decision(abs_base + inc0, (v > 0) as u32);
if v > 0 {
let ctx = abs_base + 5 + gt1.min(max_gt1);
let u = v.min(14);
for _ in 1..u {
s.decision(ctx, 1);
}
if u < 14 {
s.decision(ctx, 0);
} else {
eg_bypass(s, v - 14, 0);
}
}
s.bypass((c < 0) as u32);
if a == 1 {
eq1 += 1;
} else {
gt1 += 1;
}
}
}
fn cbf_luma(nb: &NbCtx, raster: usize) -> u32 {
let f = |v: Option<u8>| v.map_or(nb.unavail_cbf(), |n| (n != 0) as u32);
f(nb.luma_nnz(raster, true)) + 2 * f(nb.luma_nnz(raster, false))
}
fn cbf_chroma_ac(nb: &NbCtx, c: usize, b: usize) -> u32 {
let f = |v: Option<u8>| v.map_or(nb.unavail_cbf(), |n| (n != 0) as u32);
f(nb.chroma_nnz(c, b, true)) + 2 * f(nb.chroma_nnz(c, b, false))
}
fn cbf_dc(nb: &NbCtx, bit: u8) -> u32 {
let f = |m: Option<&MbInfo>| m.map_or(nb.unavail_cbf(), |m| ((m.dc_cbf >> bit) & 1) as u32);
f(nb.left) + 2 * f(nb.top)
}
/// Write the whole macroblock layer (after mb_skip_flag = 0 in P/B slices).
pub fn write_mb<S: BinSink>(s: &mut S, nb: &NbCtx, c: &MbCode) {
let cur = nb.cur;
// ---- mb_type
match c.kind {
MbKind::I4x4 | MbKind::I8x8 | MbKind::I16x16 => write_intra_type(s, nb, c),
MbKind::PInter => {
let bins: [u32; 3] = match c.part {
Part::P16x16 => [0, 0, 0],
Part::P16x8 => [0, 1, 1],
Part::P8x16 => [0, 1, 0],
Part::P8x8 => [0, 0, 1],
};
s.decision(14, bins[0]);
s.decision(15, bins[1]);
s.decision(if bins[1] != 1 { 16 } else { 17 }, bins[2]);
}
MbKind::BDirect | MbKind::BInter => write_b_bins(s, nb, B_BINS[c.b_type_number() as usize]),
_ => unreachable!("skip/none macroblocks are not written with write_mb"),
}
let is_inxn = c.kind.is_inxn();
if c.kind == MbKind::PInter && c.part == Part::P8x8 {
for _ in 0..4 {
s.decision(21, 1); // P_L0_8x8
}
}
// ---- mb_pred
if is_inxn {
if nb.t8x8_mode {
let cond = |m: Option<&MbInfo>| m.is_some_and(|m| m.t8x8) as usize;
s.decision(399 + cond(nb.left) + cond(nb.top), c.t8x8 as u32);
}
let n = if c.t8x8 { 4 } else { 16 };
for i in 0..n {
let p = c.ipred[i];
if p < 0 {
s.decision(68, 1);
} else {
s.decision(68, 0);
s.decision(69, (p & 1) as u32);
s.decision(69, ((p >> 1) & 1) as u32);
s.decision(69, ((p >> 2) & 1) as u32);
}
}
}
if c.kind.is_intra() {
let cond = |m: Option<&MbInfo>| m.is_some_and(|m| m.kind.is_intra() && m.chroma_mode != 0) as usize;
let m = c.chroma_mode as u32;
s.decision(64 + cond(nb.left) + cond(nb.top), (m > 0) as u32);
if m > 0 {
s.decision(64 + 3, (m > 1) as u32);
if m > 1 {
s.decision(64 + 3, (m > 2) as u32);
}
}
}
if matches!(c.kind, MbKind::PInter | MbKind::BInter) {
for list in 0..2 {
for p in 0..c.part.count() {
if !c.uses_list(p, list) {
continue;
}
let (px, py, _, _) = c.part.rect(p);
let r = (py / 4) * 4 + px / 4;
let mvd = c.mvd[list][p];
write_mvd(s, 0, mvd.x as i32, mvd_sum(nb, r, list, 0));
write_mvd(s, 1, mvd.y as i32, mvd_sum(nb, r, list, 1));
}
}
}
// ---- cbp, transform size, qp delta
if c.kind != MbKind::I16x16 {
write_cbp(s, nb, c.cbp);
if c.cbp & 15 != 0 && nb.t8x8_mode && !is_inxn {
let cond = |m: Option<&MbInfo>| m.is_some_and(|m| m.t8x8) as usize;
s.decision(399 + cond(nb.left) + cond(nb.top), c.t8x8 as u32);
}
}
if c.cbp == 0 && c.kind != MbKind::I16x16 {
return;
}
{
let v = c.qp_delta;
let m = if v > 0 { 2 * v as u32 - 1 } else { (-2 * v) as u32 };
s.decision(60 + nb.prev_dqp_nz as usize, (m > 0) as u32);
if m > 0 {
for k in 1..=m {
s.decision(if k == 1 { 62 } else { 63 }, (k < m) as u32);
}
}
}
// ---- residual
if c.kind == MbKind::I16x16 {
write_residual(s, Cat::LumaDc, &c.luma_dc, cbf_dc(nb, 0));
if c.cbp & 15 != 0 {
for blk in 0..16 {
let (x, y) = BLK4_XY[blk];
write_residual(s, Cat::LumaAc, &c.luma4[blk][1..16], cbf_luma(nb, y * 4 + x));
}
}
} else {
for b8 in 0..4 {
if (c.cbp >> b8) & 1 == 0 {
continue;
}
if c.t8x8 {
write_residual(s, Cat::Luma8x8, &c.luma8[b8], 0);
} else {
for i in 0..4 {
let blk = b8 * 4 + i;
let (x, y) = BLK4_XY[blk];
write_residual(s, Cat::Luma4x4, &c.luma4[blk], cbf_luma(nb, y * 4 + x));
}
}
}
}
let chroma = c.cbp >> 4;
if chroma != 0 {
for comp in 0..2 {
write_residual(s, Cat::ChromaDc, &c.cdc[comp], cbf_dc(nb, 1 + comp as u8));
}
}
if chroma == 2 {
for comp in 0..2 {
for b in 0..4 {
write_residual(s, Cat::ChromaAc, &c.cac[comp * 4 + b][1..16], cbf_chroma_ac(nb, comp, b));
}
}
}
let _ = cur;
}

262
crates/h264enc/src/cavlc.rs Normal file
View File

@ -0,0 +1,262 @@
//! CAVLC residual coding (§9.2) and macroblock-layer syntax for Baseline streams.
use filmcraft_bitstream::BitWriter;
use crate::mbinfo::{MbKind, Part};
use crate::nal::SliceType;
use crate::syntax::{MbCode, NbCtx};
use crate::tables::{
BLK4_XY, CBP_INTER_TO_CODE, CBP_INTRA_TO_CODE, COEFF_TOKEN_CDC_CODE, COEFF_TOKEN_CDC_LEN, COEFF_TOKEN_CODE, COEFF_TOKEN_LEN, RUN_BEFORE, TOTAL_ZEROS,
TOTAL_ZEROS_CDC,
};
/// Largest level magnitude representable with level_prefix <= 15 in every suffixLength state.
pub const MAX_LEVEL: i32 = 2047;
/// Write one CAVLC residual block. `nc` = -1 for chroma DC. Returns TotalCoeff.
pub fn write_residual(w: &mut BitWriter, coeffs: &[i16], nc: i32) -> u8 {
let max = coeffs.len();
// levels from highest frequency to lowest
let mut levels = [0i32; 16];
let mut pos = [0usize; 16];
let mut total = 0usize;
for i in (0..max).rev() {
if coeffs[i] != 0 {
levels[total] = coeffs[i] as i32;
pos[total] = i;
total += 1;
}
}
let mut t1 = 0usize;
while t1 < total && t1 < 3 && levels[t1].abs() == 1 {
t1 += 1;
}
// coeff_token
if nc == -1 {
w.write_bits(COEFF_TOKEN_CDC_CODE[total][t1] as u32, COEFF_TOKEN_CDC_LEN[total][t1] as u32);
} else if nc >= 8 {
if total == 0 {
w.write_bits(3, 6);
} else {
w.write_bits((((total - 1) << 2) | t1) as u32, 6);
}
} else {
let t = if nc < 2 {
0
} else if nc < 4 {
1
} else {
2
};
w.write_bits(COEFF_TOKEN_CODE[t][total][t1] as u32, COEFF_TOKEN_LEN[t][total][t1] as u32);
}
if total == 0 {
return 0;
}
let mut suffix_len: u32 = if total > 10 && t1 < 3 { 1 } else { 0 };
for i in 0..total {
let l = levels[i];
if i < t1 {
w.write_bit(l < 0);
continue;
}
let mut code = if l > 0 { 2 * l - 2 } else { -2 * l - 1 };
if i == t1 && t1 < 3 {
code -= 2;
}
let code = code as u32;
if suffix_len == 0 {
if code < 14 {
w.write_bits(1, code + 1);
} else if code < 30 {
w.write_bits(1, 15);
w.write_bits(code - 14, 4);
} else {
write_escape(w, code - 30);
}
} else if code < (15 << suffix_len) {
w.write_bits(1, (code >> suffix_len) + 1);
w.write_bits(code & ((1 << suffix_len) - 1), suffix_len);
} else {
write_escape(w, code - (15 << suffix_len));
}
if suffix_len == 0 {
suffix_len = 1;
}
if l.unsigned_abs() > (3 << (suffix_len - 1)) && suffix_len < 6 {
suffix_len += 1;
}
}
let total_zeros = pos[0] + 1 - total;
if total < max {
let (len, code) = if nc == -1 { TOTAL_ZEROS_CDC[total - 1][total_zeros] } else { TOTAL_ZEROS[total - 1][total_zeros] };
w.write_bits(code as u32, len as u32);
}
let mut zeros_left = total_zeros;
for i in 0..total - 1 {
if zeros_left == 0 {
break;
}
let run = pos[i] - pos[i + 1] - 1;
let (len, code) = RUN_BEFORE[zeros_left.min(7) - 1][run];
w.write_bits(code as u32, len as u32);
zeros_left -= run;
}
total as u8
}
/// level_prefix >= 15 escape; `rem` = levelCode minus the prefix-15 base.
fn write_escape(w: &mut BitWriter, rem: u32) {
if rem < 4096 {
w.write_bits(1, 16);
w.write_bits(rem, 12);
} else {
// level_prefix p >= 16 (High profiles): suffix of p-3 bits, value rem + 4096 - 2^(p-3)
let mut p = 16;
while rem + 4096 >= (1 << (p - 3)) * 2 {
p += 1;
}
w.write_bits(0, p);
w.write_bit(true);
w.write_bits(rem + 4096 - (1 << (p - 3)), p - 3);
}
}
fn nc_of(a: Option<u8>, b: Option<u8>) -> i32 {
match (a, b) {
(Some(a), Some(b)) => (a as i32 + b as i32 + 1) >> 1,
(Some(a), None) => a as i32,
(None, Some(b)) => b as i32,
_ => 0,
}
}
/// Write the macroblock layer in CAVLC. For P/B slices the caller writes mb_skip_run first.
pub fn write_mb(w: &mut BitWriter, nb: &NbCtx, c: &MbCode) {
let intra_base = match nb.slice_type {
SliceType::I => 0,
SliceType::P => 5,
SliceType::B => 23,
};
match c.kind {
MbKind::I4x4 | MbKind::I8x8 => w.write_ue(intra_base),
MbKind::I16x16 => {
let t = 1 + c.i16_mode as u32 + 4 * (c.cbp >> 4) as u32 + if c.cbp & 15 != 0 { 12 } else { 0 };
w.write_ue(intra_base + t);
}
MbKind::PInter => w.write_ue(match c.part {
Part::P16x16 => 0,
Part::P16x8 => 1,
Part::P8x16 => 2,
Part::P8x8 => 3,
}),
MbKind::BDirect | MbKind::BInter => w.write_ue(c.b_type_number()),
_ => unreachable!(),
}
if c.kind == MbKind::PInter && c.part == Part::P8x8 {
for _ in 0..4 {
w.write_ue(0);
}
}
if c.kind.is_inxn() {
if nb.t8x8_mode {
w.write_bit(c.t8x8);
}
let n = if c.t8x8 { 4 } else { 16 };
for i in 0..n {
let p = c.ipred[i];
if p < 0 {
w.write_bit(true);
} else {
w.write_bit(false);
w.write_bits(p as u32, 3);
}
}
}
if c.kind.is_intra() {
w.write_ue(c.chroma_mode as u32);
}
if matches!(c.kind, MbKind::PInter | MbKind::BInter) {
for list in 0..2 {
for p in 0..c.part.count() {
if c.uses_list(p, list) {
w.write_se(c.mvd[list][p].x as i32);
w.write_se(c.mvd[list][p].y as i32);
}
}
}
}
if c.kind != MbKind::I16x16 {
let code = if c.kind.is_intra() { CBP_INTRA_TO_CODE[c.cbp as usize] } else { CBP_INTER_TO_CODE[c.cbp as usize] };
let cbp_idx = (c.cbp & 15) | ((c.cbp >> 4) << 4);
debug_assert_eq!(cbp_idx, c.cbp);
w.write_ue(code as u32);
if c.cbp & 15 != 0 && nb.t8x8_mode && !c.kind.is_inxn() {
w.write_bit(c.t8x8);
}
}
if c.cbp == 0 && c.kind != MbKind::I16x16 {
return;
}
w.write_se(c.qp_delta);
let luma_nc = |r: usize| nc_of(nb.luma_nnz(r, true), nb.luma_nnz(r, false));
if c.kind == MbKind::I16x16 {
write_residual(w, &c.luma_dc, luma_nc(0));
if c.cbp & 15 != 0 {
for blk in 0..16 {
let (x, y) = BLK4_XY[blk];
write_residual(w, &c.luma4[blk][1..16], luma_nc(y * 4 + x));
}
}
} else {
for b8 in 0..4 {
if (c.cbp >> b8) & 1 == 0 {
continue;
}
for i in 0..4 {
let blk = b8 * 4 + i;
let (x, y) = BLK4_XY[blk];
if c.t8x8 {
let mut sub = [0i16; 16];
for k in 0..16 {
sub[k] = c.luma8[b8][4 * k + i];
}
write_residual(w, &sub, luma_nc(y * 4 + x));
} else {
write_residual(w, &c.luma4[blk], luma_nc(y * 4 + x));
}
}
}
}
let chroma = c.cbp >> 4;
if chroma != 0 {
for comp in 0..2 {
write_residual(w, &c.cdc[comp], -1);
}
}
if chroma == 2 {
for comp in 0..2 {
for b in 0..4 {
let ncv = nc_of(nb.chroma_nnz(comp, b, true), nb.chroma_nnz(comp, b, false));
write_residual(w, &c.cac[comp * 4 + b][1..16], ncv);
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn simple_block_matches_spec_example() {
// scan: 0,3,-1,0,0,-1,1,0,1 -> TotalCoeff 5, T1s 3.
let coeffs: [i16; 16] = [0, 3, -1, 0, 0, -1, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0];
let mut w = BitWriter::new();
let t = write_residual(&mut w, &coeffs, 0);
assert_eq!(t, 5);
// coeff_token 0000100 | signs 001 | level -1: 01 | level 3: 001 0 | total_zeros(4) 110 | runs 10 11 01 1
let bits: String = w.finish().iter().map(|b| format!("{b:08b}")).collect();
assert!(bits.starts_with("00001000010100101101011011"), "{bits}");
}
}

251
crates/h264enc/src/deblock.rs Normal file
View File

@ -0,0 +1,251 @@
//! In-loop deblocking filter (§8.7) for frame-coded 4:2:0 pictures without MBAFF.
//!
//! Operates on a band of macroblock rows that begins at a slice boundary; the top edge of the band is
//! never filtered (slices are MB-row aligned and either there is one slice, or
//! `disable_deblocking_filter_idc = 2` is used).
use crate::mbinfo::{MbInfo, Mv};
use crate::tables::{ALPHA, BETA, TC0, chroma_qp};
pub struct DeblockParams {
pub alpha_offset: i32,
pub beta_offset: i32,
pub chroma_qp_offset: i32,
/// Picture identity for (list, ref_idx 0).
pub ref_ids: [i32; 2],
}
/// A mutable band of rows of a padded plane: sample (x, y) with y relative to the band top.
pub struct BandPlane<'a> {
pub data: &'a mut [u8],
pub stride: usize,
pub pad: usize,
}
impl BandPlane<'_> {
#[inline]
fn idx(&self, x: usize, y: usize) -> usize {
y * self.stride + self.pad + x
}
}
#[inline]
fn nz8(m: &MbInfo, blk: usize) -> bool {
if m.t8x8 {
let (x, y) = ((blk & 3) >> 1, (blk >> 2) >> 1);
let b = y * 8 + x * 2;
m.nnz[b] | m.nnz[b + 1] | m.nnz[b + 4] | m.nnz[b + 5] != 0
} else {
m.nnz[blk] != 0
}
}
#[inline]
fn mvdiff(a: Mv, b: Mv) -> bool {
(a.x as i32 - b.x as i32).abs() >= 4 || (a.y as i32 - b.y as i32).abs() >= 4
}
fn bs_inter(p: &MbInfo, ip: usize, q: &MbInfo, iq: usize, ids: &[i32; 2]) -> u8 {
if nz8(p, ip) || nz8(q, iq) {
return 2;
}
let rp = [p.ref_at(0, ip), p.ref_at(1, ip)];
let rq = [q.ref_at(0, iq), q.ref_at(1, iq)];
// Only ref_idx 0 is ever used per list, so the list index identifies the picture.
let pic = |r: i8, l: usize| if r >= 0 { ids[l] } else { -1 };
let pp = [pic(rp[0], 0), pic(rp[1], 1)];
let pq = [pic(rq[0], 0), pic(rq[1], 1)];
let np = (pp[0] >= 0) as u8 + (pp[1] >= 0) as u8;
let nq = (pq[0] >= 0) as u8 + (pq[1] >= 0) as u8;
if np != nq {
return 1;
}
if np == 1 {
let (a, ma) = if pp[0] >= 0 { (pp[0], p.mv[0][ip]) } else { (pp[1], p.mv[1][ip]) };
let (b, mb) = if pq[0] >= 0 { (pq[0], q.mv[0][iq]) } else { (pq[1], q.mv[1][iq]) };
if a != b {
return 1;
}
return mvdiff(ma, mb) as u8;
}
if np == 0 {
return 0;
}
// two motion vectors each
let (mp0, mp1, mq0, mq1) = (p.mv[0][ip], p.mv[1][ip], q.mv[0][iq], q.mv[1][iq]);
let same_set = (pp[0] == pq[0] && pp[1] == pq[1]) || (pp[0] == pq[1] && pp[1] == pq[0]);
if !same_set {
return 1;
}
if pp[0] != pp[1] {
// two different pictures: compare mvs referring to the same picture
if pp[0] == pq[0] { (mvdiff(mp0, mq0) || mvdiff(mp1, mq1)) as u8 } else { (mvdiff(mp0, mq1) || mvdiff(mp1, mq0)) as u8 }
} else {
((mvdiff(mp0, mq0) || mvdiff(mp1, mq1)) && (mvdiff(mp0, mq1) || mvdiff(mp1, mq0))) as u8
}
}
#[inline]
fn bs_of(p: &MbInfo, ip: usize, q: &MbInfo, iq: usize, mb_edge: bool, ids: &[i32; 2]) -> u8 {
if p.kind.is_intra() || q.kind.is_intra() {
return if mb_edge { 4 } else { 3 };
}
bs_inter(p, ip, q, iq, ids)
}
#[inline]
fn filter_luma(d: &mut [u8], q0i: usize, step: usize, bs: u8, alpha: i32, beta: i32, tc0: i32) {
let p0 = d[q0i - step] as i32;
let q0 = d[q0i] as i32;
let p1 = d[q0i - 2 * step] as i32;
let q1 = d[q0i + step] as i32;
if (p0 - q0).abs() >= alpha || (p1 - p0).abs() >= beta || (q1 - q0).abs() >= beta {
return;
}
let p2 = d[q0i - 3 * step] as i32;
let q2 = d[q0i + 2 * step] as i32;
let ap = (p2 - p0).abs();
let aq = (q2 - q0).abs();
if bs < 4 {
let tc = tc0 + (ap < beta) as i32 + (aq < beta) as i32;
let delta = ((((q0 - p0) << 2) + (p1 - q1) + 4) >> 3).clamp(-tc, tc);
d[q0i - step] = (p0 + delta).clamp(0, 255) as u8;
d[q0i] = (q0 - delta).clamp(0, 255) as u8;
if ap < beta {
d[q0i - 2 * step] = (p1 + ((p2 + ((p0 + q0 + 1) >> 1) - (p1 << 1)) >> 1).clamp(-tc0, tc0)) as u8;
}
if aq < beta {
d[q0i + step] = (q1 + ((q2 + ((p0 + q0 + 1) >> 1) - (q1 << 1)) >> 1).clamp(-tc0, tc0)) as u8;
}
} else {
let p3 = d[q0i - 4 * step] as i32;
let q3 = d[q0i + 3 * step] as i32;
let strong = (p0 - q0).abs() < ((alpha >> 2) + 2);
if ap < beta && strong {
d[q0i - step] = ((p2 + 2 * p1 + 2 * p0 + 2 * q0 + q1 + 4) >> 3) as u8;
d[q0i - 2 * step] = ((p2 + p1 + p0 + q0 + 2) >> 2) as u8;
d[q0i - 3 * step] = ((2 * p3 + 3 * p2 + p1 + p0 + q0 + 4) >> 3) as u8;
} else {
d[q0i - step] = ((2 * p1 + p0 + q1 + 2) >> 2) as u8;
}
if aq < beta && strong {
d[q0i] = ((p1 + 2 * p0 + 2 * q0 + 2 * q1 + q2 + 4) >> 3) as u8;
d[q0i + step] = ((p0 + q0 + q1 + q2 + 2) >> 2) as u8;
d[q0i + 2 * step] = ((2 * q3 + 3 * q2 + q1 + q0 + p0 + 4) >> 3) as u8;
} else {
d[q0i] = ((2 * q1 + q0 + p1 + 2) >> 2) as u8;
}
}
}
#[inline]
fn filter_chroma(d: &mut [u8], q0i: usize, step: usize, bs: u8, alpha: i32, beta: i32, tc0: i32) {
let p0 = d[q0i - step] as i32;
let q0 = d[q0i] as i32;
let p1 = d[q0i - 2 * step] as i32;
let q1 = d[q0i + step] as i32;
if (p0 - q0).abs() >= alpha || (p1 - p0).abs() >= beta || (q1 - q0).abs() >= beta {
return;
}
if bs < 4 {
let tc = tc0 + 1;
let delta = ((((q0 - p0) << 2) + (p1 - q1) + 4) >> 3).clamp(-tc, tc);
d[q0i - step] = (p0 + delta).clamp(0, 255) as u8;
d[q0i] = (q0 - delta).clamp(0, 255) as u8;
} else {
d[q0i - step] = ((2 * p1 + p0 + q1 + 2) >> 2) as u8;
d[q0i] = ((2 * q1 + q0 + p1 + 2) >> 2) as u8;
}
}
struct EdgeQ {
alpha: i32,
beta: i32,
idx_a: usize,
}
fn edge_q(qp_p: i32, qp_q: i32, dp: &DeblockParams) -> EdgeQ {
let qav = (qp_p + qp_q + 1) >> 1;
let ia = (qav + dp.alpha_offset).clamp(0, 51) as usize;
let ib = (qav + dp.beta_offset).clamp(0, 51) as usize;
EdgeQ { alpha: ALPHA[ia] as i32, beta: BETA[ib] as i32, idx_a: ia }
}
/// Deblock `rows` macroblock rows of a band. `mbs` holds the band's MbInfo (row-major, `mbw` wide).
pub fn deblock_band<'a>(y: &mut BandPlane<'a>, u: &mut BandPlane<'a>, v: &mut BandPlane<'a>, mbs: &[MbInfo], mbw: usize, rows: usize, dp: &DeblockParams) {
let ids = &dp.ref_ids;
for my in 0..rows {
for mx in 0..mbw {
let q = &mbs[my * mbw + mx];
let cqp = |m: &MbInfo| chroma_qp(m.qp as i32 + dp.chroma_qp_offset) as i32;
// ---- vertical edges (dir 0) then horizontal (dir 1)
for dir in 0..2 {
let mut bs_all = [[0u8; 4]; 4];
for e in 0..4 {
if e == 0 && ((dir == 0 && mx == 0) || (dir == 1 && my == 0)) {
continue;
}
if e & 1 == 1 && q.t8x8 {
continue;
}
let p_mb = if e == 0 { if dir == 0 { &mbs[my * mbw + mx - 1] } else { &mbs[(my - 1) * mbw + mx] } } else { q };
let mut bs = [0u8; 4];
for s in 0..4 {
let (iq, ip) = if dir == 0 {
let iq = s * 4 + e;
(iq, if e == 0 { s * 4 + 3 } else { iq - 1 })
} else {
let iq = e * 4 + s;
(iq, if e == 0 { 12 + s } else { iq - 4 })
};
bs[s] = bs_of(p_mb, ip, q, iq, e == 0, ids);
}
bs_all[e] = bs;
if bs == [0; 4] {
continue;
}
let eq = edge_q(p_mb.qp as i32, q.qp as i32, dp);
if eq.alpha == 0 || eq.beta == 0 {
continue;
}
for k in 0..16 {
let b = bs[k >> 2];
if b == 0 {
continue;
}
let tc0 = if b < 4 { TC0[eq.idx_a][b as usize - 1] as i32 } else { 0 };
let (qi, step) = if dir == 0 { (y.idx(mx * 16 + e * 4, my * 16 + k), 1) } else { (y.idx(mx * 16 + k, my * 16 + e * 4), y.stride) };
filter_luma(y.data, qi, step, b, eq.alpha, eq.beta, tc0);
}
}
// chroma edges 0 and 2 (chroma sample offsets 0 and 4)
for e in [0usize, 2] {
if e == 0 && ((dir == 0 && mx == 0) || (dir == 1 && my == 0)) {
continue;
}
let bs = bs_all[e];
if bs == [0; 4] {
continue;
}
let p_mb = if e == 0 { if dir == 0 { &mbs[my * mbw + mx - 1] } else { &mbs[(my - 1) * mbw + mx] } } else { q };
let eq = edge_q(cqp(p_mb), cqp(q), dp);
if eq.alpha == 0 || eq.beta == 0 {
continue;
}
for plane in [&mut *u, &mut *v] {
for k in 0..8 {
let b = bs[k >> 1];
if b == 0 {
continue;
}
let tc0 = if b < 4 { TC0[eq.idx_a][b as usize - 1] as i32 } else { 0 };
let (qi, step) =
if dir == 0 { (plane.idx(mx * 8 + e * 2, my * 8 + k), 1) } else { (plane.idx(mx * 8 + k, my * 8 + e * 2), plane.stride) };
filter_chroma(plane.data, qi, step, b, eq.alpha, eq.beta, tc0);
}
}
}
}
}
}
}

172
crates/h264enc/src/dsp.rs Normal file
View File

@ -0,0 +1,172 @@
//! Distortion metrics: SAD, SATD (4x4 Hadamard), SSD.
#[inline]
pub fn sad(a: &[u8], sa: usize, b: &[u8], sb: usize, w: usize, h: usize) -> u32 {
let mut s = 0u32;
for r in 0..h {
let ra = &a[r * sa..r * sa + w];
let rb = &b[r * sb..r * sb + w];
s += ra.iter().zip(rb).map(|(&x, &y)| x.abs_diff(y) as u32).sum::<u32>();
}
s
}
/// SAD of a 16-wide block (fast path).
#[inline]
pub fn sad16(a: &[u8], sa: usize, b: &[u8], sb: usize, h: usize) -> u32 {
let mut s = 0u32;
for r in 0..h {
let ra: &[u8; 16] = a[r * sa..r * sa + 16].try_into().unwrap();
let rb: &[u8; 16] = b[r * sb..r * sb + 16].try_into().unwrap();
let mut acc = 0u16;
for i in 0..16 {
acc += ra[i].abs_diff(rb[i]) as u16;
}
s += acc as u32;
}
s
}
#[cfg(test)]
fn hadamard4x4_sum(d: &[i32; 16]) -> u32 {
let mut t = [0i32; 16];
for i in 0..4 {
let a0 = d[i * 4] + d[i * 4 + 1];
let a1 = d[i * 4] - d[i * 4 + 1];
let a2 = d[i * 4 + 2] + d[i * 4 + 3];
let a3 = d[i * 4 + 2] - d[i * 4 + 3];
t[i * 4] = a0 + a2;
t[i * 4 + 1] = a1 + a3;
t[i * 4 + 2] = a0 - a2;
t[i * 4 + 3] = a1 - a3;
}
let mut s = 0u32;
for i in 0..4 {
let a0 = t[i] + t[4 + i];
let a1 = t[i] - t[4 + i];
let a2 = t[8 + i] + t[12 + i];
let a3 = t[8 + i] - t[12 + i];
s += (a0 + a2).unsigned_abs() + (a1 + a3).unsigned_abs() + (a0 - a2).unsigned_abs() + (a1 - a3).unsigned_abs();
}
s
}
/// SATD of a strip of `N` columns (N/4 side-by-side 4x4 blocks) and 4 rows, computed lane-wise so the compiler
/// can vectorise it. Uses |p+q| + |p-q| = 2 max(|p|, |q|) for the last butterfly, which also absorbs the /2.
#[inline(always)]
fn satd_strip<const N: usize, const H: usize>(a: &[u8], sa: usize, b: &[u8], sb: usize) -> u32 {
let mut d = [[0i16; N]; 4];
for r in 0..4 {
let ra = &a[r * sa..r * sa + N];
let rb = &b[r * sb..r * sb + N];
for i in 0..N {
d[r][i] = ra[i] as i16 - rb[i] as i16;
}
}
let mut v = [[0i16; N]; 4];
for i in 0..N {
let s01 = d[0][i] + d[1][i];
let d01 = d[0][i] - d[1][i];
let s23 = d[2][i] + d[3][i];
let d23 = d[2][i] - d[3][i];
v[0][i] = s01 + s23;
v[1][i] = s01 - s23;
v[2][i] = d01 + d23;
v[3][i] = d01 - d23;
}
let mut total = 0u32;
for row in &v {
let mut hs = [0i16; H];
let mut hd = [0i16; H];
for k in 0..H {
hs[k] = row[2 * k] + row[2 * k + 1];
hd[k] = row[2 * k] - row[2 * k + 1];
}
for k in 0..H / 2 {
total += hs[2 * k].unsigned_abs().max(hs[2 * k + 1].unsigned_abs()) as u32;
total += hd[2 * k].unsigned_abs().max(hd[2 * k + 1].unsigned_abs()) as u32;
}
}
total
}
/// Sum of absolute 4x4 Hadamard-transformed differences (halved), for w, h multiples of 4.
pub fn satd(a: &[u8], sa: usize, b: &[u8], sb: usize, w: usize, h: usize) -> u32 {
let mut total = 0;
for by in (0..h).step_by(4) {
let (ao, bo) = (by * sa, by * sb);
match w {
16 => total += satd_strip::<16, 8>(&a[ao..], sa, &b[bo..], sb),
8 => total += satd_strip::<8, 4>(&a[ao..], sa, &b[bo..], sb),
4 => total += satd_strip::<4, 2>(&a[ao..], sa, &b[bo..], sb),
_ => {
for bx in (0..w).step_by(4) {
total += satd_strip::<4, 2>(&a[ao + bx..], sa, &b[bo + bx..], sb);
}
}
}
}
total
}
pub fn ssd(a: &[u8], sa: usize, b: &[u8], sb: usize, w: usize, h: usize) -> u64 {
let mut s = 0u64;
for r in 0..h {
let ra = &a[r * sa..r * sa + w];
let rb = &b[r * sb..r * sb + w];
s += ra.iter().zip(rb).map(|(&x, &y)| (x as i32 - y as i32).pow(2) as u32).sum::<u32>() as u64;
}
s
}
/// Variance*N (sum of squares minus square of sum / N) of a block.
pub fn block_var(a: &[u8], sa: usize, w: usize, h: usize) -> u32 {
let mut s = 0u32;
let mut ss = 0u32;
for r in 0..h {
for &v in &a[r * sa..r * sa + w] {
s += v as u32;
ss += (v as u32) * (v as u32);
}
}
ss - ((s as u64 * s as u64) / (w * h) as u64) as u32
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn satd_of_dc_difference() {
let a = [10u8; 16];
let b = [0u8; 16];
// DC-only difference of 10: hadamard DC = 160, halved = 80.
assert_eq!(satd(&a, 4, &b, 4, 4, 4), 80);
assert_eq!(sad(&a, 4, &b, 4, 4, 4), 160);
}
#[test]
fn satd_matches_reference() {
let mut seed = 7u32;
let mut rnd = || {
seed = seed.wrapping_mul(1103515245).wrapping_add(12345);
(seed >> 16) as u8
};
let a: Vec<u8> = (0..16 * 16).map(|_| rnd()).collect();
let b: Vec<u8> = (0..16 * 16).map(|_| rnd()).collect();
for (w, h) in [(16, 16), (8, 8), (4, 4), (16, 8), (8, 16), (12, 4)] {
let mut want = 0;
for by in (0..h).step_by(4) {
for bx in (0..w).step_by(4) {
let mut d = [0i32; 16];
for r in 0..4 {
for c in 0..4 {
d[r * 4 + c] = a[(by + r) * 16 + bx + c] as i32 - b[(by + r) * 16 + bx + c] as i32;
}
}
want += hadamard4x4_sum(&d);
}
}
assert_eq!(satd(&a, 16, &b, 16, w, h), want / 2, "{w}x{h}");
}
}
}

303
crates/h264enc/src/intra.rs Normal file
View File

@ -0,0 +1,303 @@
//! Intra prediction (§8.3): 4x4, 8x8 (with reference filtering), 16x16 and chroma.
/// Neighbouring samples of a block. `top` holds p[x,-1] for x = 0..2N (top then top-right), `left` p[-1,y].
#[derive(Clone, Copy)]
pub struct Edge<const N2: usize, const N: usize> {
pub top: [u8; N2],
pub left: [u8; N],
pub tl: u8,
pub has_top: bool,
pub has_left: bool,
pub has_tl: bool,
}
pub type Edge4 = Edge<8, 4>;
pub type Edge8 = Edge<16, 8>;
/// Which modes are usable given neighbour availability.
pub fn mode_available(mode: u8, has_top: bool, has_left: bool, has_tl: bool) -> bool {
match mode {
0 | 3 | 7 => has_top,
1 | 8 => has_left,
2 => true,
_ => has_top && has_left && has_tl,
}
}
/// Generic directional prediction for an NxN block given (already filtered, for 8x8) neighbours.
/// `t(x)` = p[x,-1] for x in -1..2N, `l(y)` = p[-1,y] for y in -1..N.
fn pred_dir<const N: usize>(mode: u8, t: &dyn Fn(i32) -> i32, l: &dyn Fn(i32) -> i32, dc: i32, out: &mut [u8]) {
let n = N as i32;
for y in 0..n {
for x in 0..n {
let v = match mode {
0 => t(x),
1 => l(y),
2 => dc,
3 => {
if x == n - 1 && y == n - 1 {
(t(2 * n - 2) + 3 * t(2 * n - 1) + 2) >> 2
} else {
(t(x + y) + 2 * t(x + y + 1) + t(x + y + 2) + 2) >> 2
}
}
4 => {
if x > y {
(t(x - y - 2) + 2 * t(x - y - 1) + t(x - y) + 2) >> 2
} else if x < y {
(l(y - x - 2) + 2 * l(y - x - 1) + l(y - x) + 2) >> 2
} else {
(t(0) + 2 * t(-1) + l(0) + 2) >> 2
}
}
5 => {
let z = 2 * x - y;
if z >= 0 && z & 1 == 0 {
(t(x - (y >> 1) - 1) + t(x - (y >> 1)) + 1) >> 1
} else if z >= 0 {
(t(x - (y >> 1) - 2) + 2 * t(x - (y >> 1) - 1) + t(x - (y >> 1)) + 2) >> 2
} else if z == -1 {
(l(0) + 2 * l(-1) + t(0) + 2) >> 2
} else {
(l(y - 2 * x - 1) + 2 * l(y - 2 * x - 2) + l(y - 2 * x - 3) + 2) >> 2
}
}
6 => {
let z = 2 * y - x;
if z >= 0 && z & 1 == 0 {
(l(y - (x >> 1) - 1) + l(y - (x >> 1)) + 1) >> 1
} else if z >= 0 {
(l(y - (x >> 1) - 2) + 2 * l(y - (x >> 1) - 1) + l(y - (x >> 1)) + 2) >> 2
} else if z == -1 {
(l(0) + 2 * l(-1) + t(0) + 2) >> 2
} else {
(t(x - 2 * y - 1) + 2 * t(x - 2 * y - 2) + t(x - 2 * y - 3) + 2) >> 2
}
}
7 => {
if y & 1 == 0 {
(t(x + (y >> 1)) + t(x + (y >> 1) + 1) + 1) >> 1
} else {
(t(x + (y >> 1)) + 2 * t(x + (y >> 1) + 1) + t(x + (y >> 1) + 2) + 2) >> 2
}
}
_ => {
let z = x + 2 * y;
let lim = 2 * n - 3;
if z < lim && z & 1 == 0 {
(l(y + (x >> 1)) + l(y + (x >> 1) + 1) + 1) >> 1
} else if z < lim {
(l(y + (x >> 1)) + 2 * l(y + (x >> 1) + 1) + l(y + (x >> 1) + 2) + 2) >> 2
} else if z == lim {
(l(n - 2) + 3 * l(n - 1) + 2) >> 2
} else {
l(n - 1)
}
}
};
out[(y * n + x) as usize] = v as u8;
}
}
}
/// Intra 4x4 prediction. `e.top[4..8]` must already be substituted with top[3] when top-right is unavailable.
pub fn pred4x4(mode: u8, e: &Edge4, out: &mut [u8; 16]) {
let dc = match (e.has_top, e.has_left) {
(true, true) => (e.top[..4].iter().map(|&v| v as i32).sum::<i32>() + e.left.iter().map(|&v| v as i32).sum::<i32>() + 4) >> 3,
(true, false) => (e.top[..4].iter().map(|&v| v as i32).sum::<i32>() + 2) >> 2,
(false, true) => (e.left.iter().map(|&v| v as i32).sum::<i32>() + 2) >> 2,
_ => 128,
};
if mode == 2 {
out.fill(dc as u8);
return;
}
let t = |x: i32| if x < 0 { e.tl as i32 } else { e.top[x as usize] as i32 };
let l = |y: i32| if y < 0 { e.tl as i32 } else { e.left[y as usize] as i32 };
pred_dir::<4>(mode, &t, &l, dc, out);
}
/// Reference sample filtering for Intra 8x8 (§8.3.2.2.1). `e.top[8..16]` must be substituted if unavailable.
pub fn filter8(e: &Edge8) -> Edge8 {
let mut f = *e;
if e.has_top {
let p = |x: usize| e.top[x] as u32;
f.top[0] = if e.has_tl { ((e.tl as u32 + 2 * p(0) + p(1) + 2) >> 2) as u8 } else { ((3 * p(0) + p(1) + 2) >> 2) as u8 };
for x in 1..15 {
f.top[x] = ((p(x - 1) + 2 * p(x) + p(x + 1) + 2) >> 2) as u8;
}
f.top[15] = ((p(14) + 3 * p(15) + 2) >> 2) as u8;
}
if e.has_tl {
let tl = e.tl as u32;
f.tl = if !e.has_top || !e.has_left {
if e.has_top {
((3 * tl + e.top[0] as u32 + 2) >> 2) as u8
} else if e.has_left {
((3 * tl + e.left[0] as u32 + 2) >> 2) as u8
} else {
e.tl
}
} else {
((e.top[0] as u32 + 2 * tl + e.left[0] as u32 + 2) >> 2) as u8
};
}
if e.has_left {
let p = |y: usize| e.left[y] as u32;
f.left[0] = if e.has_tl { ((e.tl as u32 + 2 * p(0) + p(1) + 2) >> 2) as u8 } else { ((3 * p(0) + p(1) + 2) >> 2) as u8 };
for y in 1..7 {
f.left[y] = ((p(y - 1) + 2 * p(y) + p(y + 1) + 2) >> 2) as u8;
}
f.left[7] = ((p(6) + 3 * p(7) + 2) >> 2) as u8;
}
f
}
/// Intra 8x8 prediction from *filtered* neighbours.
pub fn pred8x8(mode: u8, f: &Edge8, out: &mut [u8; 64]) {
let dc = match (f.has_top, f.has_left) {
(true, true) => (f.top[..8].iter().map(|&v| v as i32).sum::<i32>() + f.left.iter().map(|&v| v as i32).sum::<i32>() + 8) >> 4,
(true, false) => (f.top[..8].iter().map(|&v| v as i32).sum::<i32>() + 4) >> 3,
(false, true) => (f.left.iter().map(|&v| v as i32).sum::<i32>() + 4) >> 3,
_ => 128,
};
if mode == 2 {
out.fill(dc as u8);
return;
}
let t = |x: i32| if x < 0 { f.tl as i32 } else { f.top[x as usize] as i32 };
let l = |y: i32| if y < 0 { f.tl as i32 } else { f.left[y as usize] as i32 };
pred_dir::<8>(mode, &t, &l, dc, out);
}
/// Intra 16x16 prediction. Modes: 0 V, 1 H, 2 DC, 3 plane.
pub fn pred16x16(mode: u8, top: &[u8; 16], left: &[u8; 16], tl: u8, has_top: bool, has_left: bool, out: &mut [u8; 256]) {
match mode {
0 => {
for y in 0..16 {
out[y * 16..y * 16 + 16].copy_from_slice(top);
}
}
1 => {
for y in 0..16 {
out[y * 16..y * 16 + 16].fill(left[y]);
}
}
2 => {
let st: u32 = top.iter().map(|&v| v as u32).sum();
let sl: u32 = left.iter().map(|&v| v as u32).sum();
let dc = match (has_top, has_left) {
(true, true) => (st + sl + 16) >> 5,
(true, false) => (st + 8) >> 4,
(false, true) => (sl + 8) >> 4,
_ => 128,
};
out.fill(dc as u8);
}
_ => {
let p_t = |x: i32| if x < 0 { tl as i32 } else { top[x as usize] as i32 };
let p_l = |y: i32| if y < 0 { tl as i32 } else { left[y as usize] as i32 };
let mut h = 0;
let mut v = 0;
for i in 0..8 {
h += (i + 1) * (p_t(8 + i) - p_t(6 - i));
v += (i + 1) * (p_l(8 + i) - p_l(6 - i));
}
let a = 16 * (p_l(15) + p_t(15));
let b = (5 * h + 32) >> 6;
let c = (5 * v + 32) >> 6;
for y in 0..16 {
for x in 0..16 {
out[y * 16 + x] = ((a + b * (x as i32 - 7) + c * (y as i32 - 7) + 16) >> 5).clamp(0, 255) as u8;
}
}
}
}
}
/// Chroma 8x8 prediction (4:2:0). Modes: 0 DC, 1 H, 2 V, 3 plane.
pub fn pred_chroma(mode: u8, top: &[u8; 8], left: &[u8; 8], tl: u8, has_top: bool, has_left: bool, out: &mut [u8; 64]) {
match mode {
0 => {
for by in 0..2 {
for bx in 0..2 {
let st: u32 = top[bx * 4..bx * 4 + 4].iter().map(|&v| v as u32).sum();
let sl: u32 = left[by * 4..by * 4 + 4].iter().map(|&v| v as u32).sum();
let dc = if (bx == 0 && by == 0) || (bx == 1 && by == 1) {
match (has_top, has_left) {
(true, true) => (st + sl + 4) >> 3,
(true, false) => (st + 2) >> 2,
(false, true) => (sl + 2) >> 2,
_ => 128,
}
} else if bx == 1 && by == 0 {
if has_top {
(st + 2) >> 2
} else if has_left {
(sl + 2) >> 2
} else {
128
}
} else if has_left {
(sl + 2) >> 2
} else if has_top {
(st + 2) >> 2
} else {
128
};
for y in 0..4 {
for x in 0..4 {
out[(by * 4 + y) * 8 + bx * 4 + x] = dc as u8;
}
}
}
}
}
1 => {
for y in 0..8 {
out[y * 8..y * 8 + 8].fill(left[y]);
}
}
2 => {
for y in 0..8 {
out[y * 8..y * 8 + 8].copy_from_slice(top);
}
}
_ => {
let p_t = |x: i32| if x < 0 { tl as i32 } else { top[x as usize] as i32 };
let p_l = |y: i32| if y < 0 { tl as i32 } else { left[y as usize] as i32 };
let mut h = 0;
let mut v = 0;
for i in 0..4 {
h += (i + 1) * (p_t(4 + i) - p_t(2 - i));
v += (i + 1) * (p_l(4 + i) - p_l(2 - i));
}
let a = 16 * (p_l(7) + p_t(7));
let b = (34 * h + 32) >> 6;
let c = (34 * v + 32) >> 6;
for y in 0..8 {
for x in 0..8 {
out[y * 8 + x] = ((a + b * (x as i32 - 3) + c * (y as i32 - 3) + 16) >> 5).clamp(0, 255) as u8;
}
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn dc_and_vertical() {
let e = Edge4 { top: [10, 20, 30, 40, 40, 40, 40, 40], left: [1, 2, 3, 4], tl: 5, has_top: true, has_left: true, has_tl: true };
let mut o = [0u8; 16];
pred4x4(0, &e, &mut o);
assert_eq!(&o[12..16], &[10, 20, 30, 40]);
pred4x4(2, &e, &mut o);
assert_eq!(o[0], ((100 + 10 + 4) >> 3) as u8);
// horizontal up saturates at bottom-left
pred4x4(8, &e, &mut o);
assert_eq!(o[15], 4);
}
}

903
crates/h264enc/src/lib.rs Normal file
View File

@ -0,0 +1,903 @@
//! Clean-room, pure-Rust H.264/AVC encoder (ITU-T H.264), written from the public specification.
//!
//! - Profiles: Constrained Baseline (CAVLC, I/P), Main (CABAC, I/P/B) and High (CABAC, 8x8 transform, I/P/B), 4:2:0 8-bit.
//! - Intra 16x16 / 4x4 / 8x8 mode decision, hexagon motion search with quarter-sample refinement,
//! 16x16/16x8/8x16/8x8 partitions, P_Skip, B-frames with spatial direct/B_Skip, in-loop deblocking.
//! - Rate control: constant QP, constant quality (CRF-like), 1-pass ABR with VBV, 2-pass VBR.
//! - Slice-parallel encoding with rayon (feature `threads`, on by default).
//!
//! ```no_run
//! use filmcraft_h264enc::{Encoder, EncoderConfig, RateControl, YuvFrame};
//! let mut cfg = EncoderConfig::new(1280, 720, 30, 1);
//! cfg.rate = RateControl::Crf(20.0);
//! let mut enc = Encoder::new(cfg).unwrap();
//! let (y, u, v) = (vec![128u8; 1280 * 720], vec![128u8; 640 * 360], vec![128u8; 640 * 360]);
//! let frame = YuvFrame { y: &y, u: &u, v: &v, y_stride: 1280, uv_stride: 640 };
//! let mut packets = enc.encode(&frame, 0);
//! packets.extend(enc.flush());
//! ```
mod cabac;
mod cabac_mb;
mod cavlc;
mod deblock;
mod dsp;
mod intra;
mod lookahead;
mod mbinfo;
mod nal;
mod picture;
mod ratecontrol;
mod slice;
mod syntax;
mod tables;
mod transform;
use std::collections::VecDeque;
pub use ratecontrol::{FrameStat, PassStats};
use cabac::CostTable;
use deblock::{BandPlane, DeblockParams, deblock_band};
use lookahead::{Complexity, LowRes};
use mbinfo::MbInfo;
use nal::{NAL_AUD, NAL_IDR, NAL_PPS, NAL_SEI, NAL_SLICE, NAL_SPS, Pps, SliceHeader, SliceType, Sps, Vui};
use picture::{CHROMA_PAD, Frame, LUMA_PAD, RefPic, build_hpel};
use ratecontrol::RcKind;
use slice::{Band, EncParams, FrameEnc, Lambdas, SliceEnc, SliceOut};
use tables::NUM_CTX;
use transform::QuantTables;
/// Encoder errors.
#[derive(Debug, thiserror::Error)]
pub enum Error {
#[error("invalid dimensions {0}x{1} (must be even and non-zero)")]
InvalidDimensions(u32, u32),
#[error("invalid frame rate {0}/{1}")]
InvalidFrameRate(u32, u32),
#[error("invalid configuration: {0}")]
InvalidConfig(String),
#[error("input frame planes too small for the configured size")]
BadFrame,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
pub enum Profile {
/// Constrained Baseline: CAVLC, no B-frames.
Baseline,
/// Main: CABAC, B-frames.
Main,
/// High: CABAC, B-frames, 8x8 transform.
#[default]
High,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
pub enum Preset {
Speed,
#[default]
Balanced,
Quality,
}
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum RateControl {
/// Constant quality; roughly the QP of P frames (0–51, lower is better).
Crf(f32),
/// Constant bitrate target (1-pass ABR with a 1-second VBV at `kbps`).
Cbr { kbps: u32 },
/// Variable bitrate: average `target_kbps`, VBV-limited to `max_kbps`. Use [`Pass`] for 2-pass.
Vbr { target_kbps: u32, max_kbps: u32 },
/// Constant QP (B frames use QP + 2).
Qp(u8),
}
/// Multi-pass mode for bitrate-targeted rate control.
#[derive(Clone, Debug, Default, PartialEq)]
pub enum Pass {
#[default]
Single,
/// First pass: encode and collect [`PassStats`] via [`Encoder::pass_stats`].
First,
/// Second pass using the statistics of the first.
Second(PassStats),
}
/// Colour description written to the VUI (ITU-T H.273 code points).
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct ColorConfig {
pub primaries: u8,
pub transfer: u8,
pub matrix: u8,
pub full_range: bool,
}
impl Default for ColorConfig {
/// BT.709, limited range.
fn default() -> Self {
ColorConfig { primaries: 1, transfer: 1, matrix: 1, full_range: false }
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
pub enum PacketFormat {
/// Start-code prefixed NAL units (`00 00 00 01`), with AUD and in-band SPS/PPS on keyframes.
#[default]
AnnexB,
/// 4-byte big-endian length-prefixed NAL units (MP4 samples); SPS/PPS live in [`Encoder::avcc`].
LengthPrefixed,
}
#[derive(Clone, Debug)]
pub struct EncoderConfig {
pub width: u32,
pub height: u32,
pub fps_num: u32,
pub fps_den: u32,
pub profile: Profile,
pub preset: Preset,
pub rate: RateControl,
/// Maximum distance between IDR frames.
pub keyint: u32,
/// Consecutive B-frames (0–3). Ignored for Baseline.
pub bframes: u8,
pub color: ColorConfig,
/// Worker threads (0 = all available cores). Also bounds the number of slices.
pub threads: usize,
pub pass: Pass,
pub format: PacketFormat,
/// Variance adaptive quantisation strength (0 = off).
pub aq_strength: f32,
/// Sample aspect ratio (1:1 by default).
pub sar: (u16, u16),
/// Insert IDR frames at detected scene cuts.
pub scenecut: bool,
/// Number of slices per picture (0 = derived from `threads`).
pub slices: usize,
/// Access unit delimiters in Annex-B output.
pub aud: bool,
}
impl EncoderConfig {
pub fn new(width: u32, height: u32, fps_num: u32, fps_den: u32) -> Self {
EncoderConfig {
width,
height,
fps_num,
fps_den,
profile: Profile::High,
preset: Preset::Balanced,
rate: RateControl::Crf(23.0),
keyint: 250,
bframes: 2,
color: ColorConfig::default(),
threads: 0,
pass: Pass::Single,
format: PacketFormat::AnnexB,
aq_strength: 0.6,
sar: (1, 1),
scenecut: true,
slices: 0,
aud: true,
}
}
}
/// A planar 4:2:0 8-bit input picture of the configured size.
#[derive(Clone, Copy, Debug)]
pub struct YuvFrame<'a> {
pub y: &'a [u8],
pub u: &'a [u8],
pub v: &'a [u8],
pub y_stride: usize,
pub uv_stride: usize,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum FrameType {
Idr,
I,
P,
B,
}
/// One coded picture (access unit).
#[derive(Clone, Debug)]
pub struct Packet {
pub data: Vec<u8>,
pub pts: i64,
pub dts: i64,
pub keyframe: bool,
pub frame_type: FrameType,
/// QP used for the picture (slice QP).
pub qp: u8,
}
/// A reconstructed (decoded) picture, available when [`Encoder::set_recon_capture`] is on.
#[derive(Clone, Debug)]
pub struct ReconFrame {
pub pts: i64,
pub y: Vec<u8>,
pub u: Vec<u8>,
pub v: Vec<u8>,
}
struct Input {
frame: Frame,
pts: i64,
display: u64,
cplx: Complexity,
aq: Vec<f32>,
}
struct Pool {
#[cfg(feature = "threads")]
pool: Option<rayon::ThreadPool>,
}
impl Pool {
fn install<R: Send>(&self, f: impl FnOnce() -> R + Send) -> R {
#[cfg(feature = "threads")]
if let Some(p) = &self.pool {
return p.install(f);
}
f()
}
}
pub struct Encoder {
cfg: EncoderConfig,
params: EncParams,
sps: Sps,
sps_nal: Vec<u8>,
pps_nal: Vec<u8>,
qt: QuantTables,
cost_tab: CostTable,
lam: Lambdas,
cabac_i: Box<[(i8, i8); NUM_CTX]>,
cabac_pb: Box<[(i8, i8); NUM_CTX]>,
rc: ratecontrol::RateControl,
record_pass1: bool,
pending: VecDeque<Input>,
prev_lowres: Option<LowRes>,
display_count: u64,
last_idr: u64,
idr_count: u32,
next_frame_num: u32,
dpb: Vec<RefPic>,
coded_count: u64,
display_pts: VecDeque<i64>,
dts_offset: Option<i64>,
/// First two presentation times (to derive the frame interval).
all_pts: Vec<i64>,
held: Vec<Packet>,
frame_pool: Vec<Frame>,
hpel_pool: Vec<[picture::Plane; 3]>,
sei_sent: bool,
slices: usize,
pool: Pool,
recon: Option<Vec<ReconFrame>>,
ref_uid: u32,
max_refs: usize,
}
const LOG2_MAX_FRAME_NUM: u32 = 8;
const LOG2_MAX_POC_LSB: u32 = 8;
impl Encoder {
pub fn new(mut cfg: EncoderConfig) -> Result<Self, Error> {
if cfg.width == 0 || cfg.height == 0 || !cfg.width.is_multiple_of(2) || !cfg.height.is_multiple_of(2) || cfg.width > 16384 || cfg.height > 16384 {
return Err(Error::InvalidDimensions(cfg.width, cfg.height));
}
if cfg.fps_num == 0 || cfg.fps_den == 0 {
return Err(Error::InvalidFrameRate(cfg.fps_num, cfg.fps_den));
}
if cfg.bframes > 3 {
return Err(Error::InvalidConfig("bframes must be 0..=3".into()));
}
if cfg.profile == Profile::Baseline {
cfg.bframes = 0;
}
cfg.keyint = cfg.keyint.max(1);
let mbw = cfg.width.div_ceil(16) as usize;
let mbh = cfg.height.div_ceil(16) as usize;
let fps = cfg.fps_num as f64 / cfg.fps_den as f64;
let cabac = cfg.profile != Profile::Baseline;
let high = cfg.profile == Profile::High;
let params = match cfg.preset {
Preset::Speed => EncParams {
mbw,
mbh,
cabac,
t8x8: high,
i4x4: true,
i8x8: false,
partitions: false,
me_range: 8,
subpel: 2,
subpel_rounds: 0,
rd: false,
chroma_qp_offset: 0,
decimate: true,
adaptive_t8: false,
always_intra: false,
cavlc_clamp: !cabac,
},
Preset::Balanced => EncParams {
mbw,
mbh,
cabac,
t8x8: high,
i4x4: true,
i8x8: high,
partitions: true,
me_range: 16,
subpel: 2,
subpel_rounds: 1,
rd: false,
chroma_qp_offset: 0,
decimate: true,
adaptive_t8: high,
always_intra: false,
cavlc_clamp: !cabac,
},
Preset::Quality => EncParams {
mbw,
mbh,
cabac,
t8x8: high,
i4x4: true,
i8x8: high,
partitions: true,
me_range: 24,
subpel: 2,
subpel_rounds: 2,
rd: true,
chroma_qp_offset: 0,
decimate: true,
adaptive_t8: high,
always_intra: true,
cavlc_clamp: !cabac,
},
};
let max_refs = if cfg.bframes > 0 { 2 } else { 1 };
let (kbps_for_level, rc_kind, vbv) = match (&cfg.rate, &cfg.pass) {
(RateControl::Qp(q), _) => (None, RcKind::Qp((*q).min(51)), (None, None)),
(RateControl::Crf(c), _) => (None, RcKind::Crf(c.clamp(0.0, 51.0)), (None, None)),
(RateControl::Cbr { kbps }, Pass::Second(st)) => (Some(*kbps), ratecontrol::RateControl::plan_two_pass(st, *kbps), (Some(*kbps), Some(*kbps))),
(RateControl::Cbr { kbps }, Pass::First) => (Some(*kbps), RcKind::Abr { kbps: *kbps }, (None, None)),
(RateControl::Cbr { kbps }, Pass::Single) => (Some(*kbps), RcKind::Abr { kbps: *kbps }, (Some(*kbps), Some(*kbps))),
(RateControl::Vbr { target_kbps, max_kbps }, Pass::Second(st)) => {
(Some(*max_kbps), ratecontrol::RateControl::plan_two_pass(st, *target_kbps), (Some(*max_kbps), Some(max_kbps * 2)))
}
(RateControl::Vbr { target_kbps, max_kbps }, Pass::First) => (Some(*max_kbps), RcKind::Abr { kbps: *target_kbps }, (None, None)),
(RateControl::Vbr { target_kbps, max_kbps }, Pass::Single) => {
(Some(*max_kbps), RcKind::Abr { kbps: *target_kbps }, (Some(*max_kbps), Some(max_kbps * 2)))
}
};
if matches!(cfg.rate, RateControl::Cbr { kbps: 0 } | RateControl::Vbr { target_kbps: 0, .. }) {
return Err(Error::InvalidConfig("bitrate must be non-zero".into()));
}
let level = nal::pick_level(mbw as u32, mbh as u32, fps, max_refs as u32, kbps_for_level, high);
let (profile_idc, constraint) = match cfg.profile {
Profile::Baseline => (66u8, 0xC0u8), // constraint_set0 + set1 => Constrained Baseline
Profile::Main => (77, 0x40),
Profile::High => (100, 0x00),
};
let reorder = if cfg.bframes > 0 { 1 } else { 0 };
let sps = Sps {
profile_idc,
constraint_flags: constraint,
level_idc: level,
width_mbs: mbw as u32,
height_mbs: mbh as u32,
crop_right: mbw as u32 * 16 - cfg.width,
crop_bottom: mbh as u32 * 16 - cfg.height,
log2_max_frame_num: LOG2_MAX_FRAME_NUM,
log2_max_poc_lsb: LOG2_MAX_POC_LSB,
max_num_ref_frames: max_refs as u32,
vui: Vui {
sar: cfg.sar,
video_full_range: cfg.color.full_range,
colour_primaries: cfg.color.primaries,
transfer: cfg.color.transfer,
matrix: cfg.color.matrix,
num_units_in_tick: cfg.fps_den,
time_scale: cfg.fps_num.saturating_mul(2),
max_num_reorder_frames: reorder,
max_dec_frame_buffering: max_refs as u32,
},
};
let pps = Pps { cabac, pic_init_qp: 26, chroma_qp_offset: params.chroma_qp_offset, transform_8x8: params.t8x8, high };
let sps_nal = nal::nal(3, NAL_SPS, &sps.rbsp());
let pps_nal = nal::nal(3, NAL_PPS, &pps.rbsp());
let threads = if cfg.threads == 0 { std::thread::available_parallelism().map_or(1, |n| n.get()) } else { cfg.threads };
let slices = if cfg.slices > 0 { cfg.slices.min(mbh) } else { threads.min(mbh.div_ceil(4)).max(1) };
#[cfg(feature = "threads")]
let pool = Pool { pool: if threads > 1 { rayon::ThreadPoolBuilder::new().num_threads(threads).build().ok() } else { None } };
#[cfg(not(feature = "threads"))]
let pool = Pool {};
let record_pass1 = cfg.pass == Pass::First;
let rc = ratecontrol::RateControl::new(rc_kind, fps, vbv.0, vbv.1);
Ok(Encoder {
params,
sps,
sps_nal,
pps_nal,
qt: QuantTables::new(),
cost_tab: CostTable::new(),
lam: Lambdas::new(),
cabac_i: Box::new(tables::cabac_init_i()),
cabac_pb: Box::new(tables::cabac_init_pb0()),
rc,
record_pass1,
pending: VecDeque::new(),
prev_lowres: None,
display_count: 0,
last_idr: 0,
idr_count: 0,
next_frame_num: 0,
dpb: Vec::new(),
coded_count: 0,
display_pts: VecDeque::new(),
dts_offset: None,
all_pts: Vec::new(),
held: Vec::new(),
frame_pool: Vec::new(),
hpel_pool: Vec::new(),
sei_sent: false,
slices,
pool,
recon: None,
ref_uid: 0,
max_refs,
cfg,
})
}
pub fn config(&self) -> &EncoderConfig {
&self.cfg
}
/// The SPS and PPS NAL units (without start codes).
pub fn sps_pps(&self) -> (Vec<u8>, Vec<u8>) {
(self.sps_nal.clone(), self.pps_nal.clone())
}
/// AVCDecoderConfigurationRecord (`avcC` box payload) for MP4 muxing (4-byte NAL lengths).
pub fn avcc(&self) -> Vec<u8> {
nal::avcc(&self.sps_nal, &self.pps_nal, self.sps.profile_idc, self.sps.constraint_flags, self.sps.level_idc)
}
/// Number of frames the output is delayed by (B-frame reordering).
pub fn delay(&self) -> usize {
self.cfg.bframes as usize
}
/// Enable capture of reconstructed pictures (for verification).
pub fn set_recon_capture(&mut self, on: bool) {
self.recon = if on { Some(Vec::new()) } else { None };
}
/// Take the reconstructed pictures captured so far (coding order).
pub fn take_recon(&mut self) -> Vec<ReconFrame> {
self.recon.as_mut().map(std::mem::take).unwrap_or_default()
}
/// First-pass statistics (only when configured with [`Pass::First`]).
pub fn pass_stats(&self) -> Option<PassStats> {
if !self.record_pass1 {
return None;
}
Some(PassStats { frames: self.rc.pass1.clone(), fps: self.cfg.fps_num as f64 / self.cfg.fps_den as f64 })
}
/// Encode one picture. Returns zero or more packets in decoding order.
///
/// # Panics
/// If the planes are smaller than the configured size (see [`Encoder::try_encode`]).
pub fn encode(&mut self, frame: &YuvFrame, pts: i64) -> Vec<Packet> {
self.try_encode(frame, pts).expect("invalid input frame")
}
pub fn try_encode(&mut self, frame: &YuvFrame, pts: i64) -> Result<Vec<Packet>, Error> {
let src = self.import(frame)?;
let use_aq = self.cfg.aq_strength > 0.0 && !matches!(self.cfg.rate, RateControl::Qp(_));
let (mbw, mbh, strength) = (self.params.mbw, self.params.mbh, self.cfg.aq_strength);
let prev = self.prev_lowres.take();
let (lowres, cplx, aq) = self.pool.install(|| {
let lowres = LowRes::new(&src);
let cplx = lookahead::complexity(&lowres, prev.as_ref());
let aq = if use_aq { lookahead::aq_offsets(&src, mbw, mbh, strength) } else { Vec::new() };
(lowres, cplx, aq)
});
self.prev_lowres = Some(lowres);
let display = self.display_count;
self.display_count += 1;
self.display_pts.push_back(pts);
if self.all_pts.len() < 2 {
self.all_pts.push(pts);
}
let input = Input { frame: src, pts, display, cplx, aq };
let since = display - self.last_idr;
let fps = self.cfg.fps_num as f64 / self.cfg.fps_den as f64;
let min_gap = (fps / 2.0).max(1.0) as u64;
let scenecut = self.cfg.scenecut && display > 0 && since >= min_gap && cplx.inter as f64 > 0.75 * cplx.intra as f64;
let idr = display == 0 || since >= self.cfg.keyint as u64 || scenecut;
let mut out = Vec::new();
if idr {
self.flush_minigop(&mut out);
self.last_idr = display;
self.encode_batch(vec![(input, SliceType::I, true)], &mut out);
} else {
self.pending.push_back(input);
if self.pending.len() > self.cfg.bframes as usize {
self.flush_minigop(&mut out);
}
}
Ok(self.release(out, false))
}
/// Encode all buffered pictures.
pub fn flush(&mut self) -> Vec<Packet> {
let mut out = Vec::new();
self.flush_minigop(&mut out);
self.release(out, true)
}
/// Assign decoding timestamps. The k-th coded picture gets the k-th smallest presentation time minus the
/// reordering delay (one frame interval when B-frames are enabled); packets are held back until that interval is known.
fn release(&mut self, out: Vec<Packet>, flushing: bool) -> Vec<Packet> {
self.held.extend(out);
if self.dts_offset.is_none() {
if self.cfg.bframes == 0 {
self.dts_offset = Some(0);
} else if self.all_pts.len() >= 2 {
let mut v = self.all_pts.clone();
v.sort_unstable();
self.dts_offset = Some(v[1] - v[0]);
} else if flushing {
self.dts_offset = Some(0);
} else {
return Vec::new();
}
}
let off = self.dts_offset.unwrap_or(0);
let mut res = std::mem::take(&mut self.held);
for p in res.iter_mut() {
p.dts -= off;
}
res
}
fn flush_minigop(&mut self, out: &mut Vec<Packet>) {
let Some(p) = self.pending.pop_back() else { return };
let bs: Vec<(Input, SliceType, bool)> = self.pending.drain(..).map(|b| (b, SliceType::B, false)).collect();
self.encode_batch(vec![(p, SliceType::P, false)], out);
if !bs.is_empty() {
// B-frames are non-reference pictures: encode them all concurrently.
self.encode_batch(bs, out);
}
}
fn import(&mut self, f: &YuvFrame) -> Result<Frame, Error> {
let (w, h) = (self.cfg.width as usize, self.cfg.height as usize);
let (cw, ch) = (w / 2, h / 2);
if f.y_stride < w
|| f.uv_stride < cw
|| f.y.len() < f.y_stride * (h - 1) + w
|| f.u.len() < f.uv_stride * (ch - 1) + cw
|| f.v.len() < f.uv_stride * (ch - 1) + cw
{
return Err(Error::BadFrame);
}
let mut fr = self.new_frame();
for (pl, src, stride, pw, ph) in [(&mut fr.y, f.y, f.y_stride, w, h), (&mut fr.u, f.u, f.uv_stride, cw, ch), (&mut fr.v, f.v, f.uv_stride, cw, ch)] {
let aw = pl.w;
for y in 0..pl.h {
let sy = y.min(ph - 1);
let o = pl.idx(0, y as isize);
pl.data[o..o + pw].copy_from_slice(&src[sy * stride..sy * stride + pw]);
let last = pl.data[o + pw - 1];
pl.data[o + pw..o + aw].fill(last);
}
}
Ok(fr)
}
fn new_frame(&mut self) -> Frame {
self.frame_pool.pop().unwrap_or_else(|| Frame::new(self.params.mbw * 16, self.params.mbh * 16))
}
fn recycle(&mut self, f: Frame) {
if self.frame_pool.len() < 8 {
self.frame_pool.push(f);
}
}
/// Encode a batch of pictures that do not reference each other (a single I/P picture, or the B-frames of a
/// mini-GOP). Slices of all pictures in the batch are encoded in parallel.
fn encode_batch(&mut self, pics: Vec<(Input, SliceType, bool)>, out: &mut Vec<Packet>) {
struct Job {
input: Input,
st: SliceType,
idr: bool,
qpf: f64,
qp: u8,
cplx: f64,
aq: Vec<i8>,
hdr: SliceHeader,
nal_ref_idc: u8,
l0: Option<usize>,
l1: Option<usize>,
poc: i32,
rec: Option<Frame>,
mbs: Vec<MbInfo>,
results: Vec<SliceOut>,
}
let mut jobs: Vec<Job> = Vec::with_capacity(pics.len());
for (input, mut st, idr) in pics {
if idr {
let old: Vec<RefPic> = self.dpb.drain(..).collect();
for r in old {
self.recycle_ref(r);
}
self.next_frame_num = 0;
}
let poc = 2 * (input.display as i64 - self.last_idr as i64) as i32;
let by_id = self.dpb.iter().enumerate().max_by_key(|(_, r)| r.id).map(|(i, _)| i);
let (mut l0, mut l1) = match st {
SliceType::I => (None, None),
SliceType::P => (by_id, None),
SliceType::B => (
self.dpb.iter().enumerate().filter(|(_, r)| r.poc < poc).max_by_key(|(_, r)| r.poc).map(|(i, _)| i),
self.dpb.iter().enumerate().filter(|(_, r)| r.poc > poc).min_by_key(|(_, r)| r.poc).map(|(i, _)| i),
),
};
if st == SliceType::B && (l0.is_none() || l1.is_none()) {
st = SliceType::P;
l0 = by_id;
l1 = None;
}
if st == SliceType::P && l0.is_none() {
st = SliceType::I;
}
let is_ref = st != SliceType::B;
let cplx = match st {
SliceType::I => input.cplx.intra,
_ => input.cplx.inter,
} as f64;
let qpf = self.rc.frame_qp(st, cplx, self.coded_count as usize + jobs.len());
let qp = qpf.round().clamp(0.0, 51.0) as u8;
let aq: Vec<i8> = if input.aq.is_empty() {
Vec::new()
} else {
input.aq.iter().map(|&o| ((qpf + o as f64).round().clamp(0.0, 51.0) as i32 - qp as i32) as i8).collect()
};
let nal_ref_idc = if idr {
3
} else if is_ref {
2
} else {
0
};
let hdr = SliceHeader {
first_mb: 0,
slice_type: st,
nal_ref_idc,
idr,
idr_pic_id: self.idr_count % 2,
frame_num: self.next_frame_num % (1 << LOG2_MAX_FRAME_NUM),
log2_max_frame_num: LOG2_MAX_FRAME_NUM,
poc_lsb: (poc as u32) % (1 << LOG2_MAX_POC_LSB),
log2_max_poc_lsb: LOG2_MAX_POC_LSB,
l0_modification: None,
cabac: self.params.cabac,
slice_qp: qp as i32,
pic_init_qp: 26,
disable_deblocking: if self.slices > 1 { 2 } else { 0 },
alpha_offset_div2: 0,
beta_offset_div2: 0,
};
let rec = self.new_frame();
let mbs = vec![MbInfo::default(); self.params.mbw * self.params.mbh];
jobs.push(Job { input, st, idr, qpf, qp, cplx, aq, hdr, nal_ref_idc, l0, l1, poc, rec: Some(rec), mbs, results: Vec::new() });
}
// ---- parallel slice encoding (+ deblocking of each slice band)
let (mbw, mbh) = (self.params.mbw, self.params.mbh);
let n = self.slices;
let bounds: Vec<(usize, usize)> = (0..n).map(|s| (s * mbh / n, (s + 1) * mbh / n)).filter(|(a, b)| b > a).collect();
let capture = self.recon.is_some();
let mut slice_results: Vec<Vec<SliceOut>> = (0..jobs.len()).map(|_| Vec::new()).collect();
let mut bufs: Vec<(Frame, Vec<MbInfo>)> = jobs.iter_mut().map(|j| (j.rec.take().unwrap(), std::mem::take(&mut j.mbs))).collect();
{
let this = &*self;
let fencs: Vec<FrameEnc> = jobs
.iter()
.map(|j| FrameEnc {
p: &this.params,
qt: &this.qt,
cost_tab: &this.cost_tab,
st: j.hdr.slice_type,
src: &j.input.frame,
l0: j.l0.map(|i| &this.dpb[i]),
l1: j.l1.map(|i| &this.dpb[i]),
qp: j.qp,
aq: &j.aq,
hdr: j.hdr.clone(),
nal_ref_idc: j.nal_ref_idc,
nal_type: if j.idr { NAL_IDR } else { NAL_SLICE },
cabac_table: if j.hdr.slice_type == SliceType::I { &this.cabac_i } else { &this.cabac_pb },
lam: &this.lam,
})
.collect();
let dps: Vec<DeblockParams> = jobs
.iter()
.map(|j| DeblockParams {
alpha_offset: 0,
beta_offset: 0,
chroma_qp_offset: this.params.chroma_qp_offset,
ref_ids: [j.l0.map_or(-1, |i| this.dpb[i].id as i32), j.l1.map_or(-1, |i| this.dpb[i].id as i32)],
})
.collect();
let deblock: Vec<bool> = jobs.iter().map(|j| j.hdr.slice_type != SliceType::B || capture).collect();
type Task<'t> = (usize, usize, usize, usize, &'t mut [u8], &'t mut [u8], &'t mut [u8], &'t mut [MbInfo]);
let mut tasks: Vec<Task> = Vec::new();
for (pi, (rec, mbs)) in bufs.iter_mut().enumerate() {
let ys = rec.y.stride;
let cs = rec.u.stride;
let mut yrest = &mut rec.y.data[LUMA_PAD * ys..];
let mut urest = &mut rec.u.data[CHROMA_PAD * cs..];
let mut vrest = &mut rec.v.data[CHROMA_PAD * cs..];
let mut mrest = &mut mbs[..];
for (k, &(r0, r1)) in bounds.iter().enumerate() {
let rows = r1 - r0;
let (a, b) = std::mem::take(&mut yrest).split_at_mut(rows * 16 * ys);
yrest = b;
let (ua, ub) = std::mem::take(&mut urest).split_at_mut(rows * 8 * cs);
urest = ub;
let (va, vb) = std::mem::take(&mut vrest).split_at_mut(rows * 8 * cs);
vrest = vb;
let (ma, mb) = std::mem::take(&mut mrest).split_at_mut(rows * mbw);
mrest = mb;
tasks.push((pi, k, r0, r1, a, ua, va, ma));
}
}
let ys = jobs_stride_y(mbw);
let cs = jobs_stride_c(mbw);
let run = |(pi, k, r0, r1, y, u, v, m): Task| -> (usize, usize, SliceOut) {
let out = SliceEnc::new(&fencs[pi], Band { y: &mut *y, u: &mut *u, v: &mut *v, ys, cs }, &mut *m, r0, r1, k as u16).encode();
if deblock[pi] {
let mut by = BandPlane { data: y, stride: ys, pad: LUMA_PAD };
let mut bu = BandPlane { data: u, stride: cs, pad: CHROMA_PAD };
let mut bv = BandPlane { data: v, stride: cs, pad: CHROMA_PAD };
deblock_band(&mut by, &mut bu, &mut bv, m, mbw, r1 - r0, &dps[pi]);
}
(pi, k, out)
};
#[cfg(feature = "threads")]
let results: Vec<(usize, usize, SliceOut)> = {
use rayon::prelude::*;
if tasks.len() > 1 { this.pool.install(|| tasks.into_par_iter().map(run).collect()) } else { tasks.into_iter().map(run).collect() }
};
#[cfg(not(feature = "threads"))]
let results: Vec<(usize, usize, SliceOut)> = tasks.into_iter().map(run).collect();
let mut results = results;
results.sort_by_key(|r| (r.0, r.1));
for (pi, _, r) in results {
slice_results[pi].push(r);
}
}
for ((j, (rec, mbs)), res) in jobs.iter_mut().zip(bufs).zip(slice_results) {
j.rec = Some(rec);
j.mbs = mbs;
j.results = res;
}
// ---- finalise each picture in coding order
for job in jobs {
let Job { input, st, idr, qpf, qp, cplx, hdr: _, nal_ref_idc: _, l0: _, l1: _, poc, rec, mbs, results, aq: _ } = job;
let mut rec = rec.unwrap();
let is_ref = st != SliceType::B;
let annexb = self.cfg.format == PacketFormat::AnnexB;
let mut nals: Vec<Vec<u8>> = Vec::new();
if annexb && self.cfg.aud {
let t = match st {
SliceType::I => 0,
SliceType::P => 1,
SliceType::B => 2,
};
nals.push(nal::nal(0, NAL_AUD, &nal::aud_rbsp(t)));
}
if idr && annexb {
nals.push(self.sps_nal.clone());
nals.push(self.pps_nal.clone());
}
if !self.sei_sent {
self.sei_sent = true;
let text = format!("filmcraft-h264enc {} - clean-room H.264 encoder", env!("CARGO_PKG_VERSION"));
nals.push(nal::nal(0, NAL_SEI, &nal::sei_user_data_rbsp(&text)));
}
for r in results {
nals.push(r.nal);
}
let mut data = Vec::with_capacity(nals.iter().map(|n| n.len() + 4).sum());
for n in &nals {
if annexb {
data.extend_from_slice(&[0, 0, 0, 1]);
} else {
data.extend_from_slice(&(n.len() as u32).to_be_bytes());
}
data.extend_from_slice(n);
}
let total_bits = data.len() as u64 * 8;
self.rc.update(st, qpf, cplx, total_bits);
if self.record_pass1 {
self.rc.record_pass1(st, qpf, cplx, total_bits);
}
if let Some(rv) = self.recon.as_mut() {
let (w, h) = (self.cfg.width as usize, self.cfg.height as usize);
let mut y = Vec::with_capacity(w * h);
let mut u = Vec::with_capacity(w * h / 4);
let mut v = Vec::with_capacity(w * h / 4);
rec.y.copy_out(w, h, &mut y);
rec.u.copy_out(w / 2, h / 2, &mut u);
rec.v.copy_out(w / 2, h / 2, &mut v);
rv.push(ReconFrame { pts: input.pts, y, u, v });
}
if is_ref {
let hp = self.hpel_pool.pop();
let hpel = self.pool.install(|| {
rec.extend_edges();
build_hpel(&rec.y, hp)
});
self.ref_uid += 1;
let rp = RefPic { frame: rec, hpel, mbs, poc, id: self.ref_uid };
if self.dpb.len() >= self.max_refs {
// sliding window: drop the oldest
let oldest = self.dpb.iter().enumerate().min_by_key(|(_, r)| r.id).map(|(i, _)| i).unwrap();
let r = self.dpb.remove(oldest);
self.recycle_ref(r);
}
self.dpb.push(rp);
self.next_frame_num = (self.next_frame_num + 1) % (1 << LOG2_MAX_FRAME_NUM);
} else {
self.recycle(rec);
}
if idr {
self.idr_count += 1;
}
// decoding timestamp base: the k-th smallest presentation time (offset applied in `release`)
let pos = self.display_pts.iter().enumerate().min_by_key(|(_, p)| **p).map(|(i, _)| i).unwrap();
let dts_base = self.display_pts.remove(pos).unwrap();
self.coded_count += 1;
let frame_type = if idr {
FrameType::Idr
} else {
match st {
SliceType::I => FrameType::I,
SliceType::P => FrameType::P,
SliceType::B => FrameType::B,
}
};
self.recycle(input.frame);
out.push(Packet { data, pts: input.pts, dts: dts_base, keyframe: idr, frame_type, qp });
}
}
fn recycle_ref(&mut self, r: RefPic) {
self.recycle(r.frame);
if self.hpel_pool.len() < 4 {
self.hpel_pool.push(r.hpel);
}
}
}
fn jobs_stride_y(mbw: usize) -> usize {
mbw * 16 + 2 * LUMA_PAD
}
fn jobs_stride_c(mbw: usize) -> usize {
mbw * 8 + 2 * CHROMA_PAD
}

156
crates/h264enc/src/lookahead.rs Normal file
View File

@ -0,0 +1,156 @@
//! Cheap half-resolution analysis used for rate control, scene-cut detection and adaptive quantisation.
use crate::dsp::{block_var, sad, satd};
use crate::picture::Frame;
/// Half-resolution luma (2x2 box filtered) of the macroblock-aligned picture.
pub struct LowRes {
pub w: usize,
pub h: usize,
pub data: Vec<u8>,
}
impl LowRes {
pub fn new(f: &Frame) -> Self {
let (w, h) = (f.y.w / 2, f.y.h / 2);
let mut data = vec![0u8; w * h];
let row = |y: usize, out: &mut [u8]| {
let r0 = f.y.idx(0, (2 * y) as isize);
let r1 = f.y.idx(0, (2 * y + 1) as isize);
for x in 0..w {
let s = f.y.data[r0 + 2 * x] as u32 + f.y.data[r0 + 2 * x + 1] as u32 + f.y.data[r1 + 2 * x] as u32 + f.y.data[r1 + 2 * x + 1] as u32;
out[x] = ((s + 2) >> 2) as u8;
}
};
#[cfg(feature = "threads")]
{
use rayon::prelude::*;
data.par_chunks_mut(w).enumerate().for_each(|(y, r)| row(y, r));
}
#[cfg(not(feature = "threads"))]
for (y, r) in data.chunks_mut(w).enumerate() {
row(y, r);
}
LowRes { w, h, data }
}
}
/// Per-frame complexity estimates (sum over 8x8 low-res blocks, i.e. per macroblock).
#[derive(Clone, Copy, Debug, Default)]
pub struct Complexity {
pub intra: u64,
/// Inter cost against the previous frame (min with intra per block); equals `intra` when there is no previous frame.
pub inter: u64,
}
fn block_rows<F: Fn(usize) -> (u64, u64) + Sync + Send>(rows: usize, f: F) -> (u64, u64) {
#[cfg(feature = "threads")]
{
use rayon::prelude::*;
(0..rows).into_par_iter().map(f).reduce(|| (0, 0), |a, b| (a.0 + b.0, a.1 + b.1))
}
#[cfg(not(feature = "threads"))]
{
(0..rows).map(f).fold((0, 0), |a, b| (a.0 + b.0, a.1 + b.1))
}
}
pub fn complexity(cur: &LowRes, prev: Option<&LowRes>) -> Complexity {
let bw = cur.w / 8;
let bh = cur.h / 8;
let (intra, inter) = block_rows(bh, |by| {
let mut ci = 0u64;
let mut cp = 0u64;
let mut blk = [0u8; 64];
for bx in 0..bw {
let (x0, y0) = (bx * 8, by * 8);
let base = y0 * cur.w + x0;
// intra: best of DC / H / V predictions from source neighbours
let mut best = u32::MAX;
for mode in 0..3 {
if (mode == 1 && x0 == 0) || (mode == 2 && y0 == 0) {
continue;
}
for y in 0..8 {
for x in 0..8 {
blk[y * 8 + x] = match mode {
0 => {
if x0 > 0 && y0 > 0 {
((cur.data[(y0 - 1) * cur.w + x0 + x] as u32 + cur.data[(y0 + y) * cur.w + x0 - 1] as u32) / 2) as u8
} else {
128
}
}
1 => cur.data[(y0 + y) * cur.w + x0 - 1],
_ => cur.data[(y0 - 1) * cur.w + x0 + x],
};
}
}
best = best.min(satd(&cur.data[base..], cur.w, &blk, 8, 8, 8));
}
let intra_c = best as u64 + 16;
ci += intra_c;
if let Some(p) = prev {
// small diamond search on SAD, then SATD of the best
let mut bmv = (0i32, 0i32);
let cost_at = |mx: i32, my: i32| -> u32 {
let x = x0 as i32 + mx;
let y = y0 as i32 + my;
if x < 0 || y < 0 || x + 8 > p.w as i32 || y + 8 > p.h as i32 {
return u32::MAX;
}
sad(&cur.data[base..], cur.w, &p.data[y as usize * p.w + x as usize..], p.w, 8, 8)
};
let mut bc = cost_at(0, 0);
for _ in 0..8 {
let c = bmv;
for (dx, dy) in [(-1, 0), (1, 0), (0, -1), (0, 1), (-2, 0), (2, 0), (0, -2), (0, 2)] {
let k = cost_at(c.0 + dx, c.1 + dy);
if k < bc {
bc = k;
bmv = (c.0 + dx, c.1 + dy);
}
}
if bmv == c {
break;
}
}
let x = (x0 as i32 + bmv.0) as usize;
let y = (y0 as i32 + bmv.1) as usize;
let inter_c = satd(&cur.data[base..], cur.w, &p.data[y * p.w + x..], p.w, 8, 8) as u64 + 8;
cp += inter_c.min(intra_c);
} else {
cp += intra_c;
}
}
(ci, cp)
});
Complexity { intra, inter }
}
/// Variance-based adaptive quantisation offsets per macroblock (QP units, mean ~0).
pub fn aq_offsets(f: &Frame, mbw: usize, mbh: usize, strength: f32) -> Vec<f32> {
let mut e = vec![0f32; mbw * mbh];
let row = |my: usize, e: &mut [f32]| {
for mx in 0..mbw {
let i = f.y.idx((mx * 16) as isize, (my * 16) as isize);
let vy = block_var(&f.y.data[i..], f.y.stride, 16, 16) as f32 / 256.0;
let iu = f.u.idx((mx * 8) as isize, (my * 8) as isize);
let vu = block_var(&f.u.data[iu..], f.u.stride, 8, 8) as f32 / 64.0;
let iv = f.v.idx((mx * 8) as isize, (my * 8) as isize);
let vv = block_var(&f.v.data[iv..], f.v.stride, 8, 8) as f32 / 64.0;
e[mx] = (vy + 0.5 * (vu + vv) + 1.0).log2();
}
};
#[cfg(feature = "threads")]
{
use rayon::prelude::*;
e.par_chunks_mut(mbw).enumerate().for_each(|(my, r)| row(my, r));
}
#[cfg(not(feature = "threads"))]
for (my, r) in e.chunks_mut(mbw).enumerate() {
row(my, r);
}
let mean = e.iter().sum::<f32>() / e.len() as f32;
e.iter().map(|&v| (strength * (v - mean)).clamp(-8.0, 8.0)).collect()
}

126
crates/h264enc/src/mbinfo.rs Normal file
View File

@ -0,0 +1,126 @@
//! Per-macroblock state kept for neighbour context derivation, deblocking and co-located (direct) prediction.
#[derive(Clone, Copy, PartialEq, Eq, Debug, Default, Hash)]
pub struct Mv {
pub x: i16,
pub y: i16,
}
impl Mv {
pub const ZERO: Mv = Mv { x: 0, y: 0 };
#[inline]
pub const fn new(x: i32, y: i32) -> Mv {
Mv { x: x as i16, y: y as i16 }
}
}
#[derive(Clone, Copy, PartialEq, Eq, Debug, Default)]
pub enum MbKind {
#[default]
None,
I4x4,
I8x8,
I16x16,
PSkip,
PInter,
BSkip,
BDirect,
BInter,
}
impl MbKind {
#[inline]
pub fn is_intra(self) -> bool {
matches!(self, MbKind::I4x4 | MbKind::I8x8 | MbKind::I16x16)
}
#[inline]
pub fn is_skip(self) -> bool {
matches!(self, MbKind::PSkip | MbKind::BSkip)
}
#[inline]
pub fn is_inxn(self) -> bool {
matches!(self, MbKind::I4x4 | MbKind::I8x8)
}
}
/// Partition shapes for inter macroblocks.
#[derive(Clone, Copy, PartialEq, Eq, Debug, Default)]
pub enum Part {
#[default]
P16x16,
P16x8,
P8x16,
P8x8,
}
impl Part {
pub fn count(self) -> usize {
match self {
Part::P16x16 => 1,
Part::P16x8 | Part::P8x16 => 2,
Part::P8x8 => 4,
}
}
/// (x, y, w, h) in pixels of partition `i`.
pub fn rect(self, i: usize) -> (usize, usize, usize, usize) {
match self {
Part::P16x16 => (0, 0, 16, 16),
Part::P16x8 => (0, 8 * i, 16, 8),
Part::P8x16 => (8 * i, 0, 8, 16),
Part::P8x8 => (8 * (i & 1), 8 * (i >> 1), 8, 8),
}
}
}
#[derive(Clone, Copy, Debug)]
pub struct MbInfo {
pub kind: MbKind,
/// Slice number + 1 (0 = not yet coded).
pub slice: u16,
pub qp: u8,
/// Luma cbp in bits 0..3, chroma cbp in bits 4..5.
pub cbp: u8,
pub t8x8: bool,
pub chroma_mode: u8,
/// Non-zero coefficient counts per luma 4x4 (raster y*4+x).
pub nnz: [u8; 16],
/// Chroma AC non-zero counts: Cb raster 2x2 then Cr raster 2x2.
pub nnz_c: [u8; 8],
/// bit0 luma DC (Intra16x16), bit1 Cb DC, bit2 Cr DC.
pub dc_cbf: u8,
/// Intra 4x4/8x8 prediction modes per 4x4 (raster), -1 when the MB is not I_NxN.
pub imodes: [i8; 16],
/// Reference index per 8x8 (raster) per list; -1 = list unused.
pub ref_idx: [[i8; 4]; 2],
pub mv: [[Mv; 16]; 2],
/// |mvd| components per 4x4 per list, clamped to 127.
pub mvd: [[[u8; 2]; 16]; 2],
}
impl Default for MbInfo {
fn default() -> Self {
MbInfo {
kind: MbKind::None,
slice: 0,
qp: 0,
cbp: 0,
t8x8: false,
chroma_mode: 0,
nnz: [0; 16],
nnz_c: [0; 8],
dc_cbf: 0,
imodes: [-1; 16],
ref_idx: [[-1; 4]; 2],
mv: [[Mv::ZERO; 16]; 2],
mvd: [[[0; 2]; 16]; 2],
}
}
}
impl MbInfo {
#[inline]
pub fn ref_at(&self, list: usize, blk4_raster: usize) -> i8 {
let (x, y) = (blk4_raster & 3, blk4_raster >> 2);
self.ref_idx[list][(y >> 1) * 2 + (x >> 1)]
}
}

311
crates/h264enc/src/nal.rs Normal file
View File

@ -0,0 +1,311 @@
//! NAL unit and parameter-set/slice-header generation (§7.3), SEI and the AVCDecoderConfigurationRecord.
use filmcraft_bitstream::{BitWriter, escape_rbsp};
pub const NAL_SLICE: u8 = 1;
pub const NAL_IDR: u8 = 5;
pub const NAL_SEI: u8 = 6;
pub const NAL_SPS: u8 = 7;
pub const NAL_PPS: u8 = 8;
pub const NAL_AUD: u8 = 9;
/// Build a NAL unit (header + escaped RBSP), without start code.
pub fn nal(nal_ref_idc: u8, nal_type: u8, rbsp: &[u8]) -> Vec<u8> {
let mut out = Vec::with_capacity(rbsp.len() + rbsp.len() / 64 + 2);
out.push((nal_ref_idc << 5) | nal_type);
out.extend_from_slice(&escape_rbsp(rbsp));
out
}
#[derive(Clone, Debug)]
pub struct Vui {
pub sar: (u16, u16),
pub video_full_range: bool,
pub colour_primaries: u8,
pub transfer: u8,
pub matrix: u8,
pub num_units_in_tick: u32,
pub time_scale: u32,
pub max_num_reorder_frames: u32,
pub max_dec_frame_buffering: u32,
}
#[derive(Clone, Debug)]
pub struct Sps {
pub profile_idc: u8,
pub constraint_flags: u8,
pub level_idc: u8,
pub width_mbs: u32,
pub height_mbs: u32,
pub crop_right: u32,
pub crop_bottom: u32,
pub log2_max_frame_num: u32,
pub log2_max_poc_lsb: u32,
pub max_num_ref_frames: u32,
pub vui: Vui,
}
impl Sps {
pub fn rbsp(&self) -> Vec<u8> {
let mut w = BitWriter::new();
w.write_bits(self.profile_idc as u32, 8);
w.write_bits(self.constraint_flags as u32, 8);
w.write_bits(self.level_idc as u32, 8);
w.write_ue(0); // seq_parameter_set_id
if matches!(self.profile_idc, 100 | 110 | 122 | 244 | 44 | 83 | 86 | 118 | 128 | 138 | 139 | 134 | 135) {
w.write_ue(1); // chroma_format_idc 4:2:0
w.write_ue(0); // bit_depth_luma_minus8
w.write_ue(0); // bit_depth_chroma_minus8
w.write_bit(false); // qpprime_y_zero_transform_bypass_flag
w.write_bit(false); // seq_scaling_matrix_present_flag
}
w.write_ue(self.log2_max_frame_num - 4);
w.write_ue(0); // pic_order_cnt_type
w.write_ue(self.log2_max_poc_lsb - 4);
w.write_ue(self.max_num_ref_frames);
w.write_bit(false); // gaps_in_frame_num_value_allowed_flag
w.write_ue(self.width_mbs - 1);
w.write_ue(self.height_mbs - 1);
w.write_bit(true); // frame_mbs_only_flag
w.write_bit(true); // direct_8x8_inference_flag
let crop = self.crop_right > 0 || self.crop_bottom > 0;
w.write_bit(crop);
if crop {
w.write_ue(0);
w.write_ue(self.crop_right / 2);
w.write_ue(0);
w.write_ue(self.crop_bottom / 2);
}
w.write_bit(true); // vui_parameters_present_flag
let v = &self.vui;
let has_sar = v.sar != (0, 0);
w.write_bit(has_sar);
if has_sar {
if v.sar == (1, 1) {
w.write_bits(1, 8);
} else {
w.write_bits(255, 8); // Extended_SAR
w.write_bits(v.sar.0 as u32, 16);
w.write_bits(v.sar.1 as u32, 16);
}
}
w.write_bit(false); // overscan_info_present_flag
w.write_bit(true); // video_signal_type_present_flag
w.write_bits(5, 3); // video_format: unspecified
w.write_bit(v.video_full_range);
w.write_bit(true); // colour_description_present_flag
w.write_bits(v.colour_primaries as u32, 8);
w.write_bits(v.transfer as u32, 8);
w.write_bits(v.matrix as u32, 8);
w.write_bit(false); // chroma_loc_info_present_flag
w.write_bit(true); // timing_info_present_flag
w.write_bits(v.num_units_in_tick, 32);
w.write_bits(v.time_scale, 32);
w.write_bit(true); // fixed_frame_rate_flag
w.write_bit(false); // nal_hrd_parameters_present_flag
w.write_bit(false); // vcl_hrd_parameters_present_flag
w.write_bit(false); // pic_struct_present_flag
w.write_bit(true); // bitstream_restriction_flag
w.write_bit(true); // motion_vectors_over_pic_boundaries_flag
w.write_ue(0); // max_bytes_per_pic_denom
w.write_ue(0); // max_bits_per_mb_denom
w.write_ue(11); // log2_max_mv_length_horizontal (±512 samples)
w.write_ue(11); // log2_max_mv_length_vertical
w.write_ue(v.max_num_reorder_frames);
w.write_ue(v.max_dec_frame_buffering);
w.rbsp_trailing();
w.finish()
}
}
#[derive(Clone, Debug)]
pub struct Pps {
pub cabac: bool,
pub pic_init_qp: i32,
pub chroma_qp_offset: i32,
pub transform_8x8: bool,
pub high: bool,
}
impl Pps {
pub fn rbsp(&self) -> Vec<u8> {
let mut w = BitWriter::new();
w.write_ue(0); // pic_parameter_set_id
w.write_ue(0); // seq_parameter_set_id
w.write_bit(self.cabac);
w.write_bit(false); // bottom_field_pic_order_in_frame_present_flag
w.write_ue(0); // num_slice_groups_minus1
w.write_ue(0); // num_ref_idx_l0_default_active_minus1
w.write_ue(0); // num_ref_idx_l1_default_active_minus1
w.write_bit(false); // weighted_pred_flag
w.write_bits(0, 2); // weighted_bipred_idc
w.write_se(self.pic_init_qp - 26);
w.write_se(0); // pic_init_qs_minus26
w.write_se(self.chroma_qp_offset);
w.write_bit(true); // deblocking_filter_control_present_flag
w.write_bit(false); // constrained_intra_pred_flag
w.write_bit(false); // redundant_pic_cnt_present_flag
if self.high {
w.write_bit(self.transform_8x8);
w.write_bit(false); // pic_scaling_matrix_present_flag
w.write_se(self.chroma_qp_offset); // second_chroma_qp_index_offset
}
w.rbsp_trailing();
w.finish()
}
}
/// Access unit delimiter RBSP. `primary_pic_type`: 0 = I, 1 = I/P, 2 = I/P/B.
pub fn aud_rbsp(primary_pic_type: u8) -> Vec<u8> {
let mut w = BitWriter::new();
w.write_bits(primary_pic_type as u32, 3);
w.rbsp_trailing();
w.finish()
}
/// user_data_unregistered SEI carrying an encoder identification string.
pub fn sei_user_data_rbsp(text: &str) -> Vec<u8> {
const UUID: [u8; 16] = [0x7a, 0x3c, 0x1f, 0x52, 0xd4, 0x0b, 0x4e, 0x8e, 0x9d, 0x61, 0x2f, 0xc4, 0x93, 0x05, 0xb7, 0x1a];
let mut payload = UUID.to_vec();
payload.extend_from_slice(text.as_bytes());
let mut w = BitWriter::new();
w.write_bits(5, 8); // payloadType user_data_unregistered
let mut size = payload.len();
while size >= 255 {
w.write_bits(255, 8);
size -= 255;
}
w.write_bits(size as u32, 8);
w.write_bytes(&payload);
w.rbsp_trailing();
w.finish()
}
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum SliceType {
P = 0,
B = 1,
I = 2,
}
#[derive(Clone, Debug)]
pub struct SliceHeader {
pub first_mb: u32,
pub slice_type: SliceType,
pub nal_ref_idc: u8,
pub idr: bool,
pub idr_pic_id: u32,
pub frame_num: u32,
pub log2_max_frame_num: u32,
pub poc_lsb: u32,
pub log2_max_poc_lsb: u32,
/// Reorder RefPicList0 so that this picNum difference (abs_diff_pic_num_minus1, subtract) comes first.
pub l0_modification: Option<u32>,
pub cabac: bool,
pub slice_qp: i32,
pub pic_init_qp: i32,
pub disable_deblocking: u32,
pub alpha_offset_div2: i32,
pub beta_offset_div2: i32,
}
impl SliceHeader {
/// Write the slice header into `w`; for CABAC the caller must byte-align with one-bits afterwards.
pub fn write(&self, w: &mut BitWriter) {
w.write_ue(self.first_mb);
w.write_ue(self.slice_type as u32);
w.write_ue(0); // pic_parameter_set_id
w.write_bits(self.frame_num, self.log2_max_frame_num);
if self.idr {
w.write_ue(self.idr_pic_id);
}
w.write_bits(self.poc_lsb, self.log2_max_poc_lsb);
if self.slice_type == SliceType::B {
w.write_bit(true); // direct_spatial_mv_pred_flag
}
if self.slice_type != SliceType::I {
w.write_bit(false); // num_ref_idx_active_override_flag
// ref_pic_list_modification
match self.l0_modification {
Some(d) => {
w.write_bit(true);
w.write_ue(0); // modification_of_pic_nums_idc: subtract
w.write_ue(d);
w.write_ue(3);
}
None => w.write_bit(false),
}
if self.slice_type == SliceType::B {
w.write_bit(false);
}
}
if self.nal_ref_idc != 0 {
if self.idr {
w.write_bit(false); // no_output_of_prior_pics_flag
w.write_bit(false); // long_term_reference_flag
} else {
w.write_bit(false); // adaptive_ref_pic_marking_mode_flag
}
}
if self.cabac && self.slice_type != SliceType::I {
w.write_ue(0); // cabac_init_idc
}
w.write_se(self.slice_qp - self.pic_init_qp);
w.write_ue(self.disable_deblocking);
if self.disable_deblocking != 1 {
w.write_se(self.alpha_offset_div2);
w.write_se(self.beta_offset_div2);
}
}
}
/// AVCDecoderConfigurationRecord (ISO/IEC 14496-15 §5.3.3.1) with 4-byte NAL lengths.
pub fn avcc(sps_nal: &[u8], pps_nal: &[u8], profile_idc: u8, constraint: u8, level: u8) -> Vec<u8> {
let mut v = vec![1, profile_idc, constraint, level, 0xFF, 0xE1];
v.extend_from_slice(&(sps_nal.len() as u16).to_be_bytes());
v.extend_from_slice(sps_nal);
v.push(1);
v.extend_from_slice(&(pps_nal.len() as u16).to_be_bytes());
v.extend_from_slice(pps_nal);
if matches!(profile_idc, 100 | 110 | 122 | 144) {
v.push(0xFC | 1); // chroma_format 4:2:0
v.push(0xF8); // bit_depth_luma_minus8
v.push(0xF8); // bit_depth_chroma_minus8
v.push(0); // numOfSequenceParameterSetExt
}
v
}
/// Pick the lowest level satisfying frame size, macroblock rate, DPB size and (optionally) bitrate.
pub fn pick_level(width_mbs: u32, height_mbs: u32, fps: f64, dpb_frames: u32, kbps: Option<u32>, high: bool) -> u8 {
// (level_idc, MaxMBPS, MaxFS, MaxDpbMbs, MaxBR kbit/s)
const LEVELS: [(u8, u64, u64, u64, u64); 17] = [
(10, 1485, 99, 396, 64),
(11, 3000, 396, 900, 192),
(12, 6000, 396, 2376, 384),
(13, 11880, 396, 2376, 768),
(20, 11880, 396, 2376, 2000),
(21, 19800, 792, 4752, 4000),
(22, 20250, 1620, 8100, 4000),
(30, 40500, 1620, 8100, 10000),
(31, 108000, 3600, 18000, 14000),
(32, 216000, 5120, 20480, 20000),
(40, 245760, 8192, 32768, 20000),
(41, 245760, 8192, 32768, 50000),
(42, 522240, 8704, 34816, 50000),
(50, 589824, 22080, 110400, 135000),
(51, 983040, 36864, 184320, 240000),
(52, 2073600, 36864, 184320, 240000),
(62, 16711680, 139264, 696320, 800000),
];
let fs = (width_mbs * height_mbs) as u64;
let mbps = (fs as f64 * fps).ceil() as u64;
for &(l, max_mbps, max_fs, max_dpb, max_br) in &LEVELS {
let br_ok = kbps.is_none_or(|k| (k as u64) * 1000 <= max_br * if high { 1250 } else { 1000 });
let dim_ok = (width_mbs as u64).pow(2) <= 8 * max_fs && (height_mbs as u64).pow(2) <= 8 * max_fs;
if fs <= max_fs && mbps <= max_mbps && fs * dpb_frames as u64 <= max_dpb && br_ok && dim_ok {
return l;
}
}
62
}

259
crates/h264enc/src/picture.rs Normal file
View File

@ -0,0 +1,259 @@
//! Padded picture planes, reference pictures and motion-compensated prediction (§8.4.2.2).
use crate::mbinfo::{MbInfo, Mv};
pub const LUMA_PAD: usize = 64;
pub const CHROMA_PAD: usize = 32;
/// A plane with `pad` replicated samples on every side.
#[derive(Clone)]
pub struct Plane {
pub data: Vec<u8>,
pub stride: usize,
pub w: usize,
pub h: usize,
pub pad: usize,
}
impl Plane {
pub fn new(w: usize, h: usize, pad: usize) -> Self {
let stride = w + 2 * pad;
Plane { data: vec![0; stride * (h + 2 * pad)], stride, w, h, pad }
}
#[inline]
pub fn idx(&self, x: isize, y: isize) -> usize {
((y + self.pad as isize) as usize) * self.stride + (x + self.pad as isize) as usize
}
/// Replicate edge samples into the padding.
pub fn extend_edges(&mut self) {
let (w, h, pad, stride) = (self.w, self.h, self.pad, self.stride);
for y in 0..h {
let r = (y + pad) * stride;
let l = self.data[r + pad];
let rr = self.data[r + pad + w - 1];
self.data[r..r + pad].fill(l);
self.data[r + pad + w..r + stride].fill(rr);
}
let first = pad * stride;
let last = (pad + h - 1) * stride;
for y in 0..pad {
self.data.copy_within(first..first + stride, y * stride);
self.data.copy_within(last..last + stride, (pad + h + y) * stride);
}
}
/// Copy the picture area out as a tightly packed plane of `cw x ch` (cropping).
pub fn copy_out(&self, cw: usize, ch: usize, out: &mut Vec<u8>) {
for y in 0..ch {
let o = self.idx(0, y as isize);
out.extend_from_slice(&self.data[o..o + cw]);
}
}
}
/// A 4:2:0 frame with padded planes.
#[derive(Clone)]
pub struct Frame {
pub y: Plane,
pub u: Plane,
pub v: Plane,
}
impl Frame {
pub fn new(w: usize, h: usize) -> Self {
Frame { y: Plane::new(w, h, LUMA_PAD), u: Plane::new(w / 2, h / 2, CHROMA_PAD), v: Plane::new(w / 2, h / 2, CHROMA_PAD) }
}
pub fn extend_edges(&mut self) {
self.y.extend_edges();
self.u.extend_edges();
self.v.extend_edges();
}
}
/// A decoded reference picture with half-sample planes for fast sub-pel access.
pub struct RefPic {
pub frame: Frame,
/// Horizontal half (b), vertical half (h) and centre (j) planes, same geometry as luma.
pub hpel: [Plane; 3],
pub mbs: Vec<MbInfo>,
pub poc: i32,
pub id: u32,
}
#[inline]
fn clip8(v: i32) -> u8 {
v.clamp(0, 255) as u8
}
/// Compute the H, V and centre half-sample planes over the whole padded area (clamped at the buffer edge,
/// which equals the normative picture-edge clamping because the padding is replicated).
pub fn build_hpel(f: &Plane, reuse: Option<[Plane; 3]>) -> [Plane; 3] {
let [mut hp, mut vp, mut cp] = reuse.filter(|r| r[0].data.len() == f.data.len()).unwrap_or_else(|| {
let mk = || Plane { data: vec![0; f.data.len()], stride: f.stride, w: f.w, h: f.h, pad: f.pad };
[mk(), mk(), mk()]
});
let rows = f.data.len() / f.stride;
let band = 16usize;
let work = |b: usize, hrows: &mut [u8], vrows: &mut [u8], crows: &mut [u8]| {
let r0 = b * band;
let r1 = (r0 + band).min(rows);
hpel_rows(f, r0, r1, hrows, vrows, crows);
};
let chunk = band * f.stride;
#[cfg(feature = "threads")]
{
use rayon::prelude::*;
hp.data
.par_chunks_mut(chunk)
.zip(vp.data.par_chunks_mut(chunk))
.zip(cp.data.par_chunks_mut(chunk))
.enumerate()
.for_each(|(b, ((h, v), c))| work(b, h, v, c));
}
#[cfg(not(feature = "threads"))]
{
for (b, ((h, v), c)) in hp.data.chunks_mut(chunk).zip(vp.data.chunks_mut(chunk)).zip(cp.data.chunks_mut(chunk)).enumerate() {
work(b, h, v, c);
}
}
[hp, vp, cp]
}
fn hpel_rows(f: &Plane, r0: usize, r1: usize, hout: &mut [u8], vout: &mut [u8], cout: &mut [u8]) {
let stride = f.stride;
let rows = f.data.len() / stride;
let clampr = |r: isize| r.clamp(0, rows as isize - 1) as usize;
// Horizontal intermediate (unclipped, unshifted) for rows r0-2 .. r1+3.
let nrows = r1 - r0 + 5;
let mut hs = vec![0i16; nrows * stride];
let mut ext = vec![0u8; stride + 5];
for k in 0..nrows {
let r = clampr(r0 as isize - 2 + k as isize);
let row = &f.data[r * stride..r * stride + stride];
ext[0] = row[0];
ext[1] = row[0];
ext[2..2 + stride].copy_from_slice(row);
ext[stride + 2] = row[stride - 1];
ext[stride + 3] = row[stride - 1];
ext[stride + 4] = row[stride - 1];
let o = &mut hs[k * stride..k * stride + stride];
for x in 0..stride {
let e = &ext[x..x + 6];
o[x] = (e[0] as i16 + e[5] as i16) - 5 * (e[1] as i16 + e[4] as i16) + 20 * (e[2] as i16 + e[3] as i16);
}
}
for r in r0..r1 {
let lo = (r - r0) * stride;
let k = r - r0 + 2;
let hrow = &hs[k * stride..k * stride + stride];
for x in 0..stride {
hout[lo + x] = clip8((hrow[x] as i32 + 16) >> 5);
}
// vertical
let rr: [usize; 6] = std::array::from_fn(|i| clampr(r as isize - 2 + i as isize) * stride);
for x in 0..stride {
let s = f.data[rr[0] + x] as i32 + f.data[rr[5] + x] as i32 - 5 * (f.data[rr[1] + x] as i32 + f.data[rr[4] + x] as i32)
+ 20 * (f.data[rr[2] + x] as i32 + f.data[rr[3] + x] as i32);
vout[lo + x] = clip8((s + 16) >> 5);
}
// centre from horizontal intermediates
let hk: [&[i16]; 6] = std::array::from_fn(|i| &hs[(k - 2 + i) * stride..(k - 1 + i) * stride]);
for x in 0..stride {
let s = hk[0][x] as i32 + hk[5][x] as i32 - 5 * (hk[1][x] as i32 + hk[4][x] as i32) + 20 * (hk[2][x] as i32 + hk[3][x] as i32);
cout[lo + x] = clip8((s + 512) >> 10);
}
}
}
/// Plane selection per quarter-sample position: two (plane, dx, dy) sources to average (identical = no average).
/// Plane indices: 0 full, 1 horizontal half, 2 vertical half, 3 centre.
const QPEL_SRC: [[(u8, u8, u8); 2]; 16] = [
// fy = 0
[(0, 0, 0), (0, 0, 0)],
[(0, 0, 0), (1, 0, 0)],
[(1, 0, 0), (1, 0, 0)],
[(1, 0, 0), (0, 1, 0)],
// fy = 1
[(0, 0, 0), (2, 0, 0)],
[(1, 0, 0), (2, 0, 0)],
[(1, 0, 0), (3, 0, 0)],
[(1, 0, 0), (2, 1, 0)],
// fy = 2
[(2, 0, 0), (2, 0, 0)],
[(2, 0, 0), (3, 0, 0)],
[(3, 0, 0), (3, 0, 0)],
[(3, 0, 0), (2, 1, 0)],
// fy = 3
[(2, 0, 0), (0, 0, 1)],
[(2, 0, 0), (1, 0, 1)],
[(3, 0, 0), (1, 0, 1)],
[(2, 1, 0), (1, 0, 1)],
];
impl RefPic {
#[inline]
fn plane(&self, p: u8) -> &Plane {
match p {
0 => &self.frame.y,
n => &self.hpel[n as usize - 1],
}
}
/// Luma prediction of a w x h block at picture position (x, y) displaced by `mv` (quarter samples).
pub fn mc_luma(&self, x: usize, y: usize, mv: Mv, w: usize, h: usize, out: &mut [u8], ostride: usize) {
let xi = x as isize + (mv.x as isize >> 2);
let yi = y as isize + (mv.y as isize >> 2);
let q = ((mv.y & 3) * 4 + (mv.x & 3)) as usize;
let [(pa, dxa, dya), (pb, dxb, dyb)] = QPEL_SRC[q];
let a = self.plane(pa);
let stride = a.stride;
let ia = a.idx(xi + dxa as isize, yi + dya as isize);
if pa == pb && dxa == dxb && dya == dyb {
for r in 0..h {
out[r * ostride..r * ostride + w].copy_from_slice(&a.data[ia + r * stride..ia + r * stride + w]);
}
} else {
let b = self.plane(pb);
let ib = b.idx(xi + dxb as isize, yi + dyb as isize);
for r in 0..h {
let ra = &a.data[ia + r * stride..ia + r * stride + w];
let rb = &b.data[ib + r * stride..ib + r * stride + w];
let o = &mut out[r * ostride..r * ostride + w];
for i in 0..w {
o[i] = ((ra[i] as u16 + rb[i] as u16 + 1) >> 1) as u8;
}
}
}
}
/// Chroma prediction (both planes) of a w x h chroma block at chroma position (x, y); mv in luma quarter units.
pub fn mc_chroma(&self, x: usize, y: usize, mv: Mv, w: usize, h: usize, out_u: &mut [u8], out_v: &mut [u8], ostride: usize) {
let xi = x as isize + (mv.x as isize >> 3);
let yi = y as isize + (mv.y as isize >> 3);
let fx = (mv.x & 7) as u32;
let fy = (mv.y & 7) as u32;
let w00 = (8 - fx) * (8 - fy);
let w01 = fx * (8 - fy);
let w10 = (8 - fx) * fy;
let w11 = fx * fy;
for (p, out) in [(&self.frame.u, out_u), (&self.frame.v, out_v)] {
let s = p.stride;
let i0 = p.idx(xi, yi);
for r in 0..h {
let a = &p.data[i0 + r * s..i0 + r * s + w + 1];
let b = &p.data[i0 + (r + 1) * s..i0 + (r + 1) * s + w + 1];
let o = &mut out[r * ostride..r * ostride + w];
for i in 0..w {
o[i] = ((w00 * a[i] as u32 + w01 * a[i + 1] as u32 + w10 * b[i] as u32 + w11 * b[i + 1] as u32 + 32) >> 6) as u8;
}
}
}
}
}
/// Average two predictions in place (default bi-prediction).
#[inline]
pub fn avg_into(dst: &mut [u8], src: &[u8]) {
for (d, &s) in dst.iter_mut().zip(src) {
*d = ((*d as u16 + s as u16 + 1) >> 1) as u8;
}
}

262
crates/h264enc/src/ratecontrol.rs Normal file
View File

@ -0,0 +1,262 @@
//! Frame-level rate control: constant QP, constant quality (CRF-like), 1-pass ABR with VBV, and 2-pass VBR.
use crate::nal::SliceType;
pub fn qp2qscale(qp: f64) -> f64 {
0.85 * 2f64.powf((qp - 12.0) / 6.0)
}
pub fn qscale2qp(q: f64) -> f64 {
12.0 + 6.0 * (q / 0.85).log2()
}
const QCOMP: f64 = 0.6;
fn type_idx(t: SliceType) -> usize {
match t {
SliceType::I => 0,
SliceType::P => 1,
SliceType::B => 2,
}
}
/// Relative qscale factor per frame type (I frames finer, B frames coarser).
fn type_factor(t: SliceType) -> f64 {
match t {
SliceType::I => 1.0 / 1.4,
SliceType::P => 1.0,
SliceType::B => 1.3,
}
}
/// Statistics of one frame from the first pass.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct FrameStat {
pub slice_type: SliceType,
pub qscale: f64,
pub bits: u64,
pub cplx: f64,
}
/// First-pass statistics, serialisable as text.
#[derive(Clone, Debug, Default, PartialEq)]
pub struct PassStats {
pub frames: Vec<FrameStat>,
pub fps: f64,
}
impl PassStats {
pub fn to_text(&self) -> String {
let mut s = format!("h264enc-pass1 fps={}\n", self.fps);
for f in &self.frames {
let t = match f.slice_type {
SliceType::I => 'I',
SliceType::P => 'P',
SliceType::B => 'B',
};
s.push_str(&format!("{t} {} {} {}\n", f.qscale, f.bits, f.cplx));
}
s
}
pub fn from_text(text: &str) -> Option<PassStats> {
let mut lines = text.lines();
let head = lines.next()?;
let fps = head.strip_prefix("h264enc-pass1 fps=")?.trim().parse().ok()?;
let mut frames = Vec::new();
for l in lines {
let mut it = l.split_whitespace();
let t = match it.next()? {
"I" => SliceType::I,
"P" => SliceType::P,
"B" => SliceType::B,
_ => return None,
};
frames.push(FrameStat { slice_type: t, qscale: it.next()?.parse().ok()?, bits: it.next()?.parse().ok()?, cplx: it.next()?.parse().ok()? });
}
Some(PassStats { frames, fps })
}
}
#[derive(Clone, Debug)]
pub enum RcKind {
Qp(u8),
Crf(f32),
Abr { kbps: u32 },
TwoPass { kbps: u32, plan: Vec<f64>, planned_bits: Vec<f64> },
}
pub struct RateControl {
kind: RcKind,
fps: f64,
/// bits ≈ coef * cplx / qscale per frame type.
coef: [f64; 3],
coef_n: [u32; 3],
total_bits: f64,
frames: u64,
sum_s: f64,
// VBV
vbv_max_bits_per_frame: f64,
vbv_size: f64,
vbv_fill: f64,
pub pass1: Vec<FrameStat>,
model_done: f64,
inflight: std::collections::VecDeque<f64>,
last_q: f64,
prior: Option<(f64, f64)>,
}
impl RateControl {
pub fn new(kind: RcKind, fps: f64, vbv_maxrate_kbps: Option<u32>, vbv_buf_kbits: Option<u32>) -> Self {
let max = vbv_maxrate_kbps.map_or(0.0, |k| k as f64 * 1000.0 / fps);
let size = vbv_buf_kbits.map_or(0.0, |k| k as f64 * 1000.0);
RateControl {
kind,
fps,
coef: [1.6, 1.4, 1.2],
coef_n: [0; 3],
total_bits: 0.0,
frames: 0,
sum_s: 0.0,
vbv_max_bits_per_frame: max,
vbv_size: size,
vbv_fill: size * 0.9,
pass1: Vec::new(),
model_done: 0.0,
inflight: std::collections::VecDeque::new(),
last_q: 0.0,
prior: None,
}
}
/// Plan a second pass from first-pass statistics.
pub fn plan_two_pass(stats: &PassStats, kbps: u32) -> RcKind {
let n = stats.frames.len().max(1);
let target = kbps as f64 * 1000.0 * n as f64 / stats.fps;
let c: Vec<f64> = stats.frames.iter().map(|f| (f.bits as f64 * f.qscale).max(1.0)).collect();
let qs = |rf: f64| -> Vec<f64> {
stats.frames.iter().zip(&c).map(|(f, &ci)| (rf * ci.powf(1.0 - QCOMP) * type_factor(f.slice_type)).clamp(qp2qscale(0.0), qp2qscale(51.0))).collect()
};
let bits_for = |q: &[f64]| -> f64 { stats.frames.iter().zip(q).map(|(f, &qi)| f.bits as f64 * f.qscale / qi).sum() };
// bisection on log(rf)
let (mut lo, mut hi) = (-30.0f64, 30.0f64);
for _ in 0..100 {
let mid = 0.5 * (lo + hi);
if bits_for(&qs(mid.exp())) > target {
lo = mid;
} else {
hi = mid;
}
}
let plan = qs((0.5 * (lo + hi)).exp());
let planned_bits = stats.frames.iter().zip(&plan).map(|(f, &q)| f.bits as f64 * f.qscale / q).collect();
RcKind::TwoPass { kbps, plan, planned_bits }
}
/// Choose the frame QP (fractional) for frame `index` (display order) of `t` with lookahead complexity `cplx`.
pub fn frame_qp(&mut self, t: SliceType, cplx: f64, index_in_coding_order: usize) -> f64 {
let ti = type_idx(t);
let cplx = cplx.max(1.0);
let qp = match &self.kind {
RcKind::Qp(q) => {
let q = *q as f64;
return if t == SliceType::B { (q + 2.0).min(51.0) } else { q };
}
RcKind::Crf(c) => {
let c = *c as f64;
return (c + match t {
SliceType::I => -2.0,
SliceType::P => 0.0,
SliceType::B => 1.5,
})
.clamp(0.0, 51.0);
}
RcKind::Abr { kbps } => {
let bpf = *kbps as f64 * 1000.0 / self.fps;
// rate factor so that past frames at this rf would have hit the target
let s_here = self.coef[ti] * cplx.powf(QCOMP) / type_factor(t);
// Virtual history of one second of P frames (inter cost ~ 40% of the first frame's) so the first
// I frame gets a realistic share of the budget instead of a single frame's worth.
if self.prior.is_none() {
let n = self.fps.max(1.0);
let c0 = if t == SliceType::I { 0.4 * cplx } else { cplx };
self.prior = Some((n, n * self.coef[1] * c0.powf(QCOMP)));
}
let (pn0, ps0) = self.prior.unwrap();
let fade = (1.0 - self.frames as f64 / pn0).max(0.0);
let (pn, ps) = (pn0 * fade, ps0 * fade);
let s = self.sum_s + s_here + ps;
let wanted = bpf * (self.frames as f64 + 1.0 + pn);
let mut rf = s / wanted;
let abr_buffer = bpf * self.fps.max(1.0);
let overflow = (1.0 + (self.total_bits - bpf * self.frames as f64) / abr_buffer).clamp(0.5, 2.0);
rf *= overflow;
let q = rf * cplx.powf(1.0 - QCOMP) * type_factor(t);
qscale2qp(q)
}
RcKind::TwoPass { kbps, plan, planned_bits } => {
// Closed-loop second pass: `m` measures how far real frame sizes deviate from the first-pass model
// (bits ∝ 1/qscale) at the qscales actually used; the remaining plan is rescaled so its predicted
// size fits the remaining budget.
let i = index_in_coding_order.min(plan.len().saturating_sub(1));
let bpf = *kbps as f64 * 1000.0 / self.fps;
let target_total: f64 = planned_bits.iter().sum();
let smooth = 0.5 * bpf * self.fps.max(1.0);
let m = ((self.total_bits + smooth) / (self.model_done + smooth)).clamp(0.3, 3.0);
let inflight: f64 = self.inflight.iter().sum();
let remaining_plan: f64 = planned_bits[i..].iter().sum();
let remaining_budget = (target_total - self.total_bits - m * inflight).max(0.05 * remaining_plan.max(1.0));
let f = (m * remaining_plan / remaining_budget).clamp(0.25, 4.0);
let q = plan[i] * f;
self.inflight.push_back(planned_bits[i] * plan[i] / q);
qscale2qp(q)
}
};
let mut qp = qp.clamp(1.0, 51.0);
// VBV: make sure the predicted frame size fits in the buffer.
if self.vbv_size > 0.0 {
let fill = self.vbv_fill;
for _ in 0..60 {
let pred = self.coef[ti] * cplx / qp2qscale(qp);
if pred <= fill * 0.8 || qp >= 51.0 {
break;
}
qp += 0.5;
}
}
self.last_q = qp;
qp
}
/// Update after coding a frame with actual `bits`.
pub fn update(&mut self, t: SliceType, qp: f64, cplx: f64, bits: u64) {
let ti = type_idx(t);
let cplx = cplx.max(1.0);
let q = qp2qscale(qp);
let c = bits as f64 * q / cplx;
let n = self.coef_n[ti];
let w = if n == 0 { 1.0 } else { 0.35f64.max(1.0 / (n + 1) as f64) };
self.coef[ti] = self.coef[ti] * (1.0 - w) + c * w;
self.coef_n[ti] += 1;
// If a type has never been observed, keep it proportional to observed ones.
for k in 0..3 {
if self.coef_n[k] == 0 {
self.coef[k] = self.coef[ti] * [1.2, 1.0, 0.85][k] / [1.2, 1.0, 0.85][ti];
}
}
self.sum_s += self.coef[ti] * cplx.powf(QCOMP) / type_factor(t);
self.total_bits += bits as f64;
if let Some(p) = self.inflight.pop_front() {
self.model_done += p;
}
self.frames += 1;
if self.vbv_size > 0.0 {
self.vbv_fill = (self.vbv_fill - bits as f64 + self.vbv_max_bits_per_frame).min(self.vbv_size);
if self.vbv_fill < 0.0 {
self.vbv_fill = 0.0;
}
}
}
pub fn record_pass1(&mut self, t: SliceType, qp: f64, cplx: f64, bits: u64) {
self.pass1.push(FrameStat { slice_type: t, qscale: qp2qscale(qp), bits, cplx });
}
}

1805
crates/h264enc/src/slice.rs Normal file
View File

File diff suppressed because it is too large Load Diff

120
crates/h264enc/src/syntax.rs Normal file
View File

@ -0,0 +1,120 @@
//! Macroblock-layer syntax values produced by mode decision and consumed by the entropy coders.
use crate::mbinfo::{MbInfo, MbKind, Mv, Part};
use crate::nal::SliceType;
#[derive(Clone)]
pub struct MbCode {
pub kind: MbKind,
pub part: Part,
/// Prediction direction per partition: 0 = L0, 1 = L1, 2 = Bi.
pub bdir: [u8; 4],
pub i16_mode: u8,
pub chroma_mode: u8,
/// Per 4x4 (or 8x8, first four entries) block in decoding order: -1 = use predicted mode, else rem_intra_pred_mode.
pub ipred: [i8; 16],
pub t8x8: bool,
/// mvd per list per partition.
pub mvd: [[Mv; 4]; 2],
pub cbp: u8,
pub qp_delta: i32,
/// Intra16x16 DC levels in zig-zag order.
pub luma_dc: [i16; 16],
/// Luma 4x4 levels by luma4x4BlkIdx, zig-zag order (index 0 unused for Intra16x16 AC).
pub luma4: [[i16; 16]; 16],
/// Luma 8x8 levels by 8x8 block, zig-zag order.
pub luma8: [[i16; 64]; 4],
/// Chroma DC levels (Cb, Cr) in raster (c0..c3) order.
pub cdc: [[i16; 4]; 2],
/// Chroma AC levels, Cb blocks 0..3 then Cr 0..3, zig-zag order (index 0 unused).
pub cac: [[i16; 16]; 8],
}
impl Default for MbCode {
fn default() -> Self {
MbCode {
kind: MbKind::None,
part: Part::P16x16,
bdir: [0; 4],
i16_mode: 0,
chroma_mode: 0,
ipred: [-1; 16],
t8x8: false,
mvd: [[Mv::ZERO; 4]; 2],
cbp: 0,
qp_delta: 0,
luma_dc: [0; 16],
luma4: [[0; 16]; 16],
luma8: [[0; 64]; 4],
cdc: [[0; 4]; 2],
cac: [[0; 16]; 8],
}
}
}
impl MbCode {
/// B macroblock type number (Table 7-14) for inter B macroblocks.
pub fn b_type_number(&self) -> u32 {
match self.kind {
MbKind::BDirect => 0,
_ => match self.part {
Part::P16x16 => 1 + self.bdir[0] as u32,
Part::P16x8 | Part::P8x16 => {
const PAIRS: [(u8, u8); 9] = [(0, 0), (1, 1), (0, 1), (1, 0), (0, 2), (1, 2), (2, 0), (2, 1), (2, 2)];
let k = PAIRS.iter().position(|&p| p == (self.bdir[0], self.bdir[1])).unwrap() as u32;
4 + 2 * k + (self.part == Part::P8x16) as u32
}
Part::P8x8 => 22,
},
}
}
pub fn uses_list(&self, part: usize, list: usize) -> bool {
match self.kind {
MbKind::PInter => list == 0,
MbKind::BInter => self.bdir[part] == 2 || self.bdir[part] as usize == list,
_ => false,
}
}
}
/// Neighbour context for entropy coding one macroblock.
pub struct NbCtx<'a> {
pub left: Option<&'a MbInfo>,
pub top: Option<&'a MbInfo>,
pub cur: &'a MbInfo,
pub slice_type: SliceType,
pub prev_dqp_nz: bool,
pub t8x8_mode: bool,
}
impl NbCtx<'_> {
/// coded_block_flag-style condition for an unavailable neighbour.
#[inline]
pub fn unavail_cbf(&self) -> u32 {
self.cur.kind.is_intra() as u32
}
/// Luma 4x4 neighbour non-zero count (A = left when `left`, else B = above). Returns None if unavailable.
#[inline]
pub fn luma_nnz(&self, raster: usize, left: bool) -> Option<u8> {
let (x, y) = (raster & 3, raster >> 2);
if left {
if x > 0 { Some(self.cur.nnz[raster - 1]) } else { self.left.map(|m| m.nnz[y * 4 + 3]) }
} else if y > 0 {
Some(self.cur.nnz[raster - 4])
} else {
self.top.map(|m| m.nnz[12 + x])
}
}
/// Chroma AC neighbour count for component `c`, block raster `b` (2x2).
#[inline]
pub fn chroma_nnz(&self, c: usize, b: usize, left: bool) -> Option<u8> {
let (x, y) = (b & 1, b >> 1);
if left {
if x > 0 { Some(self.cur.nnz_c[c * 4 + b - 1]) } else { self.left.map(|m| m.nnz_c[c * 4 + y * 2 + 1]) }
} else if y > 0 {
Some(self.cur.nnz_c[c * 4 + b - 2])
} else {
self.top.map(|m| m.nnz_c[c * 4 + 2 + x])
}
}
}

1095
crates/h264enc/src/tables.rs Normal file
View File

File diff suppressed because it is too large Load Diff

424
crates/h264enc/src/transform.rs Normal file
View File

@ -0,0 +1,424 @@
//! Integer transforms, quantisation and the normative inverse transforms/scaling (§8.5.12).
use crate::tables::{DEQUANT4, DEQUANT8, QUANT4, pos_class4, pos_class8};
/// Forward 4x4 core transform of a residual block (raster order).
#[inline]
pub fn fdct4(d: &[i32; 16]) -> [i32; 16] {
let mut t = [0i32; 16];
for i in 0..4 {
let (a, b, c, e) = (d[i * 4], d[i * 4 + 1], d[i * 4 + 2], d[i * 4 + 3]);
let s03 = a + e;
let d03 = a - e;
let s12 = b + c;
let d12 = b - c;
t[i * 4] = s03 + s12;
t[i * 4 + 1] = 2 * d03 + d12;
t[i * 4 + 2] = s03 - s12;
t[i * 4 + 3] = d03 - 2 * d12;
}
let mut o = [0i32; 16];
for i in 0..4 {
let (a, b, c, e) = (t[i], t[4 + i], t[8 + i], t[12 + i]);
let s03 = a + e;
let d03 = a - e;
let s12 = b + c;
let d12 = b - c;
o[i] = s03 + s12;
o[4 + i] = 2 * d03 + d12;
o[8 + i] = s03 - s12;
o[12 + i] = d03 - 2 * d12;
}
o
}
/// Normative inverse 4x4 transform of scaled coefficients `d` (raster), returns residual r (after the >>6).
#[inline]
pub fn idct4(d: &[i32; 16]) -> [i32; 16] {
let mut f = [0i32; 16];
for i in 0..4 {
let (d0, d1, d2, d3) = (d[i * 4], d[i * 4 + 1], d[i * 4 + 2], d[i * 4 + 3]);
let e = d0 + d2;
let ff = d0 - d2;
let g = (d1 >> 1) - d3;
let h = d1 + (d3 >> 1);
f[i * 4] = e + h;
f[i * 4 + 1] = ff + g;
f[i * 4 + 2] = ff - g;
f[i * 4 + 3] = e - h;
}
let mut r = [0i32; 16];
for j in 0..4 {
let (f0, f1, f2, f3) = (f[j], f[4 + j], f[8 + j], f[12 + j]);
let e = f0 + f2;
let ff = f0 - f2;
let g = (f1 >> 1) - f3;
let h = f1 + (f3 >> 1);
r[j] = (e + h + 32) >> 6;
r[4 + j] = (ff + g + 32) >> 6;
r[8 + j] = (ff - g + 32) >> 6;
r[12 + j] = (e - h + 32) >> 6;
}
r
}
#[inline(always)]
fn fdct8_1d(s: [i32; 8]) -> [i32; 8] {
let a0 = s[0] + s[7];
let a1 = s[1] + s[6];
let a2 = s[2] + s[5];
let a3 = s[3] + s[4];
let b0 = s[0] - s[7];
let b1 = s[1] - s[6];
let b2 = s[2] - s[5];
let b3 = s[3] - s[4];
let e0 = a0 + a3;
let e1 = a1 + a2;
let e2 = a0 - a3;
let e3 = a1 - a2;
[
8 * (e0 + e1),
12 * b0 + 10 * b1 + 6 * b2 + 3 * b3,
8 * e2 + 4 * e3,
10 * b0 - 3 * b1 - 12 * b2 - 6 * b3,
8 * (e0 - e1),
6 * b0 - 12 * b1 + 3 * b2 + 10 * b3,
4 * e2 - 8 * e3,
3 * b0 - 6 * b1 + 10 * b2 - 12 * b3,
]
}
/// Forward 8x8 transform X = A x A^T (unnormalised; rows of A are the H.264 8x8 basis scaled by 8).
pub fn fdct8(d: &[i32; 64]) -> [i32; 64] {
let mut t = [0i32; 64];
for y in 0..8 {
let o = fdct8_1d(d[y * 8..y * 8 + 8].try_into().unwrap());
t[y * 8..y * 8 + 8].copy_from_slice(&o);
}
let mut out = [0i32; 64];
for x in 0..8 {
let col = [t[x], t[8 + x], t[16 + x], t[24 + x], t[32 + x], t[40 + x], t[48 + x], t[56 + x]];
let o = fdct8_1d(col);
for k in 0..8 {
out[k * 8 + x] = o[k];
}
}
out
}
#[cfg(test)]
const A8: [[i32; 8]; 8] = [
[8, 8, 8, 8, 8, 8, 8, 8],
[12, 10, 6, 3, -3, -6, -10, -12],
[8, 4, -4, -8, -8, -4, 4, 8],
[10, -3, -12, -6, 6, 12, 3, -10],
[8, -8, -8, 8, 8, -8, -8, 8],
[6, -12, 3, 10, -10, -3, 12, -6],
[4, -8, 8, -4, -4, 8, -8, 4],
[3, -6, 10, -12, 12, -10, 6, -3],
];
/// Normative inverse 8x8 transform (§8.5.13), returns residual after (+32)>>6.
pub fn idct8(d: &[i32; 64]) -> [i32; 64] {
fn one(s: [i32; 8]) -> [i32; 8] {
let a0 = s[0] + s[4];
let a4 = s[0] - s[4];
let a2 = (s[2] >> 1) - s[6];
let a6 = s[2] + (s[6] >> 1);
let b0 = a0 + a6;
let b2 = a4 + a2;
let b4 = a4 - a2;
let b6 = a0 - a6;
let a1 = -s[3] + s[5] - s[7] - (s[7] >> 1);
let a3 = s[1] + s[7] - s[3] - (s[3] >> 1);
let a5 = -s[1] + s[7] + s[5] + (s[5] >> 1);
let a7 = s[3] + s[5] + s[1] + (s[1] >> 1);
let b1 = a1 + (a7 >> 2);
let b7 = a7 - (a1 >> 2);
let b3 = a3 + (a5 >> 2);
let b5 = (a3 >> 2) - a5;
[b0 + b7, b2 + b5, b4 + b3, b6 + b1, b6 - b1, b4 - b3, b2 - b5, b0 - b7]
}
let mut g = [0i32; 64];
for i in 0..8 {
let row: [i32; 8] = d[i * 8..i * 8 + 8].try_into().unwrap();
let o = one(row);
g[i * 8..i * 8 + 8].copy_from_slice(&o);
}
let mut r = [0i32; 64];
for j in 0..8 {
let col = [g[j], g[8 + j], g[16 + j], g[24 + j], g[32 + j], g[40 + j], g[48 + j], g[56 + j]];
let o = one(col);
for i in 0..8 {
r[i * 8 + j] = (o[i] + 32) >> 6;
}
}
r
}
/// Per-QP quantiser tables.
pub struct QuantTables {
/// mf4[qp%6][pos]
pub mf4: [[i32; 16]; 6],
pub v4: [[i32; 16]; 6],
pub mf8: [[i64; 64]; 6],
pub v8: [[i32; 64]; 6],
}
impl QuantTables {
pub fn new() -> Self {
let mut mf4 = [[0; 16]; 6];
let mut v4 = [[0; 16]; 6];
let mut mf8 = [[0; 64]; 6];
let mut v8 = [[0; 64]; 6];
let norm8 = [512.0, 578.0, 320.0, 578.0, 512.0, 578.0, 320.0, 578.0f64];
for q in 0..6 {
for i in 0..16 {
mf4[q][i] = QUANT4[q][pos_class4(i)];
v4[q][i] = DEQUANT4[q][pos_class4(i)];
}
for i in 0..64 {
let v = DEQUANT8[q][pos_class8(i)];
v8[q][i] = v;
let (x, y) = (i & 7, i >> 3);
mf8[q][i] = ((1u64 << 36) as f64 / (norm8[x] * norm8[y] * v as f64)).round() as i64;
}
}
QuantTables { mf4, v4, mf8, v8 }
}
}
impl Default for QuantTables {
fn default() -> Self {
Self::new()
}
}
/// Deadzone rounding fractions (of 2^qbits), as numerator over 6*... we express as shift-based fractions.
#[derive(Clone, Copy)]
pub struct Deadzone {
/// rounding offset as a fraction of one quantisation step, in 1/256 units.
pub f256: i64,
}
impl Deadzone {
pub const INTRA: Deadzone = Deadzone { f256: 85 }; // ~1/3
pub const INTER: Deadzone = Deadzone { f256: 43 }; // ~1/6
}
/// Quantise a 4x4 block of transform coefficients (raster) in place into levels; returns number of non-zeros.
/// `skip_dc` leaves coefficient 0 untouched (for Intra16x16/chroma AC where DC is coded separately).
#[inline]
pub fn quant4(c: &mut [i32; 16], qt: &QuantTables, qp: u8, dz: Deadzone, skip_dc: bool) -> u32 {
let qbits = 15 + (qp / 6) as u32;
let f = ((1i64 << qbits) * dz.f256) >> 8;
let mf = &qt.mf4[(qp % 6) as usize];
let mut nz = 0;
let start = if skip_dc { 1 } else { 0 };
for i in start..16 {
let v = c[i];
let a = ((v.unsigned_abs() as i64 * mf[i] as i64 + f) >> qbits) as i32;
c[i] = if v < 0 { -a } else { a };
nz += (a != 0) as u32;
}
nz
}
/// Dequantise 4x4 levels (raster) for flat scaling matrices (§8.5.12.1). `skip_dc` leaves c[0] unchanged.
#[inline]
pub fn dequant4(c: &mut [i32; 16], qt: &QuantTables, qp: u8, skip_dc: bool) {
let v = &qt.v4[(qp % 6) as usize];
let sh = (qp / 6) as u32;
let start = if skip_dc { 1 } else { 0 };
for i in start..16 {
// LevelScale = 16*v; qP>=24: (c*LS) << (qP/6-4); else (c*LS + 2^(3-qP/6)) >> (4-qP/6) == c*v << qP/6.
c[i] = (c[i] * v[i]) << sh;
}
}
/// Quantise an 8x8 block (raster), returns non-zero count.
pub fn quant8(c: &mut [i32; 64], qt: &QuantTables, qp: u8, dz: Deadzone) -> u32 {
let qbits = 22 + (qp / 6) as u32;
let f = ((1i64 << qbits) * dz.f256) >> 8;
let mf = &qt.mf8[(qp % 6) as usize];
let mut nz = 0;
for i in 0..64 {
let v = c[i];
let a = ((v.unsigned_abs() as i64 * mf[i] + f) >> qbits) as i32;
c[i] = if v < 0 { -a } else { a };
nz += (a != 0) as u32;
}
nz
}
/// Dequantise 8x8 levels (raster), flat scaling.
pub fn dequant8(c: &mut [i32; 64], qt: &QuantTables, qp: u8) {
let v = &qt.v8[(qp % 6) as usize];
let q6 = (qp / 6) as i32;
for i in 0..64 {
let ls = 16 * v[i];
c[i] = if q6 >= 6 { (c[i] * ls) << (q6 - 6) } else { (c[i] * ls + (1 << (5 - q6))) >> (6 - q6) };
}
}
/// 4x4 Hadamard used for the Intra16x16 luma DC (unnormalised).
#[inline]
pub fn hadamard4(d: &[i32; 16]) -> [i32; 16] {
let mut t = [0i32; 16];
for i in 0..4 {
let (a, b, c, e) = (d[i * 4], d[i * 4 + 1], d[i * 4 + 2], d[i * 4 + 3]);
t[i * 4] = a + b + c + e;
t[i * 4 + 1] = a + b - c - e;
t[i * 4 + 2] = a - b - c + e;
t[i * 4 + 3] = a - b + c - e;
}
let mut o = [0i32; 16];
for i in 0..4 {
let (a, b, c, e) = (t[i], t[4 + i], t[8 + i], t[12 + i]);
o[i] = a + b + c + e;
o[4 + i] = a + b - c - e;
o[8 + i] = a - b - c + e;
o[12 + i] = a - b + c - e;
}
o
}
/// Forward quantisation of the Intra16x16 DC (input: raster 4x4 of block DC coefficients).
pub fn quant_luma_dc(dc: &[i32; 16], qt: &QuantTables, qp: u8, dz: Deadzone) -> ([i32; 16], u32) {
let h = hadamard4(dc);
let qbits = 16 + (qp / 6) as u32;
let f = ((1i64 << qbits) * dz.f256) >> 8;
let mf = qt.mf4[(qp % 6) as usize][0] as i64;
let mut out = [0i32; 16];
let mut nz = 0;
for i in 0..16 {
let v = h[i] >> 1;
let a = ((v.unsigned_abs() as i64 * mf + f) >> qbits) as i32;
out[i] = if v < 0 { -a } else { a };
nz += (a != 0) as u32;
}
(out, nz)
}
/// Normative Intra16x16 DC inverse (§8.5.10): levels (raster 4x4) -> dcY (raster 4x4).
pub fn dequant_luma_dc(c: &[i32; 16], qt: &QuantTables, qp: u8) -> [i32; 16] {
let f = hadamard4(c);
let ls = 16 * qt.v4[(qp % 6) as usize][0];
let q6 = (qp / 6) as i32;
let mut o = [0i32; 16];
for i in 0..16 {
o[i] = if q6 >= 6 { (f[i] * ls) << (q6 - 6) } else { (f[i] * ls + (1 << (5 - q6))) >> (6 - q6) };
}
o
}
/// Forward quantisation of chroma DC (2x2 raster of block DC coefficients).
pub fn quant_chroma_dc(dc: &[i32; 4], qt: &QuantTables, qp: u8, dz: Deadzone) -> ([i32; 4], u32) {
let h = [dc[0] + dc[1] + dc[2] + dc[3], dc[0] - dc[1] + dc[2] - dc[3], dc[0] + dc[1] - dc[2] - dc[3], dc[0] - dc[1] - dc[2] + dc[3]];
let qbits = 16 + (qp / 6) as u32;
let f = ((1i64 << qbits) * dz.f256) >> 8;
let mf = qt.mf4[(qp % 6) as usize][0] as i64;
let mut out = [0i32; 4];
let mut nz = 0;
for i in 0..4 {
let v = h[i];
let a = ((v.unsigned_abs() as i64 * mf + f) >> qbits) as i32;
out[i] = if v < 0 { -a } else { a };
nz += (a != 0) as u32;
}
(out, nz)
}
/// Normative chroma DC inverse (§8.5.11.2) for 4:2:0: levels (raster 2x2) -> dcC.
pub fn dequant_chroma_dc(c: &[i32; 4], qt: &QuantTables, qpc: u8) -> [i32; 4] {
let f = [c[0] + c[1] + c[2] + c[3], c[0] - c[1] + c[2] - c[3], c[0] + c[1] - c[2] - c[3], c[0] - c[1] - c[2] + c[3]];
let ls = 16 * qt.v4[(qpc % 6) as usize][0];
let q6 = (qpc / 6) as u32;
let mut o = [0i32; 4];
for i in 0..4 {
o[i] = ((f[i] * ls) << q6) >> 5;
}
o
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn roundtrip_4x4_low_qp() {
let qt = QuantTables::new();
let mut d = [0i32; 16];
for (i, v) in d.iter_mut().enumerate() {
*v = ((i * 37) % 51) as i32 - 25;
}
let mut c = fdct4(&d);
quant4(&mut c, &qt, 4, Deadzone { f256: 128 }, false);
dequant4(&mut c, &qt, 4, false);
let r = idct4(&c);
for i in 0..16 {
assert!((r[i] - d[i]).abs() <= 1, "{i}: {} vs {}", r[i], d[i]);
}
}
#[test]
fn roundtrip_8x8_low_qp() {
let qt = QuantTables::new();
let mut d = [0i32; 64];
for (i, v) in d.iter_mut().enumerate() {
*v = ((i * 53) % 71) as i32 - 35;
}
let mut c = fdct8(&d);
quant8(&mut c, &qt, 4, Deadzone { f256: 128 });
dequant8(&mut c, &qt, 4);
let r = idct8(&c);
for i in 0..64 {
assert!((r[i] - d[i]).abs() <= 1, "{i}: {} vs {}", r[i], d[i]);
}
}
#[test]
fn fdct8_matches_matrix() {
let mut d = [0i32; 64];
for (i, v) in d.iter_mut().enumerate() {
*v = ((i * 97) % 211) as i32 - 105;
}
let f = fdct8(&d);
for k in 0..8 {
for l in 0..8 {
let mut s = 0;
for y in 0..8 {
for x in 0..8 {
s += A8[k][y] * d[y * 8 + x] * A8[l][x];
}
}
assert_eq!(f[k * 8 + l], s);
}
}
}
#[test]
fn luma_dc_roundtrip() {
let qt = QuantTables::new();
// DC of a flat block of value 10 is 160 per 4x4 block.
let dc = [160i32; 16];
let (lv, _) = quant_luma_dc(&dc, &qt, 10, Deadzone { f256: 128 });
let o = dequant_luma_dc(&lv, &qt, 10);
let mut blk = [0i32; 16];
blk[0] = o[0];
let r = idct4(&blk);
assert!(r.iter().all(|&v| (v - 10).abs() <= 1), "{r:?}");
}
#[test]
fn chroma_dc_roundtrip() {
let qt = QuantTables::new();
let dc = [16 * 7, 16 * 7, 16 * 7, 16 * 7];
let (lv, _) = quant_chroma_dc(&dc, &qt, 10, Deadzone { f256: 128 });
let o = dequant_chroma_dc(&lv, &qt, 10);
let mut blk = [0i32; 16];
blk[0] = o[0];
let r = idct4(&blk);
assert!(r.iter().all(|&v| (v - 7).abs() <= 1), "{r:?}");
}
}

BIN
crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/avcc.h264 Normal file
View File

Binary file not shown.

BIN
crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/baseline_lowqp_slices.h264 Normal file
View File

Binary file not shown.

BIN
crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/baseline_speed.h264 Normal file
View File

Binary file not shown.

BIN
crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/bframes.h264 Normal file
View File

Binary file not shown.

BIN
crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/cbr.h264 Normal file
View File

Binary file not shown.

BIN
crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/high_crf_b2.h264 Normal file
View File

Binary file not shown.

BIN
crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/high_quality_b3_slices.h264 Normal file
View File

Binary file not shown.

BIN
crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/high_speed_odd_size.h264 Normal file
View File

Binary file not shown.

BIN
crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/main_b2.h264 Normal file
View File

Binary file not shown.

BIN
crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/mandelbrot.h264 Normal file
View File

Binary file not shown.

BIN
crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/psnr1080.h264 Normal file
View File

Binary file not shown.

BIN
crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/testsrc2.h264 Normal file
View File

Binary file not shown.

BIN
crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/twopass.h264 Normal file
View File

Binary file not shown.

BIN
crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/vbr1pass.h264 Normal file
View File

Binary file not shown.

154
crates/h264enc/tests/common/mod.rs Normal file
View File

@ -0,0 +1,154 @@
//! Shared helpers for the encoder tests: a synthetic content generator and ffmpeg/ffprobe oracle wrappers.
#![allow(dead_code)]
use std::path::{Path, PathBuf};
use std::process::Command;
pub struct Yuv {
pub w: usize,
pub h: usize,
pub y: Vec<u8>,
pub u: Vec<u8>,
pub v: Vec<u8>,
}
impl Yuv {
pub fn frame(&self) -> filmcraft_h264enc::YuvFrame<'_> {
filmcraft_h264enc::YuvFrame { y: &self.y, u: &self.u, v: &self.v, y_stride: self.w, uv_stride: self.w / 2 }
}
pub fn bytes(&self) -> Vec<u8> {
let mut v = self.y.clone();
v.extend_from_slice(&self.u);
v.extend_from_slice(&self.v);
v
}
}
fn hash(mut x: u32) -> u32 {
x ^= x >> 16;
x = x.wrapping_mul(0x7feb352d);
x ^= x >> 15;
x = x.wrapping_mul(0x846ca68b);
x ^= x >> 16;
x
}
/// Synthetic test picture `t`: four regions with a moving gradient, film-grain-like noise,
/// text-like hard edges and a sub-pixel panning texture.
pub fn synth(w: usize, h: usize, t: usize) -> Yuv {
let mut y = vec![0u8; w * h];
let mut u = vec![0u8; w * h / 4];
let mut v = vec![0u8; w * h / 4];
let tf = t as f32;
for yy in 0..h {
for xx in 0..w {
let region = (xx * 2 / w) + 2 * (yy * 2 / h);
let val: f32 = match region {
0 => {
// moving smooth gradient
let fx = xx as f32 / w as f32;
let fy = yy as f32 / h as f32;
128.0 + 90.0 * ((fx * 6.0 + tf * 0.07).sin() * (fy * 4.0 - tf * 0.05).cos())
}
1 => {
// mild noise over a slow ramp (like film grain)
let n = (hash((xx as u32).wrapping_mul(7919) ^ (yy as u32).wrapping_mul(104729) ^ (t as u32).wrapping_mul(1_000_003)) & 31) as f32;
60.0 + (xx as f32 * 0.1 + yy as f32 * 0.05) % 120.0 + n
}
2 => {
// text-like strokes: glyph grid scrolling slowly upwards
let sy = yy + t;
let gx = xx / 12;
let gy = sy / 16;
let lx = xx % 12;
let ly = sy % 16;
let bits = hash((gx as u32) * 31 + gy as u32 * 977);
let inr = |v: usize, a: usize, b: usize| (a..=b).contains(&v);
let bar = inr(lx, 2, 9);
let stroke = (inr(lx, 2, 3) && bits & 1 != 0)
|| (inr(lx, 8, 9) && bits & 2 != 0)
|| (inr(ly, 3, 4) && bar && bits & 4 != 0)
|| (inr(ly, 8, 9) && bar && bits & 8 != 0)
|| (inr(ly, 12, 13) && bar && bits & 16 != 0);
if ly < 14 && lx < 11 && stroke { 235.0 } else { 20.0 }
}
_ => {
// panning texture with sub-pixel motion
let px = xx as f32 + tf * 1.75;
let py = yy as f32 + tf * 0.5;
let c = ((px / 9.0).sin() + (py / 7.0).cos() + ((px + py) / 23.0).sin()) * 40.0;
let checker = if ((px / 32.0).floor() as i64 + (py / 32.0).floor() as i64) % 2 == 0 { 30.0 } else { -30.0 };
128.0 + c + checker
}
};
y[yy * w + xx] = val.clamp(0.0, 255.0) as u8;
}
}
for yy in 0..h / 2 {
for xx in 0..w / 2 {
let fx = xx as f32 / (w / 2) as f32;
let fy = yy as f32 / (h / 2) as f32;
u[yy * w / 2 + xx] = (128.0 + 60.0 * (fx * 3.0 + tf * 0.03).sin()) as u8;
v[yy * w / 2 + xx] = (128.0 + 50.0 * (fy * 5.0 - tf * 0.04).cos() * fx) as u8;
}
}
Yuv { w, h, y, u, v }
}
pub fn ffmpeg() -> Option<PathBuf> {
for p in ["/opt/homebrew/bin/ffmpeg", "/usr/local/bin/ffmpeg", "/usr/bin/ffmpeg"] {
if Path::new(p).exists() {
return Some(p.into());
}
}
None
}
pub fn ffprobe() -> Option<PathBuf> {
for p in ["/opt/homebrew/bin/ffprobe", "/usr/local/bin/ffprobe", "/usr/bin/ffprobe"] {
if Path::new(p).exists() {
return Some(p.into());
}
}
None
}
pub fn out_dir(name: &str) -> PathBuf {
let base = std::env::var("CARGO_TARGET_DIR").map(PathBuf::from).unwrap_or_else(|_| PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../../target"));
let d = base.join("fixtures").join("h264enc").join(name);
std::fs::create_dir_all(&d).unwrap();
d
}
/// Decode an Annex-B file with ffmpeg; returns (stderr, raw yuv420p bytes).
pub fn ffmpeg_decode(file: &Path) -> (String, Vec<u8>) {
let out = Command::new(ffmpeg().unwrap())
.args(["-v", "error", "-ec", "0", "-i"])
.arg(file)
.args(["-f", "rawvideo", "-pix_fmt", "yuv420p", "-"])
.output()
.expect("run ffmpeg");
(String::from_utf8_lossy(&out.stderr).into_owned(), out.stdout)
}
/// Raw frames produced by ffmpeg from a lavfi source.
pub fn ffmpeg_source(src: &str, w: usize, h: usize, frames: usize) -> Vec<Yuv> {
let out = Command::new(ffmpeg().unwrap())
.args(["-v", "error", "-f", "lavfi", "-i"])
.arg(format!("{src}=size={w}x{h}:rate=30"))
.args(["-frames:v", &frames.to_string(), "-f", "rawvideo", "-pix_fmt", "yuv420p", "-"])
.output()
.expect("run ffmpeg");
assert!(out.status.success(), "{}", String::from_utf8_lossy(&out.stderr));
let fs = w * h * 3 / 2;
out.stdout.chunks_exact(fs).map(|c| Yuv { w, h, y: c[..w * h].to_vec(), u: c[w * h..w * h * 5 / 4].to_vec(), v: c[w * h * 5 / 4..].to_vec() }).collect()
}
pub fn psnr(a: &[u8], b: &[u8]) -> f64 {
assert_eq!(a.len(), b.len());
let se: f64 = a.iter().zip(b).map(|(&x, &y)| (x as f64 - y as f64).powi(2)).sum();
if se == 0.0 {
return 99.0;
}
10.0 * (255.0f64 * 255.0 * a.len() as f64 / se).log10()
}

280
crates/h264enc/tests/ffmpeg_oracle.rs Normal file
View File

@ -0,0 +1,280 @@
//! End-to-end tests using ffmpeg/ffprobe strictly as an external decoder oracle (skipped when absent).
//!
//! Every stream must decode with zero errors, with the exact frame count, and the decoder's output must be
//! bit-identical to the encoder's own reconstruction (catches prediction/transform/deblocking mismatches).
mod common;
use common::{Yuv, ffmpeg, ffmpeg_decode, ffmpeg_source, ffprobe, out_dir, psnr, synth};
use filmcraft_h264enc::*;
use std::process::Command;
fn cfg(w: usize, h: usize, profile: Profile, preset: Preset, bframes: u8, rate: RateControl, slices: usize) -> EncoderConfig {
let mut c = EncoderConfig::new(w as u32, h as u32, 30, 1);
c.profile = profile;
c.preset = preset;
c.bframes = bframes;
c.rate = rate;
c.slices = slices;
c.threads = 4;
c
}
struct Run {
stream: Vec<u8>,
packets: Vec<Packet>,
recon: Vec<ReconFrame>,
}
fn encode(c: EncoderConfig, frames: &[Yuv]) -> Run {
let mut enc = Encoder::new(c).unwrap();
enc.set_recon_capture(true);
let mut packets = Vec::new();
for (i, f) in frames.iter().enumerate() {
packets.extend(enc.encode(&f.frame(), i as i64));
}
packets.extend(enc.flush());
let mut recon = enc.take_recon();
recon.sort_by_key(|r| r.pts);
let stream = packets.iter().flat_map(|p| p.data.iter().copied()).collect();
Run { stream, packets, recon }
}
/// Decode with ffmpeg, check zero errors, exact frame count and bit-exact match with the reconstruction.
/// Returns the decoded frames (display order).
fn verify(name: &str, run: &Run, w: usize, h: usize, n: usize) -> Vec<u8> {
let path = out_dir("oracle").join(format!("{name}.h264"));
std::fs::write(&path, &run.stream).unwrap();
let (err, dec) = ffmpeg_decode(&path);
assert!(err.trim().is_empty(), "{name}: ffmpeg reported errors: {err}");
let fs = w * h * 3 / 2;
assert_eq!(dec.len(), n * fs, "{name}: decoded frame count");
assert_eq!(run.recon.len(), n);
for (i, r) in run.recon.iter().enumerate() {
let d = &dec[i * fs..(i + 1) * fs];
assert!(d[..w * h] == r.y[..], "{name}: frame {i} luma differs from reconstruction");
assert!(d[w * h..w * h * 5 / 4] == r.u[..], "{name}: frame {i} Cb differs");
assert!(d[w * h * 5 / 4..] == r.v[..], "{name}: frame {i} Cr differs");
}
dec
}
fn check_timestamps(run: &Run, n: usize) {
let mut pts: Vec<i64> = run.packets.iter().map(|p| p.pts).collect();
for w in run.packets.windows(2) {
assert!(w[1].dts > w[0].dts, "dts must increase");
}
for p in &run.packets {
assert!(p.dts <= p.pts, "dts {} > pts {}", p.dts, p.pts);
}
pts.sort();
assert_eq!(pts, (0..n as i64).collect::<Vec<_>>());
assert!(run.packets[0].keyframe);
}
#[test]
fn bit_exact_profiles_presets() {
if ffmpeg().is_none() {
eprintln!("ffmpeg not found; skipping");
return;
}
let cases: [(&str, usize, usize, Profile, Preset, u8, RateControl, usize); 6] = [
("baseline_speed", 176, 144, Profile::Baseline, Preset::Speed, 0, RateControl::Qp(26), 1),
("baseline_lowqp_slices", 176, 144, Profile::Baseline, Preset::Balanced, 0, RateControl::Qp(8), 3),
("main_b2", 320, 240, Profile::Main, Preset::Balanced, 2, RateControl::Qp(30), 2),
("high_crf_b2", 320, 240, Profile::High, Preset::Balanced, 2, RateControl::Crf(23.0), 1),
("high_quality_b3_slices", 320, 240, Profile::High, Preset::Quality, 3, RateControl::Qp(22), 3),
("high_speed_odd_size", 330, 250, Profile::High, Preset::Speed, 1, RateControl::Qp(34), 4),
];
for (name, w, h, profile, preset, b, rate, slices) in cases {
let n = 12;
let frames: Vec<Yuv> = (0..n).map(|t| synth(w, h, t)).collect();
let run = encode(cfg(w, h, profile, preset, b, rate, slices), &frames);
let dec = verify(name, &run, w, h, n);
check_timestamps(&run, n);
// decoded display order must follow the source
let fs = w * h * 3 / 2;
for (i, f) in frames.iter().enumerate() {
let p = psnr(&dec[i * fs..i * fs + w * h], &f.y);
assert!(p > 28.0, "{name}: frame {i} psnr {p:.2} (reordering?)");
}
}
}
#[test]
fn bit_exact_lavfi_sources() {
if ffmpeg().is_none() {
return;
}
for src in ["testsrc2", "mandelbrot"] {
let (w, h, n) = (352, 288, 20);
let frames = ffmpeg_source(src, w, h, n);
let run = encode(cfg(w, h, Profile::High, Preset::Balanced, 3, RateControl::Crf(22.0), 2), &frames);
let dec = verify(src, &run, w, h, n);
let fs = w * h * 3 / 2;
let mean: f64 = (0..n).map(|i| psnr(&dec[i * fs..i * fs + w * h], &frames[i].y)).sum::<f64>() / n as f64;
assert!(mean > 38.0, "{src}: mean psnr {mean:.2}");
}
}
#[test]
fn psnr_1080p_crf20() {
if ffmpeg().is_none() {
return;
}
let (w, h, n) = (1920, 1080, 6);
let frames: Vec<Yuv> = (0..n).map(|t| synth(w, h, t)).collect();
let mut c = cfg(w, h, Profile::High, Preset::Speed, 2, RateControl::Crf(20.0), 0);
c.threads = 0;
let run = encode(c, &frames);
let dec = verify("psnr1080", &run, w, h, n);
let fs = w * h * 3 / 2;
let ps: Vec<f64> = (0..n).map(|i| psnr(&dec[i * fs..i * fs + w * h], &frames[i].y)).collect();
let mean = ps.iter().sum::<f64>() / n as f64;
eprintln!("1080p CRF 20 luma PSNR per frame: {ps:.2?} mean {mean:.2}");
assert!(mean >= 40.0, "mean luma PSNR {mean:.2} < 40 dB");
assert!(ps.iter().all(|&p| p >= 38.0), "{ps:?}");
}
fn bitrate(run: &Run, n: usize) -> f64 {
run.stream.len() as f64 * 8.0 / (n as f64 / 30.0) / 1000.0
}
#[test]
fn abr_hits_target() {
if ffmpeg().is_none() {
return;
}
let (w, h, n) = (640, 360, 90);
let frames: Vec<Yuv> = (0..n).map(|t| synth(w, h, t)).collect();
for (name, rate, target) in [("cbr", RateControl::Cbr { kbps: 1500 }, 1500.0), ("vbr1pass", RateControl::Vbr { target_kbps: 1000, max_kbps: 2000 }, 1000.0)]
{
let run = encode(cfg(w, h, Profile::High, Preset::Speed, 2, rate, 0), &frames);
verify(name, &run, w, h, n);
let k = bitrate(&run, n);
eprintln!("{name}: {k:.1} kbps (target {target})");
assert!((k - target).abs() / target <= 0.10, "{name}: {k:.1} kbps vs target {target}");
}
}
#[test]
fn two_pass_vbr_hits_target() {
if ffmpeg().is_none() {
return;
}
let (w, h, n) = (640, 360, 90);
let frames: Vec<Yuv> = (0..n).map(|t| synth(w, h, t)).collect();
let rate = RateControl::Vbr { target_kbps: 800, max_kbps: 2400 };
let mut c1 = cfg(w, h, Profile::High, Preset::Speed, 2, rate, 0);
c1.pass = Pass::First;
let mut enc = Encoder::new(c1.clone()).unwrap();
for (i, f) in frames.iter().enumerate() {
enc.encode(&f.frame(), i as i64);
}
enc.flush();
let stats = enc.pass_stats().unwrap();
assert_eq!(stats.frames.len(), n);
let p1: u64 = stats.frames.iter().map(|f| f.bits).sum();
eprintln!("pass 1: {:.1} kbps; {:?}", p1 as f64 / (n as f64 / 30.0) / 1000.0, &stats.frames[..6]);
let stats = PassStats::from_text(&stats.to_text()).unwrap();
let mut c2 = c1;
c2.pass = Pass::Second(stats);
let run = encode(c2, &frames);
verify("twopass", &run, w, h, n);
let k = bitrate(&run, n);
eprintln!("2-pass: {k:.1} kbps (target 800)");
assert!((k - 800.0).abs() / 800.0 <= 0.10, "2-pass: {k:.1} kbps vs 800");
}
#[test]
fn bframes_order_via_ffprobe() {
if ffmpeg().is_none() || ffprobe().is_none() {
return;
}
let (w, h, n) = (176, 144, 13);
let frames: Vec<Yuv> = (0..n).map(|t| synth(w, h, t)).collect();
let run = encode(cfg(w, h, Profile::High, Preset::Balanced, 3, RateControl::Qp(28), 1), &frames);
verify("bframes", &run, w, h, n);
check_timestamps(&run, n);
let path = out_dir("oracle").join("bframes.h264");
let out = Command::new(ffprobe().unwrap())
.args(["-v", "error", "-select_streams", "v:0", "-show_entries", "frame=pict_type,coded_picture_number", "-of", "csv=p=0"])
.arg(&path)
.output()
.unwrap();
let text = String::from_utf8_lossy(&out.stdout);
let types: Vec<String> = text.lines().map(|l| l.split(',').next().unwrap_or("").trim().to_string()).filter(|s| !s.is_empty()).collect();
assert_eq!(types.len(), n, "{text}");
// Display order as reported by the decoder must match what we encoded: I B B B P B B B P B B B P
let ours: Vec<char> = {
let mut v: Vec<(i64, FrameType)> = run.packets.iter().map(|p| (p.pts, p.frame_type)).collect();
v.sort_by_key(|x| x.0);
v.iter()
.map(|x| match x.1 {
FrameType::Idr | FrameType::I => 'I',
FrameType::P => 'P',
FrameType::B => 'B',
})
.collect()
};
let theirs: Vec<char> = types.iter().map(|t| t.chars().next().unwrap()).collect();
assert_eq!(ours, theirs);
assert!(theirs.contains(&'B'));
}
#[test]
fn length_prefixed_with_avcc() {
if ffmpeg().is_none() {
return;
}
let (w, h, n) = (176, 144, 8);
let frames: Vec<Yuv> = (0..n).map(|t| synth(w, h, t)).collect();
let mut c = cfg(w, h, Profile::High, Preset::Speed, 1, RateControl::Qp(26), 1);
c.format = PacketFormat::LengthPrefixed;
let mut enc = Encoder::new(c).unwrap();
enc.set_recon_capture(true);
let avcc = enc.avcc();
assert_eq!(avcc[0], 1);
assert_eq!(avcc[1], 100);
assert_eq!(avcc[4], 0xFF);
let sps_len = u16::from_be_bytes([avcc[6], avcc[7]]) as usize;
let sps = &avcc[8..8 + sps_len];
let pps_len = u16::from_be_bytes([avcc[9 + sps_len], avcc[10 + sps_len]]) as usize;
let pps = &avcc[11 + sps_len..11 + sps_len + pps_len];
let mut packets = Vec::new();
for (i, f) in frames.iter().enumerate() {
packets.extend(enc.encode(&f.frame(), i as i64));
}
packets.extend(enc.flush());
// convert to Annex-B using the avcC parameter sets
let mut stream = Vec::new();
for nal in [sps, pps] {
stream.extend_from_slice(&[0, 0, 0, 1]);
stream.extend_from_slice(nal);
}
for p in &packets {
let mut i = 0;
while i < p.data.len() {
let len = u32::from_be_bytes(p.data[i..i + 4].try_into().unwrap()) as usize;
stream.extend_from_slice(&[0, 0, 0, 1]);
stream.extend_from_slice(&p.data[i + 4..i + 4 + len]);
i += 4 + len;
}
}
let mut recon = enc.take_recon();
recon.sort_by_key(|r| r.pts);
verify("avcc", &Run { stream, packets, recon }, w, h, n);
}
#[test]
fn rejects_bad_config() {
assert!(Encoder::new(EncoderConfig::new(0, 16, 30, 1)).is_err());
assert!(Encoder::new(EncoderConfig::new(15, 16, 30, 1)).is_err());
assert!(Encoder::new(EncoderConfig::new(16, 16, 0, 1)).is_err());
let mut c = EncoderConfig::new(64, 64, 30, 1);
c.bframes = 4;
assert!(Encoder::new(c).is_err());
let mut enc = Encoder::new(EncoderConfig::new(64, 64, 30, 1)).unwrap();
let y = vec![0u8; 10];
assert!(enc.try_encode(&YuvFrame { y: &y, u: &y, v: &y, y_stride: 64, uv_stride: 32 }, 0).is_err());
}