diff --git a/crates/video/Cargo.toml b/crates/video/Cargo.toml index ea8d4773d..e5bbf1c94 100644 --- a/crates/video/Cargo.toml +++ b/crates/video/Cargo.toml @@ -25,8 +25,12 @@ apv-encode = ["craft-media/bitstream"] vp8 = ["vp8-decode", "vp8-encode"] vp8-decode = [] vp8-encode = [] -# MPEG-2 video (ITU-T H.262) and MPEG-1 video (ISO/IEC 11172-2): decoder. -mpeg2 = [] +# MPEG-2 video (ITU-T H.262) and MPEG-1 video (ISO/IEC 11172-2): decoder (`mpeg2-decode`) and +# MPEG-2 encoder (`mpeg2-encode`). The encoder reuses the decoder's code tables, IDCT and motion +# compensation, so it turns on the decoder too. +mpeg2 = ["mpeg2-decode", "mpeg2-encode"] +mpeg2-decode = [] +mpeg2-encode = ["mpeg2-decode"] # MPEG-4 Part 2 video (ISO/IEC 14496-2, the Xvid/DivX format): Simple profile encoder. mpeg4 = ["mpeg4-encode"] mpeg4-encode = ["craft-media/bitstream"] @@ -132,7 +136,7 @@ required-features = ["mpeg4-encode"] [[test]] name = "mpeg2" -required-features = ["mpeg2"] +required-features = ["mpeg2-decode"] [[test]] name = "prores" diff --git a/crates/video/README.md b/crates/video/README.md index 4aff9b8d0..74b8c6414 100644 --- a/crates/video/README.md +++ b/crates/video/README.md @@ -21,7 +21,7 @@ No codec feature is on by default. | `dnx` | [`dnx`](src/dnx/README.md) | Avid DNxHD / DNxHR (SMPTE ST 2019-1 VC-3): decoder and encoder | | `apv` | [`apv`](src/apv/README.md) | APV (Advanced Professional Video, RFC 9924): decoder (`apv-decode` alone) and encoder (`apv-encode` alone) | | `vp8` | `vp8` | VP8 (RFC 6386): decoder | -| `mpeg2` | [`mpeg2`](src/mpeg2/README.md) | MPEG-2 video (ITU-T H.262) and MPEG-1 video (ISO/IEC 11172-2): decoder | +| `mpeg2` | [`mpeg2`](src/mpeg2/README.md) | MPEG-2 video (ITU-T H.262) and MPEG-1 video (ISO/IEC 11172-2): decoder (`mpeg2-decode` alone); MPEG-2 encoder with I and P pictures (`mpeg2-encode`, which also turns on the decoder) | | `mpeg4` | [`mpeg4`](src/mpeg4/README.md) | MPEG-4 Part 2 video (ISO/IEC 14496-2, the Xvid/DivX format): Simple profile encoder (`mpeg4-encode`) | | `prores` | `prores` | Apple ProRes (SMPTE RDD 36): decoder and encoder | | `h264` | [`h264`](src/h264/decode/README.md) | ITU-T H.264 / AVC: decoder (`h264-decode` alone) and encoder (`h264-encode` alone) | diff --git a/crates/video/src/lib.rs b/crates/video/src/lib.rs index e8e9d8113..d9723a81d 100644 --- a/crates/video/src/lib.rs +++ b/crates/video/src/lib.rs @@ -8,7 +8,7 @@ //! | `dnx` | [`dnx`] | Avid DNxHD / DNxHR (SMPTE ST 2019-1): decoder and encoder | //! | `apv` | [`apv`] | APV (RFC 9924): decoder (`apv-decode`) and encoder (`apv-encode`) | //! | `vp8` | [`vp8`] | VP8 (RFC 6386): decoder (`vp8-decode`) and key-frame encoder (`vp8-encode`) | -//! | `mpeg2` | [`mpeg2`] | MPEG-2 video (ITU-T H.262) and MPEG-1 video: decoder | +//! | `mpeg2` | [`mpeg2`] | MPEG-2 video (ITU-T H.262) and MPEG-1 video: decoder (`mpeg2-decode`) and MPEG-2 I/P encoder (`mpeg2-encode`) | //! | `mpeg4` | [`mpeg4`] | MPEG-4 Part 2 video (ISO/IEC 14496-2, Xvid/DivX): Simple profile encoder (`mpeg4-encode`) | //! | `prores` | [`prores`] | Apple ProRes (SMPTE RDD 36): decoder and encoder | //! | `h264` | [`h264`] | ITU-T H.264 / AVC: decoder (`h264-decode`) and encoder (`h264-encode`) | @@ -35,7 +35,7 @@ pub mod dnx; #[cfg(any(feature = "vp8-decode", feature = "vp8-encode"))] pub mod vp8; -#[cfg(feature = "mpeg2")] +#[cfg(feature = "mpeg2-decode")] pub mod mpeg2; #[cfg(feature = "mpeg4-encode")] diff --git a/crates/video/src/mpeg2/README.md b/crates/video/src/mpeg2/README.md index 1c49d45cd..b68dc800b 100644 --- a/crates/video/src/mpeg2/README.md +++ b/crates/video/src/mpeg2/README.md @@ -1,6 +1,7 @@ -# MPEG-2 and MPEG-1 video decoder (`craft_video::mpeg2`, feature `mpeg2`) +# MPEG-2 and MPEG-1 video decoder, MPEG-2 encoder (`craft_video::mpeg2`, features `mpeg2-decode`, `mpeg2-encode`) -Clean-room MPEG-2 video and MPEG-1 video decoder in pure Rust. Depends only on +Clean-room MPEG-2 video and MPEG-1 video decoder, and the first step of an MPEG-2 video encoder +(I and P pictures, see "Encoder"), in pure Rust. Depends only on `thiserror` and (optional, feature `threads`) `rayon`. Builds for `wasm32-unknown-unknown`; no `unsafe`. No GPL/LGPL code (FFmpeg, libmpeg2, libavcodec) was consulted; FFmpeg is used only as an external fixture generator and test oracle. @@ -11,6 +12,8 @@ consulted; FFmpeg is used only as an external fixture generator and test oracle. |---|---|---| | ITU-T H.262 \| ISO/IEC 13818-2 | 02/2000 (second edition, with Amd. 1 and Cor. 1/2 as incorporated in the 2000 text) | sequence / GOP / picture headers and extensions, slices, macroblocks, VLC tables (Annex B), inverse quantisation and mismatch control (§7.4), IDCT (Annex A), motion vectors and prediction (§7.6) | | ISO/IEC 11172-2 | 1993 | MPEG-1 video: slices spanning rows, MPEG-1 escape codes, oddification, full-pel vectors, D pictures | +| ITU-T H.262 \| ISO/IEC 13818-2 §8 | 02/2000 | profile and level limits the encoder chooses from (Tables 8-8 to 8-13) | +| MPEG-2 Test Model 5 (ISO/IEC JTC1/SC29/WG11/N0400) | April 1993 | encoder: target bit allocation and adaptive quantisation by activity (§5.2) | | IEEE Std 1180 | 1990 | IDCT accuracy test procedure (run as a unit test) | | ISO/IEC TR 13818-5 (MPEG Software Simulation Group decoder behaviour) | 1997 | two decoder behaviours the normative text leaves open (see "Interpretations") | @@ -72,6 +75,58 @@ the synthetic-stream tests confirmed against FFmpeg: formula only when the prediction is out of range: a field vector's doubled vertical component predicting a following frame vector. +## Encoder (feature `mpeg2-encode`) + +The encoder is the first step of MPEG-2 encoding in this crate, tracked in +storytold/craft-libs#1133. It writes valid MPEG-2 elementary streams, but not yet DVD or +broadcast deliverables: those need the later steps listed below. + +```rust +use craft_video::mpeg2::{Encoder, EncoderConfig, RateControl, YuvFrame}; + +let mut cfg = EncoderConfig::new(720, 576, 25, 1); +cfg.rate = RateControl::Abr { kbps: 6000 }; // or RateControl::Quantiser(1..=31) +cfg.gop_size = 12; // one I picture, then P pictures +let mut enc = Encoder::new(cfg)?; +for (i, f) in frames.iter().enumerate() { + let frame = YuvFrame { y: &f.y, cb: &f.cb, cr: &f.cr, y_stride: 720, c_stride: 360 }; + for p in enc.encode(&frame, i as i64)? { /* p.data, p.pts, p.dts, p.keyframe, p.picture_type */ } +} +for p in enc.flush() { /* none yet: no B pictures */ } +es.extend_from_slice(enc.end_of_sequence()); +``` + +What it does: + +- Main Profile (4:2:0) at Low, Main, High-1440 and High Level, and 4:2:2 Profile at Main and + High Level. The lowest level that holds the size, frame rate and (average) bitrate is chosen; + sizes beyond High Level are refused. Any even picture size; edges are padded by repetition. +- Progressive frame pictures, I and P, in closed GOPs. Every GOP starts with a sequence header, + sequence extension and GOP header (with a time code), so every I picture is a random access + point. Packets come out in decoding order, which equals display order without B pictures. +- Forward DCT on the transpose of the decoder's IDCT basis; intra quantisation with a small + dead zone, non-intra quantisation by truncation; default or custom quantiser matrices + (sequence header); linear or non-linear quantiser scale; intra DC precision 8-10 bits; + intra VLC table B.15. +- P pictures: integer pattern search (steps 8, 4, 2, 1) from candidate vectors (left neighbour, + predictor, the previous picture's vectors), then half-sample refinement; cost SAD + λ·bits. + Each macroblock is coded intra, motion-compensated with or without residual, "no MC" with + residual, or skipped. The search range sets f_code (1-4). +- Adaptive quantisation by macroblock activity (Test Model 5 §5.2.3). +- Rate control: a constant quantiser, or single-pass average bitrate with Test Model 5's + target bit allocation at picture level. Streams are signalled as variable bitrate + (`vbv_delay` 0xFFFF, the level's maximum bitrate and VBV size in the sequence header); the VBV + buffer is not modelled. +- The reconstruction uses the decoder's own dequantisation, mismatch control, IDCT and motion + compensation, so this crate's decoder outputs exactly the encoder's reconstruction. +- Slices (one per macroblock row) are coded in parallel on rayon with `threads`; the stream is + identical either way. + +Later steps (#1133): B pictures and reordering; CBR and VBR with a VBV guarantee; DVD presets +(sizes, aspect, GOP length, 9.8 Mb/s, 23.976 soft telecine); interlaced coding (field/frame DCT, +field prediction, alternate scan) and 4:2:2 High Level CBR for XDCAM HD422. MPEG-2 essence in the +MXF writer and a program stream muxer belong to `craft-media`. + ## Accuracy (vs FFmpeg as an external oracle) MPEG-2's IDCT is not bit-exact by definition (any IEEE 1180-accurate IDCT conforms, and the @@ -131,3 +186,14 @@ fell to 77–150 fps). These are lower bounds. 3000 random mutations never panic. - `synth_tests.rs`, `tests/mpeg2/oracle.rs`: see above (single-threaded decode identical to the threaded one). +- `encode/`: every encoder code decodes to its value with the decoder's tables; motion vector + deltas survive the decoder's wrap for f_code 1-4; forward/inverse DCT and quantiser round trips; + rate control converges on its target. `encode/tests.rs`: our decoder outputs exactly the + encoder's reconstruction (intra only, a 30-picture P GOP, 4:2:2, linear and non-linear + quantisers, quantisers 1 and 31, DC precision 8 and 10, custom matrices, odd macroblock counts, + average bitrate); quality and size follow the quantiser; repeated pictures skip; levels and + aspect in the sequence header; bad configurations and short or inconsistent planes are refused. +- `tests/mpeg2/encode.rs`: FFmpeg decodes our streams of `testsrc2` (CIF and 576-line 4:2:2) + without errors and within the decoder suite's tolerance of our reconstruction, ffprobe reports + the profile, level, pixel format and I/P picture types, and average bitrate lands within 10 % of + 800 and 2500 kbit/s targets. diff --git a/crates/video/src/mpeg2/encode/bitwriter.rs b/crates/video/src/mpeg2/encode/bitwriter.rs new file mode 100644 index 000000000..75b363d47 --- /dev/null +++ b/crates/video/src/mpeg2/encode/bitwriter.rs @@ -0,0 +1,92 @@ +//! Big-endian bit writer for start-code units. + +pub(crate) struct BitWriter { + out: Vec, + acc: u64, + /// Bits held in `acc` (always below 8 between calls). + n: u32, +} + +impl BitWriter { + pub fn new() -> Self { + Self { out: Vec::new(), acc: 0, n: 0 } + } + + /// Append the low `bits` (0..=32) bits of `v`, most significant first. + #[inline(always)] + pub fn put(&mut self, v: u32, bits: u32) { + if bits == 0 { + return; + } + let mask = if bits == 32 { u32::MAX } else { (1u32 << bits) - 1 }; + self.acc = (self.acc << bits) | u64::from(v & mask); + self.n += bits; + while self.n >= 8 { + self.n -= 8; + self.out.push((self.acc >> self.n) as u8); + } + self.acc &= (1u64 << self.n) - 1; + } + + #[inline(always)] + pub fn bit(&mut self, b: bool) { + self.put(b as u32, 1); + } + + /// Pad with zero bits to a byte boundary. + pub fn align(&mut self) { + if self.n > 0 { + self.put(0, 8 - self.n); + } + } + + /// Byte-align and write the start code `00 00 01 code`. + pub fn start_code(&mut self, code: u8) { + self.align(); + self.out.extend_from_slice(&[0, 0, 1, code]); + } + + /// Bits written so far. + #[cfg(test)] + pub fn bits(&self) -> u64 { + self.out.len() as u64 * 8 + u64::from(self.n) + } + + /// The bytes, padded to a byte boundary. + pub fn finish(mut self) -> Vec { + self.align(); + self.out + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::mpeg2::bits::Bits; + + #[test] + fn round_trips_through_the_reader() { + let mut w = BitWriter::new(); + w.put(0b101, 3); + w.put(0x1234_5678, 32); + w.bit(true); + w.put(0x3FF, 10); + assert_eq!(w.bits(), 46); + let d = w.finish(); + assert_eq!(d.len(), 6); + let mut b = Bits::new(&d); + assert_eq!(b.read(3), 0b101); + assert_eq!(b.read(32), 0x1234_5678); + assert!(b.bit()); + assert_eq!(b.read(10), 0x3FF); + assert_eq!(b.read(2), 0); + } + + #[test] + fn start_codes_are_byte_aligned() { + let mut w = BitWriter::new(); + w.put(1, 1); + w.start_code(0xB3); + assert_eq!(w.finish(), vec![0x80, 0, 0, 1, 0xB3]); + } +} diff --git a/crates/video/src/mpeg2/encode/mod.rs b/crates/video/src/mpeg2/encode/mod.rs new file mode 100644 index 000000000..fcd1bd8d1 --- /dev/null +++ b/crates/video/src/mpeg2/encode/mod.rs @@ -0,0 +1,624 @@ +//! MPEG-2 video encoder (ITU-T H.262 | ISO/IEC 13818-2), feature `mpeg2-encode`. +//! +//! This is the first step of MPEG-2 encoding in this crate (storytold/craft-libs#1133): I and P +//! pictures, progressive frame pictures, 4:2:0 (Main Profile) and 4:2:2 (4:2:2 Profile), closed +//! GOPs, constant quantiser or single-pass average bitrate. B pictures, a VBV guarantee, +//! interlaced coding and DVD / broadcast presets come later. +//! +//! ```no_run +//! use craft_video::mpeg2::encode::{Encoder, EncoderConfig, RateControl, YuvFrame}; +//! +//! let mut cfg = EncoderConfig::new(720, 576, 25, 1); +//! cfg.rate = RateControl::Abr { kbps: 6000 }; +//! let mut enc = Encoder::new(cfg)?; +//! let (y, c) = (vec![128u8; 720 * 576], vec![128u8; 360 * 288]); +//! let mut es = Vec::new(); +//! for i in 0..50 { +//! let frame = YuvFrame { y: &y, cb: &c, cr: &c, y_stride: 720, c_stride: 360 }; +//! for p in enc.encode(&frame, i)? { +//! es.extend_from_slice(&p.data); +//! } +//! } +//! for p in enc.flush() { +//! es.extend_from_slice(&p.data); +//! } +//! es.extend_from_slice(enc.end_of_sequence()); +//! # Ok::<(), craft_video::mpeg2::encode::EncodeError>(()) +//! ``` + +mod bitwriter; +mod motion; +mod picture; +mod rate; +mod tables; +mod transform; + +#[cfg(test)] +mod tests; + +use crate::mpeg2::ChromaFormat; +use crate::mpeg2::headers::{DEFAULT_INTRA, PictureType, SCAN}; +use crate::mpeg2::mc::Frame; + +use bitwriter::BitWriter; +use picture::{PicCtx, RowOut, RowPlanes, encode_row, qcode_for, qscale}; +use rate::RateCtl; + +#[derive(Debug, thiserror::Error)] +pub enum EncodeError { + #[error("unsupported MPEG-2 encoder configuration: {0}")] + Config(String), + #[error("bad input frame: {0}")] + BadFrame(String), +} + +/// aspect_ratio_information (H.262 Table 6-3). +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum AspectRatio { + /// Square samples. + Square, + /// 4:3 display. + Dar4x3, + /// 16:9 display. + Dar16x9, + /// 2.21:1 display. + Dar221x100, +} + +impl AspectRatio { + fn code(self) -> u32 { + match self { + AspectRatio::Square => 1, + AspectRatio::Dar4x3 => 2, + AspectRatio::Dar16x9 => 3, + AspectRatio::Dar221x100 => 4, + } + } +} + +#[derive(Clone, Copy, Debug, PartialEq)] +pub enum RateControl { + /// A constant quantiser_scale_code (1..=31; mapped through H.262 Table 7-6 with + /// `nonlinear_quant`). Adaptive quantisation still varies it per macroblock when enabled. + Quantiser(u8), + /// Single-pass average bitrate in kbit/s. The stream is signalled as variable bitrate and + /// is not checked against the VBV buffer. + Abr { kbps: u32 }, +} + +#[derive(Clone, Debug)] +pub struct EncoderConfig { + /// Picture size: even, at most 1920×1152 (4:2:0) or 1920×1088 (4:2:2). + pub width: u32, + pub height: u32, + /// Frame rate: one of H.262 Table 6-4 (23.976, 24, 25, 29.97, 30, 50, 59.94, 60). + pub fps_num: u32, + pub fps_den: u32, + /// 4:2:0 (Main Profile) or 4:2:2 (4:2:2 Profile). + pub chroma: ChromaFormat, + pub aspect: AspectRatio, + /// Pictures per GOP: one I picture, then P pictures. 1 codes every picture as I. 1..=1023. + pub gop_size: u32, + pub rate: RateControl, + /// q_scale_type: the non-linear quantiser scale (finer steps at low quantisers). + pub nonlinear_quant: bool, + /// Vary the quantiser per macroblock with its activity (finer in flat areas). + pub adaptive_quant: bool, + /// Intra DC precision in bits: 8, 9 or 10. + pub intra_dc_bits: u8, + /// Quantiser matrices (raster order, values 1..=255); `None`: the defaults. + pub intra_matrix: Option<[u8; 64]>, + pub non_intra_matrix: Option<[u8; 64]>, + /// Largest motion vector searched, in luma samples (1..=63). + pub search_range: u32, + /// Code slices on the rayon pool (with the `threads` feature). + pub threads: bool, +} + +impl EncoderConfig { + pub fn new(width: u32, height: u32, fps_num: u32, fps_den: u32) -> Self { + Self { + width, + height, + fps_num, + fps_den, + chroma: ChromaFormat::Yuv420, + aspect: AspectRatio::Square, + gop_size: 12, + rate: RateControl::Quantiser(4), + nonlinear_quant: true, + adaptive_quant: true, + intra_dc_bits: 9, + intra_matrix: None, + non_intra_matrix: None, + search_range: 31, + threads: cfg!(feature = "threads"), + } + } +} + +/// Planar 8-bit input: luma `width × height`, chroma `width/2 × height/2` (4:2:0) or +/// `width/2 × height` (4:2:2), rows `y_stride` / `c_stride` bytes apart. +#[derive(Clone, Copy, Debug)] +pub struct YuvFrame<'a> { + pub y: &'a [u8], + pub cb: &'a [u8], + pub cr: &'a [u8], + pub y_stride: usize, + pub c_stride: usize, +} + +/// One coded picture, in decoding order. +#[derive(Clone, Debug)] +pub struct Packet { + /// Elementary stream bytes: the sequence header and GOP header before every I picture, then + /// the picture. + pub data: Vec, + pub pts: i64, + /// Decoding timestamp; equal to `pts` while there are no B pictures. + pub dts: i64, + /// An I picture starting a closed GOP (random access point). + pub keyframe: bool, + pub picture_type: PictureType, + /// Average quantiser_scale of the picture's macroblocks. + pub quantiser_scale: f32, +} + +/// The encoder's reconstruction of a picture (what a decoder outputs), cropped. +#[derive(Clone, Debug)] +pub struct ReconFrame { + pub width: u32, + pub height: u32, + pub y: Vec, + pub cb: Vec, + pub cr: Vec, +} + +/// The sequence_end_code. +const SEQUENCE_END: [u8; 4] = [0, 0, 1, 0xB7]; + +/// A profile and level (H.262 Tables 8-8 to 8-13): limits and profile_and_level_indication. +#[derive(Clone, Copy, Debug)] +struct Level { + indication: u8, + max_w: u32, + max_h: u32, + max_fps: u32, + luma_rate: u64, + bit_rate: u64, + vbv_bits: u64, +} + +const MAIN_LEVELS: [Level; 4] = [ + Level { indication: 0x4A, max_w: 352, max_h: 288, max_fps: 30, luma_rate: 3_041_280, bit_rate: 4_000_000, vbv_bits: 475_136 }, + Level { indication: 0x48, max_w: 720, max_h: 576, max_fps: 30, luma_rate: 10_368_000, bit_rate: 15_000_000, vbv_bits: 1_835_008 }, + Level { indication: 0x46, max_w: 1440, max_h: 1152, max_fps: 60, luma_rate: 47_001_600, bit_rate: 60_000_000, vbv_bits: 7_340_032 }, + Level { indication: 0x44, max_w: 1920, max_h: 1152, max_fps: 60, luma_rate: 62_668_800, bit_rate: 80_000_000, vbv_bits: 9_781_248 }, +]; + +const PROFILE_422_LEVELS: [Level; 2] = [ + Level { indication: 0x85, max_w: 720, max_h: 608, max_fps: 30, luma_rate: 11_059_200, bit_rate: 50_000_000, vbv_bits: 9_437_184 }, + Level { indication: 0x82, max_w: 1920, max_h: 1088, max_fps: 60, luma_rate: 62_668_800, bit_rate: 300_000_000, vbv_bits: 47_185_920 }, +]; + +/// frame_rate_code for an exact frame rate (Table 6-4). +fn frame_rate_code(num: u32, den: u32) -> Option { + if num == 0 || den == 0 { + return None; + } + let rates: [(u64, u64); 8] = [(24000, 1001), (24, 1), (25, 1), (30000, 1001), (30, 1), (50, 1), (60000, 1001), (60, 1)]; + rates.iter().position(|&(n, d)| u64::from(num) * d == n * u64::from(den)).map(|i| i as u32 + 1) +} + +/// The motion search's f_code: the smallest whose range [-8f, 8f-1] samples covers `range`. +fn f_code_for(range: u32) -> u8 { + (1..=4u8).find(|&fc| 8 * (1u32 << (fc - 1)) > range).unwrap_or(4) +} + +/// Pick the lowest level of the profile that holds the stream. +fn choose_level(cfg: &EncoderConfig) -> Result { + let levels: &[Level] = match cfg.chroma { + ChromaFormat::Yuv420 => &MAIN_LEVELS, + ChromaFormat::Yuv422 => &PROFILE_422_LEVELS, + ChromaFormat::Yuv444 => return Err(EncodeError::Config("4:4:4 has no MPEG-2 profile in use".into())), + }; + let fps = f64::from(cfg.fps_num) / f64::from(cfg.fps_den); + let luma_rate = f64::from(cfg.width) * f64::from(cfg.height) * fps; + let bps = match cfg.rate { + RateControl::Abr { kbps } => u64::from(kbps) * 1000, + RateControl::Quantiser(_) => 0, + }; + levels + .iter() + .find(|l| { + cfg.width <= l.max_w + && cfg.height <= l.max_h + && fps <= f64::from(l.max_fps) + 0.01 + && luma_rate <= l.luma_rate as f64 + && bps <= l.bit_rate + }) + .copied() + .ok_or_else(|| { + EncodeError::Config(format!( + "{}x{} at {:.3} fps{} is beyond every level of the profile", + cfg.width, + cfg.height, + fps, + if bps > 0 { format!(", {bps} bit/s") } else { String::new() } + )) + }) +} + +fn validate(cfg: &EncoderConfig) -> Result<(), EncodeError> { + let bad = |s: String| Err(EncodeError::Config(s)); + if cfg.width == 0 || cfg.height == 0 || !cfg.width.is_multiple_of(2) || !cfg.height.is_multiple_of(2) { + return bad(format!("picture size {}x{} must be even and non-zero", cfg.width, cfg.height)); + } + if frame_rate_code(cfg.fps_num, cfg.fps_den).is_none() { + return bad(format!("frame rate {}/{} is not one of H.262 Table 6-4", cfg.fps_num, cfg.fps_den)); + } + if !(1..=1023).contains(&cfg.gop_size) { + return bad(format!("GOP size {} (1..=1023)", cfg.gop_size)); + } + match cfg.rate { + RateControl::Quantiser(q) if !(1..=31).contains(&q) => return bad(format!("quantiser_scale_code {q} (1..=31)")), + RateControl::Abr { kbps: 0 } => return bad("bitrate 0".into()), + _ => {} + } + if !(8..=10).contains(&cfg.intra_dc_bits) { + return bad(format!("intra DC precision {} bits (8..=10)", cfg.intra_dc_bits)); + } + if [cfg.intra_matrix, cfg.non_intra_matrix].iter().flatten().any(|m| m.contains(&0)) { + return bad("quantiser matrix with a zero entry".into()); + } + if !(1..=63).contains(&cfg.search_range) { + return bad(format!("search range {} (1..=63)", cfg.search_range)); + } + Ok(()) +} + +/// Write a quantiser matrix (raster order) in zigzag order. +fn put_matrix(w: &mut BitWriter, m: &[u8; 64]) { + for &pos in &SCAN[0] { + w.put(u32::from(m[usize::from(pos)]), 8); + } +} + +/// A stateful encoder. Feed pictures in display order; packets come out in decoding order. +pub struct Encoder { + cfg: EncoderConfig, + level: Level, + mb_w: usize, + mb_h: usize, + shx: u32, + shy: u32, + f_code: u8, + intra_matrix: [u8; 64], + inter_matrix: [u8; 64], + /// Sequence header and sequence extension, sent before every GOP. + seq_header: Vec, + /// The reconstruction of the previous picture. + reference: Option, + prev_mvs: Vec<[i32; 2]>, + /// Pictures coded so far. + count: u64, + /// Position of the next picture in its GOP. + gop_pos: u32, + rc: Option, + recon_capture: bool, + recon: Vec, +} + +impl Encoder { + pub fn new(cfg: EncoderConfig) -> Result { + validate(&cfg)?; + let level = choose_level(&cfg)?; + let (shx, shy) = if cfg.chroma == ChromaFormat::Yuv422 { (1, 0) } else { (1, 1) }; + let nonlinear = cfg.nonlinear_quant; + let q_range = (f64::from(qscale(nonlinear, 1)), f64::from(qscale(nonlinear, 31))); + let rc = match cfg.rate { + RateControl::Abr { kbps } => Some(RateCtl::new(f64::from(kbps) * 1000.0, f64::from(cfg.fps_num) / f64::from(cfg.fps_den), q_range)), + RateControl::Quantiser(_) => None, + }; + let mut enc = Encoder { + level, + mb_w: (cfg.width as usize).div_ceil(16), + mb_h: (cfg.height as usize).div_ceil(16), + shx, + shy, + f_code: f_code_for(cfg.search_range), + intra_matrix: cfg.intra_matrix.unwrap_or(DEFAULT_INTRA), + inter_matrix: cfg.non_intra_matrix.unwrap_or([16; 64]), + seq_header: Vec::new(), + reference: None, + prev_mvs: Vec::new(), + count: 0, + gop_pos: 0, + rc, + recon_capture: false, + recon: Vec::new(), + cfg, + }; + enc.seq_header = enc.sequence_header_bytes(); + Ok(enc) + } + + pub fn config(&self) -> &EncoderConfig { + &self.cfg + } + + /// profile_and_level_indication of the stream (e.g. 0x48 Main Profile @ Main Level). + pub fn profile_and_level(&self) -> u8 { + self.level.indication + } + + /// The sequence header and sequence extension (for containers that store them apart). + pub fn sequence_header(&self) -> &[u8] { + &self.seq_header + } + + /// The sequence_end_code to append after the last packet of an elementary stream. + pub fn end_of_sequence(&self) -> &'static [u8] { + &SEQUENCE_END + } + + /// Keep each picture's reconstruction for [`Encoder::take_recon`]. + pub fn set_recon_capture(&mut self, on: bool) { + self.recon_capture = on; + } + + pub fn take_recon(&mut self) -> Vec { + std::mem::take(&mut self.recon) + } + + /// Pictures held for reordering (none yet: there are no B pictures). + pub fn flush(&mut self) -> Vec { + Vec::new() + } + + fn sequence_header_bytes(&self) -> Vec { + let c = &self.cfg; + let mut w = BitWriter::new(); + let bit_rate = self.level.bit_rate.div_ceil(400); + let vbv = self.level.vbv_bits / 16384; + w.start_code(0xB3); + w.put(c.width & 0xFFF, 12); + w.put(c.height & 0xFFF, 12); + w.put(c.aspect.code(), 4); + w.put(frame_rate_code(c.fps_num, c.fps_den).unwrap_or(3), 4); + w.put((bit_rate & 0x3FFFF) as u32, 18); + w.bit(true); // marker_bit + w.put((vbv & 0x3FF) as u32, 10); + w.bit(false); // constrained_parameters_flag + for m in [&c.intra_matrix, &c.non_intra_matrix] { + w.bit(m.is_some()); + if let Some(m) = m { + put_matrix(&mut w, m); + } + } + w.start_code(0xB5); + w.put(1, 4); // sequence extension + w.put(u32::from(self.level.indication), 8); + w.bit(true); // progressive_sequence + w.put(if c.chroma == ChromaFormat::Yuv422 { 2 } else { 1 }, 2); + w.put(c.width >> 12, 2); + w.put(c.height >> 12, 2); + w.put((bit_rate >> 18) as u32 & 0xFFF, 12); + w.bit(true); // marker_bit + w.put((vbv >> 10) as u32 & 0xFF, 8); + w.bit(true); // low_delay: no B pictures + w.put(0, 2); // frame_rate_extension_n + w.put(0, 5); // frame_rate_extension_d + w.finish() + } + + /// group_of_pictures_header with the time code of picture `self.count`. + fn gop_header(&self, w: &mut BitWriter) { + let fps = u64::from(self.cfg.fps_num.div_ceil(self.cfg.fps_den)).max(1); + let n = self.count; + let secs = n / fps; + w.start_code(0xB8); + w.bit(false); // drop_frame_flag + w.put(((secs / 3600) % 24) as u32, 5); + w.put(((secs / 60) % 60) as u32, 6); + w.bit(true); // marker_bit + w.put((secs % 60) as u32, 6); + w.put((n % fps) as u32, 6); + w.bit(true); // closed_gop + w.bit(false); // broken_link + } + + fn picture_header(&self, w: &mut BitWriter, intra: bool) { + w.start_code(0x00); + w.put(self.gop_pos & 0x3FF, 10); // temporal_reference: display order within the GOP + w.put(if intra { 1 } else { 2 }, 3); + w.put(0xFFFF, 16); // vbv_delay: variable bitrate + if !intra { + w.bit(false); // full_pel_forward_vector + w.put(7, 3); // forward_f_code (MPEG-2: in the extension) + } + w.bit(false); // extra_bit_picture + w.start_code(0xB5); + w.put(8, 4); // picture coding extension + let fc = if intra { 15 } else { u32::from(self.f_code) }; + for v in [fc, fc, 15, 15] { + w.put(v, 4); + } + w.put(u32::from(self.cfg.intra_dc_bits - 8), 2); + w.put(3, 2); // frame picture + w.bit(false); // top_field_first + w.bit(true); // frame_pred_frame_dct + w.bit(false); // concealment_motion_vectors + w.bit(self.cfg.nonlinear_quant); // q_scale_type + w.bit(true); // intra_vlc_format + w.bit(false); // alternate_scan + w.bit(false); // repeat_first_field + w.bit(self.cfg.chroma == ChromaFormat::Yuv420); // chroma_420_type = progressive_frame + w.bit(true); // progressive_frame + w.bit(false); // composite_display_flag + } + + /// Copy the input into a macroblock-aligned frame, repeating the last column and row. + fn load(&self, f: &YuvFrame) -> Result { + let (w, h) = (self.cfg.width as usize, self.cfg.height as usize); + let (cw, ch) = (w >> self.shx, h >> self.shy); + let (pw, ph) = (self.mb_w * 16, self.mb_h * 16); + let check = |name: &str, plane: &[u8], stride: usize, pw: usize, ph: usize| { + let need = stride.checked_mul(ph - 1).and_then(|n| n.checked_add(pw)); + if stride < pw || need.is_none_or(|n| plane.len() < n) { + return Err(EncodeError::BadFrame(format!("{name}: {} bytes with stride {stride} for {pw}x{ph}", plane.len()))); + } + Ok(()) + }; + check("luma", f.y, f.y_stride, w, h)?; + check("Cb", f.cb, f.c_stride, cw, ch)?; + check("Cr", f.cr, f.c_stride, cw, ch)?; + let mut out = Frame::new(pw, ph, pw >> self.shx, ph >> self.shy); + for (k, (src, stride, sw, sh)) in [(f.y, f.y_stride, w, h), (f.cb, f.c_stride, cw, ch), (f.cr, f.c_stride, cw, ch)].into_iter().enumerate() { + let dw = if k == 0 { pw } else { out.cw }; + let dh = if k == 0 { ph } else { out.ch }; + let dst = &mut out.planes[k]; + for r in 0..dh { + let s = &src[r.min(sh - 1) * stride..][..sw]; + let d = &mut dst[r * dw..(r + 1) * dw]; + d[..sw].copy_from_slice(s); + let last = s[sw - 1]; + d[sw..].fill(last); + } + } + Ok(out) + } + + /// quantiser_scale_code per macroblock: `base` scaled by each macroblock's activity + /// (H.262 Test Model 5 §5.2.3) with adaptive quantisation. + fn mb_quant(&self, src: &Frame, base: f64) -> Vec { + let n = self.mb_w * self.mb_h; + let nl = self.cfg.nonlinear_quant; + if !self.cfg.adaptive_quant { + return vec![qcode_for(nl, base); n]; + } + let acts: Vec = (0..n).map(|a| activity(src, a % self.mb_w, a / self.mb_w)).collect(); + let avg = acts.iter().sum::() / n.max(1) as f64; + acts.iter().map(|&a| qcode_for(nl, base * (2.0 * a + avg) / (a + 2.0 * avg))).collect() + } + + /// Encode one picture. Returns the packets that became ready (one per picture for now). + pub fn encode(&mut self, frame: &YuvFrame, pts: i64) -> Result, EncodeError> { + let src = self.load(frame)?; + let reference = if self.gop_pos == 0 { None } else { self.reference.take() }; + let intra = reference.is_none(); + if intra { + self.gop_pos = 0; + if let Some(rc) = self.rc.as_mut() { + rc.start_gop(self.cfg.gop_size); + } + } + let base = match (&self.rc, self.cfg.rate) { + (Some(rc), _) => rc.qscale(intra), + (None, RateControl::Quantiser(q)) => f64::from(qscale(self.cfg.nonlinear_quant, q)), + (None, _) => 8.0, + }; + let mb_qcode = self.mb_quant(&src, base); + let mut recon = Frame::new(src.w, src.h, src.cw, src.ch); + let ctx = PicCtx { + src: &src, + reference: reference.as_ref(), + prev_mvs: &self.prev_mvs, + mb_w: self.mb_w, + shx: self.shx, + shy: self.shy, + intra_matrix: &self.intra_matrix, + inter_matrix: &self.inter_matrix, + nonlinear: self.cfg.nonlinear_quant, + dc_prec: self.cfg.intra_dc_bits - 8, + f_code: self.f_code, + search: self.cfg.search_range as i32, + mb_qcode: &mb_qcode, + }; + let rows = encode_rows(&ctx, &mut recon, self.cfg.threads); + let mut w = BitWriter::new(); + if intra { + self.gop_header(&mut w); + } + self.picture_header(&mut w, intra); + let mut data = if intra { self.seq_header.clone() } else { Vec::new() }; + data.extend_from_slice(&w.finish()); + let mut qsum = 0u64; + let mut mvs = Vec::with_capacity(self.mb_w * self.mb_h); + for r in rows { + data.extend_from_slice(&r.bytes); + qsum += r.stats.qsum; + mvs.extend_from_slice(&r.mvs); + } + let avg_q = qsum as f64 / (self.mb_w * self.mb_h).max(1) as f64; + if let Some(rc) = self.rc.as_mut() { + rc.update(intra, data.len() as u64 * 8, avg_q); + } + if self.recon_capture { + self.recon.push(self.crop(&recon)); + } + self.reference = Some(recon); + self.prev_mvs = mvs; + self.count += 1; + self.gop_pos = (self.gop_pos + 1) % self.cfg.gop_size; + let picture_type = if intra { PictureType::I } else { PictureType::P }; + Ok(vec![Packet { data, pts, dts: pts, keyframe: intra, picture_type, quantiser_scale: avg_q as f32 }]) + } + + fn crop(&self, f: &Frame) -> ReconFrame { + let (w, h) = (self.cfg.width as usize, self.cfg.height as usize); + let (cw, ch) = (w >> self.shx, h >> self.shy); + let crop = |p: &[u8], stride: usize, w: usize, h: usize| (0..h).flat_map(|r| p[r * stride..r * stride + w].iter().copied()).collect(); + ReconFrame { + width: w as u32, + height: h as u32, + y: crop(&f.planes[0], f.w, w, h), + cb: crop(&f.planes[1], f.cw, cw, ch), + cr: crop(&f.planes[2], f.cw, cw, ch), + } + } +} + +/// Macroblock activity: one plus the smallest variance of its four luma blocks (TM5). +fn activity(f: &Frame, mbx: usize, mby: usize) -> f64 { + let mut min = f64::MAX; + for b in 0..4 { + let (x0, y0) = (mbx * 16 + (b & 1) * 8, mby * 16 + (b >> 1) * 8); + let (mut s, mut s2) = (0u64, 0u64); + for r in 0..8 { + for &v in &f.planes[0][(y0 + r) * f.w + x0..][..8] { + s += u64::from(v); + s2 += u64::from(v) * u64::from(v); + } + } + let mean = s as f64 / 64.0; + min = min.min(s2 as f64 / 64.0 - mean * mean); + } + 1.0 + min.max(0.0) +} + +/// Code every macroblock row of the picture into `recon`, in parallel when allowed. +fn encode_rows(ctx: &PicCtx, recon: &mut Frame, threads: bool) -> Vec { + let (w, cw) = (recon.w, recon.cw); + let clines = 16 >> ctx.shy; + let [y, cb, cr] = &mut recon.planes; + let run = |r: usize, y: &mut [u8], cb: &mut [u8], cr: &mut [u8]| encode_row(ctx, r, &mut RowPlanes { y, cb, cr }); + #[cfg(feature = "threads")] + if threads { + use rayon::prelude::*; + return y + .par_chunks_mut(16 * w) + .zip(cb.par_chunks_mut(clines * cw)) + .zip(cr.par_chunks_mut(clines * cw)) + .enumerate() + .map(|(r, ((y, cb), cr))| run(r, y, cb, cr)) + .collect(); + } + let _ = threads; + y.chunks_mut(16 * w) + .zip(cb.chunks_mut(clines * cw)) + .zip(cr.chunks_mut(clines * cw)) + .enumerate() + .map(|(r, ((y, cb), cr))| run(r, y, cb, cr)) + .collect() +} diff --git a/crates/video/src/mpeg2/encode/motion.rs b/crates/video/src/mpeg2/encode/motion.rs new file mode 100644 index 000000000..a7bedc80b --- /dev/null +++ b/crates/video/src/mpeg2/encode/motion.rs @@ -0,0 +1,191 @@ +//! Motion search for frame prediction: integer pattern search from predictor candidates, then +//! half-sample refinement. Costs are SAD + λ·(motion vector bits). + +use crate::mpeg2::mc::{Block, Frame, predict}; + +use super::tables::{enc_tables, wrap_delta}; + +pub(crate) struct Search<'a> { + pub reference: &'a Frame, + /// The source macroblock's luma (16×16). + pub src: &'a [u8; 256], + pub mbx: usize, + pub mby: usize, + pub f_code: u8, + /// Largest integer displacement searched, in samples. + pub range: i32, + /// Motion vector predictor (half-sample units). + pub pmv: [i32; 2], + pub lambda: u32, +} + +/// The chosen vector (half-sample units) and its SAD. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) struct Found { + pub mv: [i32; 2], + pub sad: u32, +} + +/// Steps of the integer pattern search, coarse to fine. +const STEPS: [i32; 4] = [8, 4, 2, 1]; +/// Moves per step before moving on (bounds the search on pathological content). +const MAX_MOVES: usize = 16; + +impl Search<'_> { + fn origin(&self) -> (i32, i32) { + ((self.mbx * 16) as i32, (self.mby * 16) as i32) + } + + /// Integer displacements allowed per axis: inside the reference and the f_code range. + fn bounds(&self) -> ([i32; 2], [i32; 2]) { + let f = 1i32 << self.f_code.saturating_sub(1); + let (ox, oy) = self.origin(); + let (w, h) = (self.reference.w as i32, self.reference.h as i32); + let lo_lim = -(8 * f).min(self.range); + let hi_lim = (8 * f - 1).min(self.range); + ([lo_lim.max(-ox), lo_lim.max(-oy)], [hi_lim.min(w - 16 - ox), hi_lim.min(h - 16 - oy)]) + } + + /// A half-sample vector the bitstream can carry whose prediction stays inside the reference. + fn valid(&self, mv: [i32; 2]) -> bool { + let f = 1i32 << self.f_code.saturating_sub(1); + let (ox, oy) = self.origin(); + let size = [self.reference.w as i32, self.reference.h as i32]; + let o = [ox, oy]; + (0..2).all(|c| { + let start = o[c] + (mv[c] >> 1); + mv[c] >= -16 * f && mv[c] < 16 * f && start >= 0 && start + 16 + (mv[c] & 1) <= size[c] + }) + } + + fn mv_cost(&self, mv: [i32; 2]) -> u32 { + let t = enc_tables(); + let bits: u32 = (0..2).map(|c| t.motion_bits(wrap_delta(mv[c], self.pmv[c], self.f_code), self.f_code)).sum(); + self.lambda * bits + } + + fn sad_int(&self, d: [i32; 2]) -> u32 { + let (ox, oy) = self.origin(); + let stride = self.reference.w; + let x = (ox + d[0]) as usize; + let y = (oy + d[1]) as usize; + let plane = &self.reference.planes[0]; + let mut sad = 0u32; + for r in 0..16 { + let at = (y + r) * stride + x; + let Some(row) = plane.get(at..at + 16) else { return u32::MAX }; + for (a, b) in row.iter().zip(&self.src[r * 16..r * 16 + 16]) { + sad += u32::from(a.abs_diff(*b)); + } + } + sad + } + + fn sad_half(&self, mv: [i32; 2]) -> u32 { + let (ox, oy) = self.origin(); + let mut tmp = [0u8; 256]; + let blk = Block { x: ox + (mv[0] >> 1), y: oy + (mv[1] >> 1), hx: mv[0] & 1 != 0, hy: mv[1] & 1 != 0, bw: 16, bh: 16 }; + predict(&mut tmp, 0, 16, self.reference.plane(0, None), blk, false); + tmp.iter().zip(self.src.iter()).map(|(a, b)| u32::from(a.abs_diff(*b))).sum() + } + + /// Best vector among the candidates (half-sample units), refined. + pub fn run(&self, candidates: &[[i32; 2]]) -> Found { + let (lo, hi) = self.bounds(); + let clamp = |d: [i32; 2]| [d[0].clamp(lo[0], hi[0].max(lo[0])), d[1].clamp(lo[1], hi[1].max(lo[1]))]; + let eval = |d: [i32; 2]| { + let sad = self.sad_int(d); + (sad.saturating_add(self.mv_cost([2 * d[0], 2 * d[1]])), sad) + }; + let mut best_d = [0, 0]; + let (mut best_cost, mut best_sad) = eval(best_d); + for c in candidates { + let d = clamp([c[0] >> 1, c[1] >> 1]); + if d != best_d { + let (cost, sad) = eval(d); + if cost < best_cost { + (best_d, best_cost, best_sad) = (d, cost, sad); + } + } + } + for step in STEPS.into_iter().filter(|&s| s <= self.range.max(1)) { + for _ in 0..MAX_MOVES { + let mut moved = false; + for off in [[step, 0], [-step, 0], [0, step], [0, -step]] { + let d = clamp([best_d[0] + off[0], best_d[1] + off[1]]); + if d == best_d { + continue; + } + let (cost, sad) = eval(d); + if cost < best_cost { + (best_d, best_cost, best_sad) = (d, cost, sad); + moved = true; + } + } + if !moved { + break; + } + } + } + // half-sample refinement around the integer result + let centre = [2 * best_d[0], 2 * best_d[1]]; + let mut best = Found { mv: centre, sad: best_sad }; + for dy in -1..=1 { + for dx in -1..=1 { + let mv = [centre[0] + dx, centre[1] + dy]; + if (dx == 0 && dy == 0) || !self.valid(mv) { + continue; + } + let sad = self.sad_half(mv); + let cost = sad.saturating_add(self.mv_cost(mv)); + if cost < best_cost { + best_cost = cost; + best = Found { mv, sad }; + } + } + } + best + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// A textured reference and a source macroblock cut from it at a known offset. + fn scene(dx: i32, dy: i32) -> (Frame, [u8; 256]) { + let mut f = Frame::new(64, 64, 32, 32); + for y in 0..64 { + for x in 0..64 { + let (x, y) = (x as f32, y as f32); + f.planes[0][y as usize * 64 + x as usize] = + (128.0 + 60.0 * (x * 0.31).sin() * (y * 0.23).cos() + 30.0 * ((x + y) * 0.13).sin()) as u8; + } + } + let mut src = [0u8; 256]; + for r in 0..16 { + for c in 0..16 { + src[r * 16 + c] = f.planes[0][(16 + dy + r as i32) as usize * 64 + (16 + dx + c as i32) as usize]; + } + } + (f, src) + } + + #[test] + fn finds_an_integer_displacement() { + let (f, src) = scene(3, -2); + let s = Search { reference: &f, src: &src, mbx: 1, mby: 1, f_code: 3, range: 15, pmv: [0, 0], lambda: 1 }; + let found = s.run(&[]); + assert_eq!(found, Found { mv: [6, -4], sad: 0 }); + } + + #[test] + fn stays_inside_the_reference() { + let (f, src) = scene(0, 0); + for (mbx, mby) in [(0, 0), (3, 3), (0, 3)] { + let s = Search { reference: &f, src: &src, mbx, mby, f_code: 4, range: 63, pmv: [0, 0], lambda: 0 }; + let found = s.run(&[[-200, -200], [200, 200]]); + assert!(s.valid(found.mv), "{found:?} at {mbx},{mby}"); + } + } +} diff --git a/crates/video/src/mpeg2/encode/picture.rs b/crates/video/src/mpeg2/encode/picture.rs new file mode 100644 index 000000000..e7bb84b7e --- /dev/null +++ b/crates/video/src/mpeg2/encode/picture.rs @@ -0,0 +1,412 @@ +//! Slices and macroblocks of one frame picture (H.262 §6.2.4-§6.2.6), with a reconstruction +//! that matches the decoder exactly. One slice per macroblock row; slices only depend on the +//! reference picture, so rows can be coded in parallel. + +use crate::mpeg2::headers::SCAN; +use crate::mpeg2::mc::{Block, Frame, predict}; +use crate::mpeg2::slice::NON_LINEAR_Q; +use crate::mpeg2::vlc::{MB_FWD, MB_INTRA, MB_PATTERN, MB_QUANT}; + +use super::bitwriter::BitWriter; +use super::motion::Search; +use super::tables::{EncTables, enc_tables, wrap_delta}; +use super::transform::{fdct, quant_inter, quant_intra, reconstruct}; + +/// Intra coding is chosen when its cost (the luma's deviation from its block means) undercuts +/// the best prediction's SAD by this much. +const INTRA_BIAS: u32 = 256; + +/// quantiser_scale for quantiser_scale_code `code` (Table 7-6 when `nonlinear`). +pub(crate) fn qscale(nonlinear: bool, code: u8) -> i32 { + if nonlinear { i32::from(NON_LINEAR_Q[usize::from(code & 31)]) } else { 2 * i32::from(code & 31) } +} + +/// The quantiser_scale_code (1..=31) whose quantiser_scale is closest to `qs`. +pub(crate) fn qcode_for(nonlinear: bool, qs: f64) -> u8 { + let mut best = (1u8, f64::MAX); + for code in 1..=31u8 { + let d = (f64::from(qscale(nonlinear, code)) - qs).abs(); + if d < best.1 { + best = (code, d); + } + } + best.0 +} + +/// What every slice of one picture shares (read-only). +pub(crate) struct PicCtx<'a> { + /// The padded source picture. + pub src: &'a Frame, + /// The reconstructed previous picture (P pictures); `None` codes an I picture. + pub reference: Option<&'a Frame>, + /// Vectors of the previous picture (motion search candidates), one per macroblock. + pub prev_mvs: &'a [[i32; 2]], + pub mb_w: usize, + pub shx: u32, + pub shy: u32, + pub intra_matrix: &'a [u8; 64], + pub inter_matrix: &'a [u8; 64], + pub nonlinear: bool, + /// intra_dc_precision (0..=2: 8..=10 bits). + pub dc_prec: u8, + pub f_code: u8, + pub search: i32, + /// quantiser_scale_code per macroblock. + pub mb_qcode: &'a [u8], +} + +impl PicCtx<'_> { + fn cmb(&self) -> (usize, usize) { + (16 >> self.shx, 16 >> self.shy) + } + fn blocks(&self) -> usize { + if self.shy == 1 { 6 } else { 8 } + } + fn dc_reset(&self) -> i32 { + 1 << (7 + self.dc_prec) + } +} + +/// The part of the reconstructed picture one macroblock row covers. +pub(crate) struct RowPlanes<'b> { + pub y: &'b mut [u8], + pub cb: &'b mut [u8], + pub cr: &'b mut [u8], +} + +#[derive(Clone, Copy, Debug, Default)] +pub(crate) struct RowStats { + pub intra: u32, + pub skipped: u32, + /// Sum of the quantiser_scale of every macroblock. + pub qsum: u64, +} + +pub(crate) struct RowOut { + /// The slice, from its start code, byte aligned. + pub bytes: Vec, + /// The vector each macroblock used (zero for intra and skipped ones). + pub mvs: Vec<[i32; 2]>, + pub stats: RowStats, +} + +/// One macroblock's samples: luma 16×16, chroma rows of `16 >> shx`. +#[derive(Clone)] +struct MbPix { + y: [u8; 256], + c: [[u8; 256]; 2], +} + +impl MbPix { + fn load(f: &Frame, mbx: usize, mby: usize, (cbw, cbh): (usize, usize)) -> MbPix { + let mut m = MbPix { y: [0; 256], c: [[0; 256]; 2] }; + for r in 0..16 { + let at = (mby * 16 + r) * f.w + mbx * 16; + m.y[r * 16..r * 16 + 16].copy_from_slice(&f.planes[0][at..at + 16]); + } + for c in 0..2 { + for r in 0..cbh { + let at = (mby * cbh + r) * f.cw + mbx * cbw; + m.c[c][r * cbw..r * cbw + cbw].copy_from_slice(&f.planes[1 + c][at..at + cbw]); + } + } + m + } + + /// Plane, offset and stride of block `i` (H.262 Figure 6-10 order: four luma blocks, then + /// Cb, Cr, and for 4:2:2 the lower Cb, Cr). + fn block_at(i: usize, cbw: usize) -> (usize, usize, usize) { + if i < 4 { (0, (i >> 1) * 8 * 16 + (i & 1) * 8, 16) } else { (1 + (i & 1), ((i - 4) >> 1) * 8 * cbw, cbw) } + } + + fn plane(&self, p: usize) -> &[u8; 256] { + if p == 0 { &self.y } else { &self.c[p - 1] } + } + + fn plane_mut(&mut self, p: usize) -> &mut [u8; 256] { + if p == 0 { &mut self.y } else { &mut self.c[p - 1] } + } + + fn block(&self, i: usize, cbw: usize) -> [i32; 64] { + let (p, off, stride) = Self::block_at(i, cbw); + let d = self.plane(p); + std::array::from_fn(|k| i32::from(d[off + (k >> 3) * stride + (k & 7)])) + } + + /// Write `v` (intra samples) or add it to what is there (inter residual), saturated. + fn store(&mut self, i: usize, cbw: usize, v: &[i32; 64], add: bool) { + let (p, off, stride) = Self::block_at(i, cbw); + let d = self.plane_mut(p); + for (k, &x) in v.iter().enumerate() { + let at = off + (k >> 3) * stride + (k & 7); + let base = if add { i32::from(d[at]) } else { 0 }; + d[at] = (base + x).clamp(0, 255) as u8; + } + } + + fn put(&self, out: &mut RowPlanes, mbx: usize, (cbw, cbh): (usize, usize), (w, cw): (usize, usize)) { + for r in 0..16 { + let at = r * w + mbx * 16; + if let Some(d) = out.y.get_mut(at..at + 16) { + d.copy_from_slice(&self.y[r * 16..r * 16 + 16]); + } + } + for (c, plane) in [&mut *out.cb, &mut *out.cr].into_iter().enumerate() { + for r in 0..cbh { + let at = r * cw + mbx * cbw; + if let Some(d) = plane.get_mut(at..at + cbw) { + d.copy_from_slice(&self.c[c][r * cbw..r * cbw + cbw]); + } + } + } + } +} + +/// Slice coding state (what the decoder tracks too). +struct SliceSt { + qcode: u8, + dc: [i32; 3], + pmv: [i32; 2], + /// Skipped macroblocks not yet covered by an address increment. + skipped: usize, + /// The vector of the macroblock to the left (motion search candidate). + left_mv: [i32; 2], +} + +/// One macroblock being coded. +struct MbIn<'x> { + src: &'x MbPix, + mbx: usize, + mby: usize, + qcode: u8, +} + +/// Code macroblock row `row` as one slice, writing its reconstruction into `out`. +pub(crate) fn encode_row(ctx: &PicCtx, row: usize, out: &mut RowPlanes) -> RowOut { + let mut w = BitWriter::new(); + w.start_code(1 + row as u8); + let q0 = ctx.mb_qcode.get(row * ctx.mb_w).copied().unwrap_or(1); + w.put(u32::from(q0), 5); + w.bit(false); // extra_bit_slice + let mut st = SliceSt { qcode: q0, dc: [ctx.dc_reset(); 3], pmv: [0; 2], skipped: 0, left_mv: [0; 2] }; + let mut ro = RowOut { bytes: Vec::new(), mvs: vec![[0; 2]; ctx.mb_w], stats: RowStats::default() }; + for mbx in 0..ctx.mb_w { + macroblock(ctx, &mut st, &mut w, (mbx, row), out, &mut ro); + } + ro.bytes = w.finish(); + ro +} + +fn macroblock(ctx: &PicCtx, st: &mut SliceSt, w: &mut BitWriter, (mbx, mby): (usize, usize), out: &mut RowPlanes, ro: &mut RowOut) { + let addr = mby * ctx.mb_w + mbx; + let qcode = ctx.mb_qcode.get(addr).copied().unwrap_or(st.qcode); + let src = MbPix::load(ctx.src, mbx, mby, ctx.cmb()); + let mb = MbIn { src: &src, mbx, mby, qcode }; + let qs = qscale(ctx.nonlinear, qcode); + ro.stats.qsum += qs as u64; + let planes = (ctx.src.w, ctx.src.cw); + let Some(reference) = ctx.reference else { + intra_mb(ctx, st, w, &mb).put(out, mbx, ctx.cmb(), planes); + ro.stats.intra += 1; + return; + }; + let mut candidates = vec![st.left_mv, st.pmv]; + for a in [Some(addr), addr.checked_sub(ctx.mb_w), Some(addr + 1), Some(addr + ctx.mb_w)].into_iter().flatten() { + if let Some(&v) = ctx.prev_mvs.get(a) { + candidates.push(v); + } + } + let lambda = ((qs * 3 + 4) / 8).max(1) as u32; + let search = Search { reference, src: &src.y, mbx, mby, f_code: ctx.f_code, range: ctx.search, pmv: st.pmv, lambda }; + let found = search.run(&candidates); + if intra_cost(&src.y).saturating_add(INTRA_BIAS) < found.sad { + intra_mb(ctx, st, w, &mb).put(out, mbx, ctx.cmb(), planes); + ro.stats.intra += 1; + st.left_mv = [0; 2]; + return; + } + let (rec, skipped) = inter_mb(ctx, st, w, &mb, reference, found.mv); + rec.put(out, mbx, ctx.cmb(), planes); + ro.stats.skipped += skipped as u32; + ro.mvs[mbx] = found.mv; + st.left_mv = found.mv; +} + +/// Σ |sample − block mean| over the four luma blocks: the cost of coding them intra. +fn intra_cost(y: &[u8; 256]) -> u32 { + let mut cost = 0u32; + for b in 0..4 { + let off = (b >> 1) * 128 + (b & 1) * 8; + let at = |k: usize| u32::from(y[off + (k >> 3) * 16 + (k & 7)]); + let mean = ((0..64).map(at).sum::() + 32) / 64; + cost += (0..64).map(|k| at(k).abs_diff(mean)).sum::(); + } + cost +} + +fn put_increment(w: &mut BitWriter, t: &EncTables, mut n: usize) { + while n > 33 { + t.mba_escape.put(w); + n -= 33; + } + t.mba[n.max(1)].put(w); +} + +/// Write the coefficients of a block (raster-order levels) in zigzag order from `start`, then +/// end of block. +fn put_coefficients(w: &mut BitWriter, t: &EncTables, levels: &[i32; 64], table: usize, start: usize) { + let mut run = 0; + let mut first = start == 0; + for &pos in &SCAN[0][start..] { + let l = levels[usize::from(pos)]; + if l == 0 { + run += 1; + continue; + } + t.put_coefficient(w, table, run, l, first); + first = false; + run = 0; + } + t.dct_eob[table].put(w); +} + +fn put_mb_type(w: &mut BitWriter, t: &EncTables, predicted: bool, flags: i16) { + let table = if predicted { &t.mbtype_p } else { &t.mbtype_i }; + if let Some(c) = table.get(flags as usize).copied().flatten() { + c.put(w); + } +} + +/// Code an intra macroblock; returns its reconstruction. +fn intra_mb(ctx: &PicCtx, st: &mut SliceSt, w: &mut BitWriter, mb: &MbIn) -> MbPix { + let t = enc_tables(); + let quant = mb.qcode != st.qcode; + put_increment(w, t, st.skipped + 1); + st.skipped = 0; + put_mb_type(w, t, ctx.reference.is_some(), if quant { MB_INTRA | MB_QUANT } else { MB_INTRA }); + if quant { + w.put(u32::from(mb.qcode), 5); + st.qcode = mb.qcode; + } + st.pmv = [0; 2]; + let qs = qscale(ctx.nonlinear, st.qcode); + let dc_mult = 8 >> ctx.dc_prec; + let dc_max = (1 << (8 + ctx.dc_prec)) - 1; + let (cbw, _) = ctx.cmb(); + let mut rec = mb.src.clone(); + for i in 0..ctx.blocks() { + let cc = if i < 4 { 0 } else { 1 + (i & 1) }; + let f = fdct(&mb.src.block(i, cbw)); + let mut levels = [0i32; 64]; + levels[0] = ((f[0] / dc_mult as f32).round() as i32).clamp(0, dc_max); + for pos in 1..64 { + levels[pos] = quant_intra(f[pos], ctx.intra_matrix[pos], qs); + } + let diff = levels[0] - st.dc[cc]; + st.dc[cc] = levels[0]; + let size = 32 - diff.unsigned_abs().leading_zeros(); + let size_codes = if cc == 0 { &t.dc_luma } else { &t.dc_chroma }; + size_codes[size as usize].put(w); + if size > 0 { + let v = if diff > 0 { diff } else { diff + (1 << size) - 1 }; + w.put(v as u32, size); + } + put_coefficients(w, t, &levels, 1, 1); + rec.store(i, cbw, &reconstruct(&levels, ctx.intra_matrix, qs, true, dc_mult), false); + } + rec +} + +/// Code a predicted macroblock with vector `mv` (or skip it); returns its reconstruction and +/// whether it was skipped. +fn inter_mb(ctx: &PicCtx, st: &mut SliceSt, w: &mut BitWriter, mb: &MbIn, reference: &Frame, mv: [i32; 2]) -> (MbPix, bool) { + let t = enc_tables(); + let mut rec = predict_mb(ctx, reference, mv, mb.mbx, mb.mby); + let qs = qscale(ctx.nonlinear, mb.qcode); + let (cbw, _) = ctx.cmb(); + let n = ctx.blocks(); + let mut levels = [[0i32; 64]; 8]; + let mut cbp = 0u32; + for i in 0..n { + let (s, p) = (mb.src.block(i, cbw), rec.block(i, cbw)); + let f = fdct(&std::array::from_fn(|k| s[k] - p[k])); + for pos in 0..64 { + levels[i][pos] = quant_inter(f[pos], ctx.inter_matrix[pos], qs); + } + if levels[i].iter().any(|&l| l != 0) { + cbp |= 1 << (n - 1 - i); + } + } + st.dc = [ctx.dc_reset(); 3]; + // the first and last macroblock of a slice are always coded + let edge = mb.mbx == 0 || mb.mbx + 1 == ctx.mb_w; + if mv == [0, 0] && cbp == 0 && !edge { + st.skipped += 1; + st.pmv = [0; 2]; + return (rec, true); + } + put_increment(w, t, st.skipped + 1); + st.skipped = 0; + let quant = cbp != 0 && mb.qcode != st.qcode; + let mut flags = if cbp == 0 { + MB_FWD + } else if mv == [0, 0] { + MB_PATTERN + } else { + MB_FWD | MB_PATTERN + }; + if quant { + flags |= MB_QUANT; + } + put_mb_type(w, t, true, flags); + if quant { + w.put(u32::from(mb.qcode), 5); + st.qcode = mb.qcode; + } + if flags & MB_FWD != 0 { + for c in 0..2 { + t.put_motion(w, wrap_delta(mv[c], st.pmv[c], ctx.f_code), ctx.f_code); + } + st.pmv = mv; + } else { + st.pmv = [0; 2]; // "No MC" resets the predictors + } + if cbp != 0 { + if n == 8 { + t.cbp[(cbp >> 2) as usize].put(w); + w.put(cbp & 3, 2); + } else { + t.cbp[cbp as usize].put(w); + } + for i in 0..n { + if cbp & (1 << (n - 1 - i)) != 0 { + put_coefficients(w, t, &levels[i], 0, 0); + rec.store(i, cbw, &reconstruct(&levels[i], ctx.inter_matrix, qs, false, 0), true); + } + } + } + (rec, false) +} + +/// The frame prediction of a macroblock from `f` with vector `mv`, formed as the decoder forms +/// it (§7.6.4: chroma vectors are the luma ones halved towards zero). +fn predict_mb(ctx: &PicCtx, f: &Frame, mv: [i32; 2], mbx: usize, mby: usize) -> MbPix { + let (cbw, cbh) = ctx.cmb(); + let mut p = MbPix { y: [0; 256], c: [[0; 256]; 2] }; + let blk = + Block { x: (mbx * 16) as i32 + (mv[0] >> 1), y: (mby * 16) as i32 + (mv[1] >> 1), hx: mv[0] & 1 != 0, hy: mv[1] & 1 != 0, bw: 16, bh: 16 }; + predict(&mut p.y, 0, 16, f.plane(0, None), blk, false); + let cmv = [if ctx.shx == 1 { mv[0] / 2 } else { mv[0] }, if ctx.shy == 1 { mv[1] / 2 } else { mv[1] }]; + let cb = Block { + x: (mbx * cbw) as i32 + (cmv[0] >> 1), + y: (mby * cbh) as i32 + (cmv[1] >> 1), + hx: cmv[0] & 1 != 0, + hy: cmv[1] & 1 != 0, + bw: cbw, + bh: cbh, + }; + for c in 0..2 { + predict(&mut p.c[c], 0, cbw, f.plane(1 + c, None), cb, false); + } + p +} diff --git a/crates/video/src/mpeg2/encode/rate.rs b/crates/video/src/mpeg2/encode/rate.rs new file mode 100644 index 000000000..65c23efc3 --- /dev/null +++ b/crates/video/src/mpeg2/encode/rate.rs @@ -0,0 +1,102 @@ +//! Single-pass average bitrate control at picture level, after the target bit allocation of the +//! MPEG-2 Test Model 5: each picture type keeps a complexity `X = bits · average quantiser`, the +//! bits left for the GOP are shared out by complexity, and a picture's quantiser is `X / target`. +//! There is no VBV model: the average converges on the target, single pictures may not. + +pub(crate) struct RateCtl { + bps: f64, + fps: f64, + /// Complexity of I and P pictures. + xi: f64, + xp: f64, + /// Bits left for the current GOP (carries over between GOPs). + remaining: f64, + /// P pictures left in the current GOP. + np: u32, + /// Range of quantiser_scale values. + q_min: f64, + q_max: f64, +} + +impl RateCtl { + pub fn new(bps: f64, fps: f64, q_range: (f64, f64)) -> Self { + // TM5's starting complexities + Self { bps, fps, xi: 160.0 * bps / 115.0, xp: 60.0 * bps / 115.0, remaining: 0.0, np: 0, q_min: q_range.0, q_max: q_range.1 } + } + + /// A GOP of `pictures` pictures (one I, the rest P) starts. + pub fn start_gop(&mut self, pictures: u32) { + self.remaining += self.bps * f64::from(pictures) / self.fps; + self.np = pictures.saturating_sub(1); + } + + /// quantiser_scale for the next picture. + pub fn qscale(&self, intra: bool) -> f64 { + let floor = self.bps / (8.0 * self.fps); + let (target, x) = if intra { + (self.remaining / (1.0 + f64::from(self.np) * self.xp / self.xi), self.xi) + } else { + (self.remaining / f64::from(self.np.max(1)), self.xp) + }; + (x / target.max(floor)).clamp(self.q_min, self.q_max) + } + + /// The picture used `bits` at average quantiser_scale `avg_q`. + pub fn update(&mut self, intra: bool, bits: u64, avg_q: f64) { + let x = bits as f64 * avg_q; + if intra { + self.xi = x.max(1.0); + } else { + self.xp = x.max(1.0); + self.np = self.np.saturating_sub(1); + } + self.remaining -= bits as f64; + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// A picture whose size is `complexity / q` (the model the controller assumes). + fn run(complexity_i: f64, complexity_p: f64, bps: f64, gops: u32) -> f64 { + let fps = 25.0; + let mut rc = RateCtl::new(bps, fps, (1.0, 112.0)); + let mut total = 0.0; + for _ in 0..gops { + rc.start_gop(12); + for k in 0..12 { + let intra = k == 0; + let q = rc.qscale(intra); + let bits = (if intra { complexity_i } else { complexity_p }) / q; + rc.update(intra, bits as u64, q); + total += bits; + } + } + total * fps / (12.0 * f64::from(gops)) + } + + #[test] + fn converges_on_the_target() { + // complexities whose target quantiser is inside the allowed range + for (ci, cp) in [(4e7, 1e7), (8e7, 2e6), (1e7, 9e6)] { + let got = run(ci, cp, 4e6, 20); + assert!((got / 4e6 - 1.0).abs() < 0.03, "{ci} {cp}: {got}"); + } + } + + #[test] + fn the_quantiser_stays_in_range() { + let mut rc = RateCtl::new(1e6, 25.0, (2.0, 62.0)); + rc.start_gop(2); + // a picture far larger than its budget pushes the next quantiser to the top + rc.update(true, 50_000_000, 10.0); + rc.start_gop(2); + assert_eq!(rc.qscale(true), 62.0); + // and a tiny one to the bottom + let mut rc = RateCtl::new(1e6, 25.0, (2.0, 62.0)); + rc.start_gop(2); + rc.update(true, 10, 2.0); + assert_eq!(rc.qscale(true), 2.0); + } +} diff --git a/crates/video/src/mpeg2/encode/tables.rs b/crates/video/src/mpeg2/encode/tables.rs new file mode 100644 index 000000000..b79dd92ea --- /dev/null +++ b/crates/video/src/mpeg2/encode/tables.rs @@ -0,0 +1,264 @@ +//! Encoder-side variable-length codes, built from the decoder's code lists (H.262 Annex B) so the +//! two cannot disagree. + +use std::sync::OnceLock; + +use crate::mpeg2::vlc::{ + CBP, DC_CHROMA, DC_LUMA, DCT_EOB, DCT_ESCAPE, DCT_ONE, DCT_ZERO, DCT_ZERO_LONG, MBA, MBA_ESCAPE, MBTYPE_I, MBTYPE_P, MOTION, dct_table, + parse_code, +}; + +use super::bitwriter::BitWriter; + +/// A code: `len` bits of `bits`, most significant first. +#[derive(Clone, Copy, Default, Debug, PartialEq, Eq)] +pub(crate) struct Code { + pub bits: u32, + pub len: u8, +} + +impl Code { + fn of(s: &str) -> Code { + let (bits, len) = parse_code(s); + Code { bits, len: len as u8 } + } + + #[inline(always)] + pub fn put(self, w: &mut BitWriter) { + w.put(self.bits, u32::from(self.len)); + } +} + +/// Longest run with a table code (table zero and one: run 31 at level 1). +const MAX_RUN: usize = 31; +/// Largest level with a table code (run 0: level 40). +const MAX_LEVEL: usize = 40; + +pub(crate) struct EncTables { + /// macroblock_address_increment 1..=33 (index 0 unused) and the escape. + pub mba: [Code; 34], + pub mba_escape: Code, + /// macroblock_type by flags: I pictures, P pictures (index: flags, 0..32). + pub mbtype_i: [Option; 32], + pub mbtype_p: [Option; 32], + /// coded_block_pattern_420 (the first six blocks). + pub cbp: [Code; 64], + /// motion_code -16..=16 at index code + 16. + pub motion: [Code; 33], + pub dc_luma: [Code; 12], + pub dc_chroma: [Code; 12], + /// DCT coefficients `[table][run][level]` without the sign bit (table 0: B.14, table 1: B.15). + pub dct: [Box<[[Option; MAX_LEVEL + 1]; MAX_RUN + 1]>; 2], + pub dct_eob: [Code; 2], + pub dct_escape: Code, +} + +fn by_value(list: &[(&str, i16)], offset: i16) -> [Code; N] { + let mut out = [Code::default(); N]; + for &(s, v) in list { + if let Some(slot) = usize::try_from(v + offset).ok().and_then(|i| out.get_mut(i)) { + *slot = Code::of(s); + } + } + out +} + +fn mb_types(list: &[(&str, i16)]) -> [Option; 32] { + let mut out = [None; 32]; + for &(s, v) in list { + if let Some(slot) = usize::try_from(v).ok().and_then(|i| out.get_mut(i)) { + *slot = Some(Code::of(s)); + } + } + out +} + +fn dct( + short: &'static [(&'static str, i16)], + extra: &'static [(&'static str, i16)], +) -> (Box<[[Option; MAX_LEVEL + 1]; MAX_RUN + 1]>, Code, Code) { + let mut t = Box::new([[None; MAX_LEVEL + 1]; MAX_RUN + 1]); + let (mut eob, mut esc) = (Code::default(), Code::default()); + for (s, v) in dct_table(short, extra) { + match v { + DCT_EOB => eob = Code::of(s), + DCT_ESCAPE => esc = Code::of(s), + v => { + let (run, level) = ((v & 63) as usize, (v >> 6) as usize); + if let Some(slot) = t.get_mut(run).and_then(|r| r.get_mut(level)) { + *slot = Some(Code::of(s)); + } + } + } + } + (t, eob, esc) +} + +pub(crate) fn enc_tables() -> &'static EncTables { + static T: OnceLock = OnceLock::new(); + T.get_or_init(|| { + let (zero, eob0, esc) = dct(DCT_ZERO, DCT_ZERO_LONG); + let (one, eob1, _) = dct(DCT_ONE, &[]); + let mba_escape = MBA.iter().find(|e| e.1 == MBA_ESCAPE).map_or(Code::default(), |e| Code::of(e.0)); + EncTables { + mba: by_value(MBA, 0), + mba_escape, + mbtype_i: mb_types(MBTYPE_I), + mbtype_p: mb_types(MBTYPE_P), + cbp: by_value(CBP, 0), + motion: by_value(MOTION, 16), + dc_luma: by_value(DC_LUMA, 0), + dc_chroma: by_value(DC_CHROMA, 0), + dct: [zero, one], + dct_eob: [eob0, eob1], + dct_escape: esc, + } + }) +} + +impl EncTables { + /// One DCT coefficient. `first_inter`: the first coefficient of a non-intra block, where + /// run 0 level ±1 is "1s". + #[inline] + pub fn put_coefficient(&self, w: &mut BitWriter, table: usize, run: usize, level: i32, first_inter: bool) { + let mag = level.unsigned_abs() as usize; + let neg = level < 0; + if first_inter && run == 0 && mag == 1 { + w.put(0b10 | neg as u32, 2); + return; + } + if let Some(c) = self.dct[table & 1].get(run).and_then(|r| r.get(mag)).copied().flatten() { + c.put(w); + w.bit(neg); + return; + } + // escape: 6-bit run, 12-bit two's complement level (-2047..=2047) + self.dct_escape.put(w); + w.put(run as u32, 6); + w.put((level.clamp(-2047, 2047) & 0xFFF) as u32, 12); + } + + /// Bits of one motion vector component delta (`delta` already in the wrapped range). + #[inline] + pub fn motion_bits(&self, delta: i32, f_code: u8) -> u32 { + let r_size = u32::from(f_code.saturating_sub(1)); + if delta == 0 { + return u32::from(self.motion[16].len); + } + let code = ((delta.abs() - 1) >> r_size) + 1; + let idx = (16 + code.min(16)) as usize; + u32::from(self.motion[idx].len) + r_size + } + + /// Write one motion vector component delta (`delta` in the wrapped range). + #[inline] + pub fn put_motion(&self, w: &mut BitWriter, delta: i32, f_code: u8) { + let r_size = u32::from(f_code.saturating_sub(1)); + if delta == 0 { + self.motion[16].put(w); + return; + } + let a = delta.abs() - 1; + let code = ((a >> r_size) + 1).min(16); + let idx = (16 + code * delta.signum()) as usize; + self.motion[idx].put(w); + w.put((a & ((1 << r_size) - 1)) as u32, r_size); + } +} + +/// `mv - pred` brought into the range a motion code can carry ([-16f, 16f-1], §7.6.3.1). +#[inline] +pub(crate) fn wrap_delta(mv: i32, pred: i32, f_code: u8) -> i32 { + let f = 1i32 << f_code.saturating_sub(1); + let d = mv - pred; + if d > 16 * f - 1 { + d - 32 * f + } else if d < -16 * f { + d + 32 * f + } else { + d + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::mpeg2::bits::Bits; + use crate::mpeg2::slice::wrap_vector; + use crate::mpeg2::vlc::{MB_FWD, MB_INTRA, MB_PATTERN, MB_QUANT, rl, tables}; + + /// Macroblock types the encoder sends (frame pictures, frame prediction, no B pictures). + const ENC_MB_TYPES: [i16; 7] = + [MB_INTRA, MB_INTRA | MB_QUANT, MB_FWD | MB_PATTERN, MB_PATTERN, MB_FWD, MB_QUANT | MB_FWD | MB_PATTERN, MB_QUANT | MB_PATTERN]; + + #[test] + fn every_code_decodes_to_its_value() { + let t = enc_tables(); + let d = tables(); + let check = |c: Code, vlc: &crate::mpeg2::vlc::Vlc, v: i16| { + let mut w = BitWriter::new(); + c.put(&mut w); + let bytes = w.finish(); + let mut b = Bits::new(&bytes); + assert_eq!(vlc.decode(&mut b), Some(v)); + }; + for n in 1..=33 { + check(t.mba[n], &d.mba, n as i16); + } + check(t.mba_escape, &d.mba, MBA_ESCAPE); + for v in 0..64 { + check(t.cbp[v], &d.cbp, v as i16); + } + for v in -16..=16i16 { + check(t.motion[(v + 16) as usize], &d.motion, v); + } + for s in 0..12 { + check(t.dc_luma[s], &d.dc_luma, s as i16); + check(t.dc_chroma[s], &d.dc_chroma, s as i16); + } + for f in ENC_MB_TYPES { + // the I table holds only the intra types + if let Some(c) = t.mbtype_i[f as usize] { + check(c, &d.mbtype[0], f); + } + check(t.mbtype_p[f as usize].unwrap_or_default(), &d.mbtype[1], f); + } + for (k, vlc) in [&d.dct_zero, &d.dct_one].into_iter().enumerate() { + for run in 0..=MAX_RUN { + for level in 1..=MAX_LEVEL { + if let Some(c) = t.dct[k][run][level] { + check(c, vlc, rl(run as i16, level as i16)); + } + } + } + check(t.dct_eob[k], vlc, DCT_EOB); + } + } + + #[test] + fn motion_deltas_round_trip_through_the_decoder_wrap() { + let t = enc_tables(); + let d = tables(); + for f_code in 1..=4u8 { + let f = 1i32 << (f_code - 1); + for pred in [-16 * f, -3, 0, 5, 16 * f - 1] { + for mv in -16 * f..16 * f { + let delta = wrap_delta(mv, pred, f_code); + let mut w = BitWriter::new(); + t.put_motion(&mut w, delta, f_code); + let bits = w.bits() as u32; + assert_eq!(bits, t.motion_bits(delta, f_code)); + let bytes = w.finish(); + let mut b = Bits::new(&bytes); + // the decoder's reading of motion_code and motion_residual + let code = d.motion.decode(&mut b).unwrap_or(99) as i32; + let r_size = (f_code - 1) as u32; + let residual = if f_code != 1 && code != 0 { b.read(r_size) as i32 } else { 0 }; + let dd = if f == 1 || code == 0 { code } else { ((code.abs() - 1) * f + residual + 1) * code.signum() }; + assert_eq!(dd, delta); + assert_eq!(wrap_vector(pred, dd, f), mv, "f_code {f_code} pred {pred} mv {mv}"); + } + } + } + } +} diff --git a/crates/video/src/mpeg2/encode/tests.rs b/crates/video/src/mpeg2/encode/tests.rs new file mode 100644 index 000000000..5352161b3 --- /dev/null +++ b/crates/video/src/mpeg2/encode/tests.rs @@ -0,0 +1,250 @@ +//! Encoder tests against this crate's decoder: the decoded pictures must equal the encoder's +//! reconstruction exactly (any difference would drift along a GOP of P pictures). The ffmpeg +//! comparison is in `tests/mpeg2/encode.rs`. + +use super::*; +use crate::mpeg2::{Decoder, access_units, probe}; + +/// A moving, textured test sequence: a gradient with a sharp pattern drifting by a fraction of a +/// sample per picture, and a block moving across it. +pub(crate) fn scene(w: usize, h: usize, cw: usize, ch: usize, t: usize) -> (Vec, Vec, Vec) { + let tf = t as f32; + let mut y = vec![0u8; w * h]; + for r in 0..h { + for c in 0..w { + let (x, yy) = (c as f32 + tf * 1.5, r as f32 + tf * 0.75); + let v = 60.0 + + 0.3 * x + + 0.2 * yy + + 40.0 * ((x * 0.21).sin() * (yy * 0.17).cos()) + + if ((x as i32 / 8) + (yy as i32 / 8)) % 2 == 0 { 25.0 } else { 0.0 }; + y[r * w + c] = v.clamp(0.0, 255.0) as u8; + } + } + let (bx, by) = ((t * 3) % w.max(1), (t * 2) % h.max(1)); + for r in by..(by + 12).min(h) { + for c in bx..(bx + 12).min(w) { + y[r * w + c] = 230; + } + } + let mut cb = vec![0u8; cw * ch]; + let mut cr = vec![0u8; cw * ch]; + for r in 0..ch { + for c in 0..cw { + cb[r * cw + c] = (128.0 + 50.0 * ((c as f32 + tf) * 0.13).sin()) as u8; + cr[r * cw + c] = (128.0 + 40.0 * ((r as f32 - tf) * 0.11).cos()) as u8; + } + } + (y, cb, cr) +} + +/// Encode `frames` pictures of the scene; returns the stream and the reconstruction. +fn encode(cfg: EncoderConfig, frames: usize) -> (Vec, Vec, Vec) { + let (w, h) = (cfg.width as usize, cfg.height as usize); + let (cw, ch) = (w / 2, if cfg.chroma == ChromaFormat::Yuv422 { h } else { h / 2 }); + let mut enc = Encoder::new(cfg).unwrap(); + enc.set_recon_capture(true); + let mut es = Vec::new(); + let mut packets = Vec::new(); + for t in 0..frames { + let (y, cb, cr) = scene(w, h, cw, ch, t); + let f = YuvFrame { y: &y, cb: &cb, cr: &cr, y_stride: w, c_stride: cw }; + for p in enc.encode(&f, t as i64).unwrap() { + es.extend_from_slice(&p.data); + packets.push(p); + } + } + packets.extend(enc.flush()); + es.extend_from_slice(enc.end_of_sequence()); + (es, enc.take_recon(), packets) +} + +fn decode(es: &[u8]) -> Vec { + let mut d = Decoder::new(); + let mut out = Vec::new(); + for (i, r) in access_units(es).into_iter().enumerate() { + out.extend(d.decode(&es[r], i as i64).unwrap()); + } + out.extend(d.flush()); + assert_eq!(d.errors(), 0, "decoder errors: {:?}", d.last_error()); + out +} + +fn assert_decodes_to_recon(name: &str, cfg: EncoderConfig, frames: usize) -> Vec { + let (es, recon, packets) = encode(cfg, frames); + let pics = decode(&es); + assert_eq!(pics.len(), frames, "{name}: picture count"); + for (i, (p, r)) in pics.iter().zip(&recon).enumerate() { + assert_eq!(p.pts, i as i64, "{name}: display order"); + for (plane, (a, b)) in [("Y", (&p.y, &r.y)), ("Cb", (&p.cb, &r.cb)), ("Cr", (&p.cr, &r.cr))] { + let first = a.iter().zip(b.iter()).position(|(x, y)| x != y); + assert!(first.is_none() && a.len() == b.len(), "{name}: picture {i} ({:?}) {plane} differs at {first:?}", p.picture_type); + } + } + packets +} + +fn psnr(a: &[u8], b: &[u8]) -> f64 { + let sse: f64 = a.iter().zip(b).map(|(x, y)| (f64::from(*x) - f64::from(*y)).powi(2)).sum(); + 10.0 * (255.0f64.powi(2) * a.len() as f64 / sse.max(1e-9)).log10() +} + +#[test] +fn intra_only_decodes_to_the_reconstruction() { + let mut cfg = EncoderConfig::new(64, 48, 25, 1); + cfg.gop_size = 1; + let packets = assert_decodes_to_recon("intra", cfg, 3); + assert!(packets.iter().all(|p| p.keyframe && p.picture_type == PictureType::I)); +} + +#[test] +fn long_p_gop_does_not_drift() { + let mut cfg = EncoderConfig::new(96, 64, 25, 1); + cfg.gop_size = 30; + let packets = assert_decodes_to_recon("long gop", cfg, 30); + assert_eq!(packets.iter().filter(|p| p.keyframe).count(), 1); + assert!(packets[1..].iter().all(|p| p.picture_type == PictureType::P)); +} + +#[test] +fn every_coding_option_decodes_to_the_reconstruction() { + let base = EncoderConfig { gop_size: 6, ..EncoderConfig::new(80, 48, 30000, 1001) }; + let mut matrix = [0u8; 64]; + for (i, m) in matrix.iter_mut().enumerate() { + *m = (8 + i * 2) as u8; + } + let cases = [ + ("422", EncoderConfig { chroma: ChromaFormat::Yuv422, ..base.clone() }), + ("linear q", EncoderConfig { nonlinear_quant: false, rate: RateControl::Quantiser(3), ..base.clone() }), + ("coarse", EncoderConfig { rate: RateControl::Quantiser(31), ..base.clone() }), + ("fine", EncoderConfig { rate: RateControl::Quantiser(1), adaptive_quant: false, ..base.clone() }), + ("dc 8 and 10", EncoderConfig { intra_dc_bits: 10, ..base.clone() }), + ("dc 8", EncoderConfig { intra_dc_bits: 8, ..base.clone() }), + ("matrices", EncoderConfig { intra_matrix: Some(matrix), non_intra_matrix: Some(matrix), ..base.clone() }), + ("small range", EncoderConfig { search_range: 3, ..base.clone() }), + ("abr", EncoderConfig { rate: RateControl::Abr { kbps: 300 }, ..base.clone() }), + ("odd size", EncoderConfig { width: 50, height: 34, ..base.clone() }), + ("422 odd size", EncoderConfig { width: 70, height: 30, chroma: ChromaFormat::Yuv422, ..base.clone() }), + ]; + for (name, cfg) in cases { + assert_decodes_to_recon(name, cfg, 8); + } +} + +#[cfg(feature = "threads")] +#[test] +fn threaded_slices_give_the_same_stream() { + let cfg = EncoderConfig { gop_size: 5, threads: false, ..EncoderConfig::new(96, 80, 25, 1) }; + let (a, _, _) = encode(cfg.clone(), 6); + let (b, _, _) = encode(EncoderConfig { threads: true, ..cfg }, 6); + assert_eq!(a, b); +} + +#[test] +fn quality_follows_the_quantiser() { + let mut last = f64::MAX; + let mut last_size = usize::MAX; + for q in [1u8, 4, 12, 31] { + let cfg = EncoderConfig { rate: RateControl::Quantiser(q), gop_size: 4, ..EncoderConfig::new(96, 64, 25, 1) }; + let (es, recon, _) = encode(cfg, 4); + let p: f64 = recon.iter().enumerate().map(|(t, r)| psnr(&scene(96, 64, 48, 32, t).0, &r.y)).sum::() / recon.len() as f64; + assert!(p < last && es.len() < last_size, "q {q}: {p:.2} dB, {} bytes", es.len()); + if q == 1 { + assert!(p > 45.0, "q 1: {p:.2} dB"); + } + if q == 4 { + assert!(p > 36.0, "q 4: {p:.2} dB"); + } + (last, last_size) = (p, es.len()); + } +} + +#[test] +fn p_pictures_are_smaller_and_static_content_skips() { + let cfg = EncoderConfig { gop_size: 8, ..EncoderConfig::new(128, 64, 25, 1) }; + let mut enc = Encoder::new(cfg).unwrap(); + let (y, cb, cr) = scene(128, 64, 64, 32, 0); + let f = YuvFrame { y: &y, cb: &cb, cr: &cr, y_stride: 128, c_stride: 64 }; + let i = enc.encode(&f, 0).unwrap().remove(0); + let p1 = enc.encode(&f, 1).unwrap().remove(0); + let p2 = enc.encode(&f, 2).unwrap().remove(0); + assert_eq!(p1.picture_type, PictureType::P); + // the first P picture refines the I picture's reconstruction; the next repeat has almost + // nothing left to code and skips (every macroblock but the first and last of a slice) + assert!(p1.data.len() * 5 < i.data.len(), "I {} bytes, P {} bytes", i.data.len(), p1.data.len()); + assert!(p2.data.len() * 20 < i.data.len(), "I {} bytes, second P {} bytes", i.data.len(), p2.data.len()); +} + +#[test] +fn the_sequence_header_describes_the_stream() { + let cases = [ + (EncoderConfig::new(352, 288, 25, 1), 0x4A, "Main@Low"), + (EncoderConfig::new(720, 576, 25, 1), 0x48, "Main@Main"), + (EncoderConfig::new(1440, 1080, 25, 1), 0x46, "Main@High-1440"), + (EncoderConfig::new(1920, 1080, 30000, 1001), 0x44, "Main@High"), + (EncoderConfig { chroma: ChromaFormat::Yuv422, ..EncoderConfig::new(720, 608, 25, 1) }, 0x85, "4:2:2@Main"), + (EncoderConfig { chroma: ChromaFormat::Yuv422, ..EncoderConfig::new(1920, 1080, 25, 1) }, 0x82, "4:2:2@High"), + (EncoderConfig { rate: RateControl::Abr { kbps: 20_000 }, ..EncoderConfig::new(720, 576, 25, 1) }, 0x46, "Main@High-1440"), + ]; + for (cfg, pli, name) in cases { + let (w, h, aspect) = (cfg.width, cfg.height, cfg.aspect); + let enc = Encoder::new(cfg).unwrap(); + assert_eq!(enc.profile_and_level(), pli, "{name}"); + let info = probe(enc.sequence_header()).unwrap(); + assert_eq!((info.width, info.height, info.mpeg2, info.progressive_sequence, info.low_delay), (w, h, true, true, true)); + assert_eq!(info.codec_name(), format!("MPEG-2 Video ({name})")); + assert_eq!(aspect, AspectRatio::Square); + assert_eq!(info.sar, (1, 1)); + } + let enc = Encoder::new(EncoderConfig { aspect: AspectRatio::Dar16x9, ..EncoderConfig::new(720, 576, 30000, 1001) }).unwrap(); + let info = probe(enc.sequence_header()).unwrap(); + assert_eq!((info.dar, info.frame_rate), (Some((16, 9)), Some((30000, 1001)))); +} + +#[test] +fn bad_configurations_are_refused() { + let ok = EncoderConfig::new(64, 64, 25, 1); + let bad = [ + EncoderConfig { width: 63, ..ok.clone() }, + EncoderConfig { height: 0, ..ok.clone() }, + EncoderConfig { width: 3840, height: 2160, ..ok.clone() }, + EncoderConfig { chroma: ChromaFormat::Yuv444, ..ok.clone() }, + EncoderConfig { chroma: ChromaFormat::Yuv422, width: 1920, height: 1152, ..ok.clone() }, + EncoderConfig { fps_num: 26, ..ok.clone() }, + EncoderConfig { fps_den: 0, ..ok.clone() }, + EncoderConfig { gop_size: 0, ..ok.clone() }, + EncoderConfig { gop_size: 1024, ..ok.clone() }, + EncoderConfig { rate: RateControl::Quantiser(0), ..ok.clone() }, + EncoderConfig { rate: RateControl::Quantiser(32), ..ok.clone() }, + EncoderConfig { rate: RateControl::Abr { kbps: 0 }, ..ok.clone() }, + EncoderConfig { rate: RateControl::Abr { kbps: 90_000 }, ..ok.clone() }, + EncoderConfig { intra_dc_bits: 11, ..ok.clone() }, + EncoderConfig { intra_matrix: Some([0; 64]), ..ok.clone() }, + EncoderConfig { search_range: 0, ..ok.clone() }, + EncoderConfig { search_range: 64, ..ok.clone() }, + ]; + for cfg in bad { + let desc = format!("{cfg:?}"); + assert!(matches!(Encoder::new(cfg), Err(EncodeError::Config(_))), "{desc}"); + } +} + +#[test] +fn short_or_inconsistent_planes_are_refused() { + let mut enc = Encoder::new(EncoderConfig::new(32, 32, 25, 1)).unwrap(); + let (y, c) = (vec![0u8; 32 * 32], vec![0u8; 16 * 16]); + let ok = YuvFrame { y: &y, cb: &c, cr: &c, y_stride: 32, c_stride: 16 }; + assert!(enc.encode(&ok, 0).is_ok()); + let bad = [ + YuvFrame { y: &y[..32 * 31], ..ok }, + YuvFrame { cb: &c[..10], ..ok }, + YuvFrame { y_stride: 31, ..ok }, + YuvFrame { c_stride: usize::MAX, ..ok }, + YuvFrame { y_stride: usize::MAX / 2, ..ok }, + ]; + for f in bad { + assert!(matches!(enc.encode(&f, 1), Err(EncodeError::BadFrame(_)))); + } + // a refused frame leaves the encoder usable + assert_eq!(enc.encode(&ok, 1).unwrap()[0].picture_type, PictureType::P); +} diff --git a/crates/video/src/mpeg2/encode/transform.rs b/crates/video/src/mpeg2/encode/transform.rs new file mode 100644 index 000000000..d3a6ce4d3 --- /dev/null +++ b/crates/video/src/mpeg2/encode/transform.rs @@ -0,0 +1,122 @@ +//! Forward DCT, quantisation and the decoder-exact reconstruction of one block. + +use crate::mpeg2::idct::{basis, idct}; +use crate::mpeg2::slice::dequant_mpeg2; + +/// Forward 8×8 DCT of `p` (raster order), the transpose of the decoder's IDCT basis: +/// `F[v][u] = Σ_y Σ_x M[x][u] M[y][v] p[y][x]`, so a flat block of value `a` has `F[0][0] = 8a`. +pub(crate) fn fdct(p: &[i32; 64]) -> [f32; 64] { + let m = basis(); + let mut t = [0f32; 64]; + for y in 0..8 { + let row: [f32; 8] = std::array::from_fn(|x| p[y * 8 + x] as f32); + for u in 0..8 { + let mut s = 0f32; + for x in 0..8 { + s += m[x][u] * row[x]; + } + t[y * 8 + u] = s; + } + } + let mut out = [0f32; 64]; + for v in 0..8 { + for u in 0..8 { + let mut s = 0f32; + for y in 0..8 { + s += m[y][v] * t[y * 8 + u]; + } + out[v * 8 + u] = s; + } + } + out +} + +/// Rounding offset of intra AC quantisation (in units of the quantiser step): below one half, +/// a small dead zone. +const INTRA_ROUND: f32 = 0.375; + +/// Quantise intra AC coefficient `f` with weight `w` and quantiser_scale `qs` (inverse of +/// `F = 2·QF·w·qs / 32`). +#[inline] +pub(crate) fn quant_intra(f: f32, w: u8, qs: i32) -> i32 { + let y = 16.0 * f.abs() / (f32::from(w) * qs as f32); + let q = ((y + INTRA_ROUND).floor() as i32).min(2047); + if f < 0.0 { -q } else { q } +} + +/// Quantise non-intra coefficient `f` (inverse of `F = (2·QF + sign)·w·qs / 32`): truncation, +/// which leaves the usual wide dead zone around zero. +#[inline] +pub(crate) fn quant_inter(f: f32, w: u8, qs: i32) -> i32 { + let y = 16.0 * f.abs() / (f32::from(w) * qs as f32); + let q = (y.floor() as i32).min(2047); + if f < 0.0 { -q } else { q } +} + +/// Reconstruct a block as the decoder does: dequantise `levels` (raster order; for intra blocks +/// `levels[0]` is the quantised DC), mismatch control, IDCT. Returns the residual (inter) or +/// the samples before clamping (intra). +pub(crate) fn reconstruct(levels: &[i32; 64], matrix: &[u8; 64], qs: i32, intra: bool, dc_mult: i32) -> [i32; 64] { + let mut blk = [0i32; 64]; + let mut sum = 0i32; + for (i, &l) in levels.iter().enumerate() { + if l == 0 { + continue; + } + let v = if intra && i == 0 { l * dc_mult } else { dequant_mpeg2(l, qs, i32::from(matrix[i]), intra) }; + blk[i] = v; + sum += v; + } + if sum & 1 == 0 { + blk[63] ^= 1; + } + let dc_only = blk[1..].iter().all(|&v| v == 0); + idct(&mut blk, dc_only); + blk +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn flat_block_has_only_dc() { + let f = fdct(&[100; 64]); + assert!((f[0] - 800.0).abs() < 0.01); + assert!(f[1..].iter().all(|v| v.abs() < 0.01)); + } + + #[test] + fn forward_then_inverse_is_close() { + let p: [i32; 64] = std::array::from_fn(|i| ((i * 37 + i / 8 * 11) % 200) as i32 - 60); + let f = fdct(&p); + let mut b: [i32; 64] = std::array::from_fn(|i| f[i].round() as i32); + idct(&mut b, false); + for i in 0..64 { + assert!((b[i] - p[i]).abs() <= 1, "{i}: {} vs {}", b[i], p[i]); + } + } + + #[test] + fn quantisers_invert_the_dequantiser() { + for qs in [1, 2, 8, 31, 62, 112] { + for w in [8u8, 16, 40, 83] { + // only where the dequantiser's integer division is exact + if (i32::from(w) * qs) % 32 != 0 { + continue; + } + for level in [-40, -3, -1, 1, 2, 17] { + let f = dequant_mpeg2(level, qs, i32::from(w), true) as f32; + // without the clamp the intra quantiser returns the level it came from + if f.abs() < 2047.0 { + assert_eq!(quant_intra(f, w, qs), level, "intra qs {qs} w {w}"); + } + let f = dequant_mpeg2(level, qs, i32::from(w), false) as f32; + if f.abs() < 2047.0 { + assert_eq!(quant_inter(f, w, qs), level, "inter qs {qs} w {w}"); + } + } + } + } + } +} diff --git a/crates/video/src/mpeg2/idct.rs b/crates/video/src/mpeg2/idct.rs index f94b935f0..3828d05a6 100644 --- a/crates/video/src/mpeg2/idct.rs +++ b/crates/video/src/mpeg2/idct.rs @@ -6,7 +6,7 @@ use std::sync::OnceLock; /// `M[x][u] = C(u)/2 · cos((2x+1)uπ/16)`, `C(0) = 1/√2`, `C(u>0) = 1`. -fn basis() -> &'static [[f32; 8]; 8] { +pub(crate) fn basis() -> &'static [[f32; 8]; 8] { static B: OnceLock<[[f32; 8]; 8]> = OnceLock::new(); B.get_or_init(|| { let mut m = [[0f32; 8]; 8]; diff --git a/crates/video/src/mpeg2/mod.rs b/crates/video/src/mpeg2/mod.rs index 5d88ea4e5..ae2bb1682 100644 --- a/crates/video/src/mpeg2/mod.rs +++ b/crates/video/src/mpeg2/mod.rs @@ -1,5 +1,6 @@ //! Clean-room MPEG-2 video (ITU-T H.262 | ISO/IEC 13818-2) and MPEG-1 video (ISO/IEC 11172-2) -//! decoder. +//! decoder, and the start of an MPEG-2 encoder ([`encode`], feature `mpeg2-encode`: I and P +//! pictures). //! //! - Main Profile at Low/Main/High-1440/High Level, 4:2:2 Profile (XDCAM HD422, IMX / D-10), //! Simple Profile, and MPEG-1 (including D pictures and constrained-parameters streams). @@ -26,12 +27,17 @@ #![allow(clippy::needless_range_loop, clippy::type_complexity)] mod bits; +#[cfg(feature = "mpeg2-encode")] +pub mod encode; pub mod headers; mod idct; mod mc; mod slice; mod vlc; +#[cfg(feature = "mpeg2-encode")] +pub use encode::{AspectRatio, EncodeError, Encoder, EncoderConfig, Packet, RateControl, ReconFrame, YuvFrame}; + use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; diff --git a/crates/video/src/mpeg2/slice.rs b/crates/video/src/mpeg2/slice.rs index f7a6389f0..706439af4 100644 --- a/crates/video/src/mpeg2/slice.rs +++ b/crates/video/src/mpeg2/slice.rs @@ -519,8 +519,7 @@ pub(crate) fn wrap_vector(pred: i32, delta: i32, f: i32) -> i32 { #[inline(always)] fn dequant(p: &PicParams, level: i32, qs: i32, w: i32, intra: bool) -> i32 { if p.mpeg2 { - let k = if intra { 0 } else { level.signum() }; - ((2 * level + k) * w * qs / 32).clamp(-2048, 2047) + dequant_mpeg2(level, qs, w, intra) } else { let mut v = if intra { 2 * level * qs * w / 16 } else { (2 * level + level.signum()) * qs * w / 16 }; if v & 1 == 0 { @@ -530,6 +529,14 @@ fn dequant(p: &PicParams, level: i32, qs: i32, w: i32, intra: bool) -> i32 { } } +/// MPEG-2 inverse quantisation of one AC (or non-intra) coefficient (§7.4.2.3), saturated +/// (§7.4.3). Shared with the encoder, whose reconstruction must match the decoder exactly. +#[inline(always)] +pub(crate) fn dequant_mpeg2(level: i32, qs: i32, w: i32, intra: bool) -> i32 { + let k = if intra { 0 } else { level.signum() }; + ((2 * level + k) * w * qs / 32).clamp(-2048, 2047) +} + /// The reference field (or frame) for direction `s` and field `sel` (field prediction). fn reference<'a>(p: &PicParams<'a>, s: usize, sel: usize) -> Option<&'a Frame> { if s == 1 { diff --git a/crates/video/src/mpeg2/vlc.rs b/crates/video/src/mpeg2/vlc.rs index 192f86be0..3f8015bb5 100644 --- a/crates/video/src/mpeg2/vlc.rs +++ b/crates/video/src/mpeg2/vlc.rs @@ -22,7 +22,7 @@ pub(crate) struct Vlc { table: Vec, } -fn parse_code(s: &str) -> (u32, u32) { +pub(crate) fn parse_code(s: &str) -> (u32, u32) { let mut v = 0u32; let mut n = 0u32; for c in s.chars() { diff --git a/crates/video/tests/mpeg2/encode.rs b/crates/video/tests/mpeg2/encode.rs new file mode 100644 index 000000000..5cb2377dc --- /dev/null +++ b/crates/video/tests/mpeg2/encode.rs @@ -0,0 +1,206 @@ +//! The MPEG-2 encoder against ffmpeg as an external decoder: ffmpeg must decode our streams +//! without errors, to within the decoder suite's tolerance of the encoder's own reconstruction +//! (ffmpeg's integer IDCT differs from ours), and report the profile, level and picture types +//! we signal. Source pictures come from ffmpeg's `testsrc2`. + +use std::path::{Path, PathBuf}; +use std::process::{Command, Stdio}; + +use craft_video::mpeg2::{ChromaFormat, Encoder, EncoderConfig, RateControl, ReconFrame, YuvFrame}; + +use crate::common::{Stats, fixture_dir, pix_fmt_of}; +use crate::oracle::{MAX_DIFFERING, MIN_PSNR, TOL_MAX}; + +/// Raw source frames from ffmpeg's `testsrc2`. +fn source(ff: &Path, w: u32, h: u32, frames: usize, pix_fmt: &str) -> Vec { + let o = Command::new(ff) + .args([ + "-hide_banner", + "-loglevel", + "error", + "-f", + "lavfi", + "-i", + &format!("testsrc2=size={w}x{h}:rate=25"), + "-frames:v", + &frames.to_string(), + ]) + .args(["-pix_fmt", pix_fmt, "-f", "rawvideo", "-"]) + .stdin(Stdio::null()) + .output() + .expect("run ffmpeg"); + assert!(o.status.success(), "ffmpeg: {}", String::from_utf8_lossy(&o.stderr)); + o.stdout +} + +fn sizes(cfg: &EncoderConfig) -> (usize, usize, usize) { + let (w, h) = (cfg.width as usize, cfg.height as usize); + let ch = if cfg.chroma == ChromaFormat::Yuv422 { h } else { h / 2 }; + (w * h, w / 2 * ch, w / 2) +} + +/// Encode `raw` frames; returns the elementary stream and the reconstruction. +fn encode(cfg: &EncoderConfig, raw: &[u8]) -> (Vec, Vec) { + let (ys, cs, cw) = sizes(cfg); + let fsize = ys + 2 * cs; + let mut enc = Encoder::new(cfg.clone()).expect("config"); + enc.set_recon_capture(true); + let mut es = Vec::new(); + for (i, f) in raw.chunks_exact(fsize).enumerate() { + let frame = YuvFrame { y: &f[..ys], cb: &f[ys..ys + cs], cr: &f[ys + cs..], y_stride: cfg.width as usize, c_stride: cw }; + for p in enc.encode(&frame, i as i64).expect("encode") { + es.extend_from_slice(&p.data); + } + } + for p in enc.flush() { + es.extend_from_slice(&p.data); + } + es.extend_from_slice(enc.end_of_sequence()); + (es, enc.take_recon()) +} + +/// ffmpeg's decode of `file`, failing on any decoding error. +fn ffmpeg_decode(ff: &Path, file: &Path, pix_fmt: &str) -> Vec { + let o = Command::new(ff) + .args(["-hide_banner", "-loglevel", "error", "-xerror", "-err_detect", "explode", "-i"]) + .arg(file) + .args(["-fps_mode", "passthrough", "-f", "rawvideo", "-pix_fmt", pix_fmt, "-"]) + .stdin(Stdio::null()) + .output() + .expect("run ffmpeg"); + assert!(o.status.success() && o.stderr.is_empty(), "ffmpeg rejects {}: {}", file.display(), String::from_utf8_lossy(&o.stderr)); + o.stdout +} + +/// ffprobe's stream description and per-frame picture types. +fn ffprobe(file: &Path) -> Option<(String, Vec)> { + let fp = craft_media_testkit::ffprobe()?; + let stream = Command::new(&fp) + .args(["-v", "error", "-select_streams", "v:0", "-show_entries", "stream=codec_name,profile,level,pix_fmt", "-of", "default=nw=1"]) + .arg(file) + .output() + .ok()?; + let frames = Command::new(&fp) + .args(["-v", "error", "-select_streams", "v:0", "-show_entries", "frame=pict_type", "-of", "csv=p=0"]) + .arg(file) + .output() + .ok()?; + let types = String::from_utf8_lossy(&frames.stdout).lines().map(|l| l.trim_end_matches(',').to_string()).filter(|l| !l.is_empty()).collect(); + // "codec_name=…\nprofile=…" in ffprobe's order → "codec,profile,level,pix_fmt" + let text = String::from_utf8_lossy(&stream.stdout).to_string(); + let field = |k: &str| text.lines().find_map(|l| l.strip_prefix(&format!("{k}="))).unwrap_or("").to_string(); + Some((["codec_name", "profile", "level", "pix_fmt"].map(field).join(","), types)) +} + +fn write(name: &str, es: &[u8]) -> PathBuf { + let path = fixture_dir().join(format!("enc_{name}.m2v")); + std::fs::write(&path, es).expect("write stream"); + path +} + +fn psnr_y(recon: &[ReconFrame], raw: &[u8], ys: usize, fsize: usize) -> f64 { + let mut s = Stats::default(); + for (r, f) in recon.iter().zip(raw.chunks_exact(fsize)) { + s.add(&r.y, &f[..ys]); + } + s.psnr() +} + +struct Case { + name: &'static str, + cfg: EncoderConfig, + frames: usize, + /// ffprobe's "codec,profile,level,pix_fmt". + probe: &'static str, + /// Lowest luma PSNR against the source. + min_psnr: f64, +} + +fn check(ff: &Path, c: &Case) -> (Vec, f64) { + let pix_fmt = pix_fmt_of(c.cfg.chroma); + let raw = source(ff, c.cfg.width, c.cfg.height, c.frames, pix_fmt); + let (ys, cs, _) = sizes(&c.cfg); + let fsize = ys + 2 * cs; + let (es, recon) = encode(&c.cfg, &raw); + let path = write(c.name, &es); + let decoded = ffmpeg_decode(ff, &path, pix_fmt); + assert_eq!(decoded.len(), c.frames * fsize, "{}: ffmpeg's frame count", c.name); + let mut all = Stats::default(); + for (i, (r, f)) in recon.iter().zip(decoded.chunks_exact(fsize)).enumerate() { + let mut s = Stats::default(); + s.add(&r.y, &f[..ys]); + s.add(&r.cb, &f[ys..ys + cs]); + s.add(&r.cr, &f[ys + cs..]); + assert!(s.max <= TOL_MAX && s.psnr() >= MIN_PSNR, "{} frame {i}: max {} psnr {:.2} against ffmpeg", c.name, s.max, s.psnr()); + all.merge(&s); + } + assert!((all.diff as f64) / (all.n as f64) <= MAX_DIFFERING, "{}: {} of {} samples differ from ffmpeg", c.name, all.diff, all.n); + let quality = psnr_y(&recon, &raw, ys, fsize); + assert!(quality >= c.min_psnr, "{}: luma PSNR {quality:.2} dB against the source", c.name); + if let Some((stream, types)) = ffprobe(&path) { + assert_eq!(stream, c.probe, "{}: ffprobe", c.name); + let gop = c.cfg.gop_size as usize; + let want: Vec<&str> = (0..c.frames).map(|i| if i % gop == 0 { "I" } else { "P" }).collect(); + assert_eq!(types, want, "{}: picture types", c.name); + } + let kbps = es.len() as f64 * 8.0 * 25.0 / c.frames as f64 / 1000.0; + println!( + "{}: {} frames, {kbps:.0} kbit/s, luma PSNR {quality:.2} dB, max diff to ffmpeg {}, {:.3}% differ", + c.name, + c.frames, + all.max, + 100.0 * all.diff as f64 / all.n as f64 + ); + (es, kbps) +} + +#[test] +fn ffmpeg_decodes_our_streams() { + let Some(ff) = craft_media_testkit::ffmpeg_or_skip("mpeg2 encoder oracle") else { return }; + let cif = EncoderConfig { gop_size: 12, ..EncoderConfig::new(352, 288, 25, 1) }; + let cases = [ + Case { name: "cif_q4", cfg: cif.clone(), frames: 25, probe: "mpeg2video,Main,10,yuv420p", min_psnr: 36.0 }, + Case { + name: "cif_intra", + cfg: EncoderConfig { gop_size: 1, rate: RateControl::Quantiser(8), ..cif.clone() }, + frames: 6, + probe: "mpeg2video,Main,10,yuv420p", + min_psnr: 32.0, + }, + Case { + name: "cif_linear_coarse", + cfg: EncoderConfig { nonlinear_quant: false, adaptive_quant: false, rate: RateControl::Quantiser(20), gop_size: 15, ..cif.clone() }, + frames: 20, + probe: "mpeg2video,Main,10,yuv420p", + min_psnr: 26.0, + }, + Case { + name: "sd_422", + cfg: EncoderConfig { chroma: ChromaFormat::Yuv422, intra_dc_bits: 10, ..EncoderConfig::new(720, 576, 25, 1) }, + frames: 8, + // 4:2:2@Main is indication 0x85; ffprobe prints its low four bits as the level + probe: "mpeg2video,4:2:2,5,yuv422p", + min_psnr: 36.0, + }, + ]; + for c in &cases { + check(&ff, c); + } +} + +#[test] +fn average_bitrate_reaches_its_target() { + let Some(ff) = craft_media_testkit::ffmpeg_or_skip("mpeg2 encoder bitrate") else { return }; + for kbps in [800u32, 2500] { + let c = Case { + name: if kbps == 800 { "cif_abr_800" } else { "cif_abr_2500" }, + cfg: EncoderConfig { rate: RateControl::Abr { kbps }, gop_size: 12, ..EncoderConfig::new(352, 288, 25, 1) }, + frames: 72, + probe: "mpeg2video,Main,10,yuv420p", + min_psnr: 26.0, + }; + let (_, got) = check(&ff, &c); + let err = got / f64::from(kbps) - 1.0; + assert!(err.abs() < 0.10, "{kbps} kbit/s target: got {got:.0} kbit/s ({:+.1}%)", 100.0 * err); + } +} diff --git a/crates/video/tests/mpeg2/main.rs b/crates/video/tests/mpeg2/main.rs index 1864b4a22..acbbf662e 100644 --- a/crates/video/tests/mpeg2/main.rs +++ b/crates/video/tests/mpeg2/main.rs @@ -1,4 +1,6 @@ mod common; +#[cfg(feature = "mpeg2-encode")] +mod encode; mod fuzz_hostile; mod hostile_dims; mod oracle;