mirror of
https://github.com/storytold/craft-libs.git
synced 2026-11-13 22:12:12 +00:00
Merge pull request #1136 from alehel/craft-video/mpeg2-encode-ip
craft-video: MPEG-2 encoder, first step: I and P pictures (#1133)
This commit is contained in:
18 files changed
+2361
-13
No files matched your search
@@ -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"
|
||||
|
||||
@@ -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) |
|
||||
|
||||
@@ -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")]
|
||||
|
||||
@@ -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.
|
||||
@@ -0,0 +1,92 @@
|
||||
//! Big-endian bit writer for start-code units.
|
||||
|
||||
pub(crate) struct BitWriter {
|
||||
out: Vec<u8>,
|
||||
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<u8> {
|
||||
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]);
|
||||
}
|
||||
}
|
||||
@@ -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<u8>,
|
||||
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<u8>,
|
||||
pub cb: Vec<u8>,
|
||||
pub cr: Vec<u8>,
|
||||
}
|
||||
|
||||
/// 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<u32> {
|
||||
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<Level, EncodeError> {
|
||||
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<u8>,
|
||||
/// The reconstruction of the previous picture.
|
||||
reference: Option<Frame>,
|
||||
prev_mvs: Vec<[i32; 2]>,
|
||||
/// Pictures coded so far.
|
||||
count: u64,
|
||||
/// Position of the next picture in its GOP.
|
||||
gop_pos: u32,
|
||||
rc: Option<RateCtl>,
|
||||
recon_capture: bool,
|
||||
recon: Vec<ReconFrame>,
|
||||
}
|
||||
|
||||
impl Encoder {
|
||||
pub fn new(cfg: EncoderConfig) -> Result<Self, EncodeError> {
|
||||
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<ReconFrame> {
|
||||
std::mem::take(&mut self.recon)
|
||||
}
|
||||
|
||||
/// Pictures held for reordering (none yet: there are no B pictures).
|
||||
pub fn flush(&mut self) -> Vec<Packet> {
|
||||
Vec::new()
|
||||
}
|
||||
|
||||
fn sequence_header_bytes(&self) -> Vec<u8> {
|
||||
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<Frame, EncodeError> {
|
||||
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<u8> {
|
||||
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<f64> = (0..n).map(|a| activity(src, a % self.mb_w, a / self.mb_w)).collect();
|
||||
let avg = acts.iter().sum::<f64>() / 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<Vec<Packet>, 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<RowOut> {
|
||||
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()
|
||||
}
|
||||
@@ -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}");
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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<u8>,
|
||||
/// 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::<u32>() + 32) / 64;
|
||||
cost += (0..64).map(|k| at(k).abs_diff(mean)).sum::<u32>();
|
||||
}
|
||||
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
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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<Code>; 32],
|
||||
pub mbtype_p: [Option<Code>; 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<Code>; MAX_LEVEL + 1]; MAX_RUN + 1]>; 2],
|
||||
pub dct_eob: [Code; 2],
|
||||
pub dct_escape: Code,
|
||||
}
|
||||
|
||||
fn by_value<const N: usize>(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<Code>; 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<Code>; 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<EncTables> = 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}");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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<u8>, Vec<u8>, Vec<u8>) {
|
||||
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<u8>, Vec<ReconFrame>, Vec<Packet>) {
|
||||
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<crate::mpeg2::Picture> {
|
||||
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<Packet> {
|
||||
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::<f64>() / 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);
|
||||
}
|
||||
@@ -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}");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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];
|
||||
|
||||
@@ -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};
|
||||
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -22,7 +22,7 @@ pub(crate) struct Vlc {
|
||||
table: Vec<Entry>,
|
||||
}
|
||||
|
||||
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() {
|
||||
|
||||
@@ -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<u8> {
|
||||
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<u8>, Vec<ReconFrame>) {
|
||||
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<u8> {
|
||||
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<String>)> {
|
||||
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<u8>, 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);
|
||||
}
|
||||
}
|
||||
@@ -1,4 +1,6 @@
|
||||
mod common;
|
||||
#[cfg(feature = "mpeg2-encode")]
|
||||
mod encode;
|
||||
mod fuzz_hostile;
|
||||
mod hostile_dims;
|
||||
mod oracle;
|
||||
|
||||
Reference in new issue
Block a user