pub mod bitacct;
mod cabac;
mod config;
mod lookahead;
mod mb16;
mod mbtree;
pub fn mbtree_satd_calls() -> u64 {
mbtree::SATD_CALLS.load(std::sync::atomic::Ordering::Relaxed)
}
pub fn mbtree_satd_reset() {
mbtree::SATD_CALLS.store(0, std::sync::atomic::Ordering::Relaxed)
}
mod mvd_cost_tab;
mod params;
mod rc;
mod slice;
pub use crate::mb16::{EXT_MV, ME_PROBE, MVCMP, MVCMP_FRAME};
pub use config::{EncoderConfig, LookaheadMode, Preset};
pub use params::{Pps, Sps};
pub use rc::RateControl;
use rusty_h264_common::{BitWriter, ChromaFormat, NalUnit, NalUnitType, Profile, YuvFrame};
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum EncodeError {
Unsupported(&'static str),
FrameMismatch,
}
impl core::fmt::Display for EncodeError {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
match self {
EncodeError::Unsupported(s) => write!(f, "unsupported: {s}"),
EncodeError::FrameMismatch => write!(f, "frame dimensions do not match encoder config"),
}
}
}
impl std::error::Error for EncodeError {}
#[derive(Debug)]
pub struct Encoder {
cfg: EncoderConfig,
sps: Sps,
pps: Pps,
frame_index: u32,
next_frame_num: u32,
gop_index: u32,
refs: Vec<RefFrame>,
rc: Option<RateControl>,
pending_qpo: Option<Vec<i32>>,
la_queue: Vec<YuvFrame>,
}
impl Drop for Encoder {
fn drop(&mut self) {
debug_assert!(
self.la_queue.is_empty() || std::thread::panicking(),
"Encoder dropped with {} frame(s) still in the lookahead queue — call flush()",
self.la_queue.len()
);
}
}
#[derive(Clone, Debug)]
#[allow(dead_code)]
pub(crate) struct RefFrame {
pub y: rusty_h264_common::aligned::AlignedBytes,
pub u: rusty_h264_common::aligned::AlignedBytes,
pub v: rusty_h264_common::aligned::AlignedBytes,
pub poc: i32,
pub frame_num: u32,
pub mv: Vec<(i32, i32)>,
pub ref_idx: Vec<i32>,
pub w4: usize,
pub hpel: std::sync::OnceLock<std::sync::Arc<rusty_h264_common::inter::HpelPlanes>>,
}
impl RefFrame {
pub(crate) fn hpel(&self, cw: usize, ch: usize) -> &rusty_h264_common::inter::HpelPlanes {
self.hpel.get_or_init(|| {
std::sync::Arc::new(rusty_h264_common::inter::build_hpel_planes(&self.y, cw, ch))
})
}
}
#[cfg(feature = "profile")]
pub fn satdpath_snapshot() -> Vec<u64> { crate::mb16::satdpath::snapshot() }
#[cfg(not(feature = "profile"))]
pub fn satdpath_snapshot() -> Vec<u64> { Vec::new() }
#[cfg(feature = "profile")]
pub fn satdpath_reset() { crate::mb16::satdpath::reset() }
#[cfg(not(feature = "profile"))]
pub fn satdpath_reset() {}
#[cfg(feature = "profile")]
pub fn spstats_redundant() -> u64 { crate::mb16::spstats::redundant_count() }
#[cfg(not(feature = "profile"))]
pub fn spstats_redundant() -> u64 { 0 }
#[cfg(feature = "profile")]
pub fn spstats_snapshot() -> (Vec<u64>, Vec<u64>) { crate::mb16::spstats::snapshot() }
#[cfg(not(feature = "profile"))]
pub fn spstats_snapshot() -> (Vec<u64>, Vec<u64>) { (Vec::new(), Vec::new()) }
#[cfg(feature = "profile")]
pub fn spstats_reset() { crate::mb16::spstats::reset() }
#[cfg(not(feature = "profile"))]
pub fn spstats_reset() {}
pub const DIA_DEFAULT_MASK: u32 = crate::mb16::DIA_DEFAULT;
pub fn set_dia_mask(m: u32) { crate::mb16::set_dia_mask(m) }
pub fn set_me_sadfp(on: bool) { crate::mb16::set_me_sadfp(on) }
pub fn set_me_sadfp_mode(m: u32) { crate::mb16::set_me_sadfp_mode(m) }
pub fn set_me_fc(on: bool) { crate::mb16::set_me_fc(on) }
pub fn set_split_mg(milli: u32) { crate::mb16::set_split_mg(milli) }
pub fn set_mv_smooth(on: bool) { crate::mb16::set_mv_smooth(on) }
pub fn set_mv_smooth_mode(m: u32) { crate::mb16::set_mv_smooth_mode(m) }
pub fn set_sp_fc(on: bool) { crate::mb16::set_sp_fc(on) }
pub fn set_subme(level: u32) {
let (pat, cap) = match level {
1 => (3, 0),
2 => (2, 0),
3 => (0, 2),
4 => (0, 3),
_ => (0, 0),
};
set_subpel_pattern(pat);
crate::mb16::set_sp_maxit(cap);
}
pub fn set_turbo(on: bool) {
set_split_t(if on { 10_000_000 } else { 0 });
}
pub fn set_sp_maxit(n: u32) { crate::mb16::set_sp_maxit(n) }
#[cfg(feature = "profile")]
pub fn diastats_snapshot() -> Vec<(u64, u64)> { crate::mb16::diastats::snapshot() }
#[cfg(not(feature = "profile"))]
pub fn diastats_snapshot() -> Vec<(u64, u64)> { Vec::new() }
#[cfg(feature = "profile")]
pub fn diastats_reset() { crate::mb16::diastats::reset() }
#[cfg(not(feature = "profile"))]
pub fn diastats_reset() {}
pub fn set_defer_subpel(on: bool) {
crate::mb16::DEFER_SUBPEL.store(if on { 1 } else { 0 }, std::sync::atomic::Ordering::Relaxed);
}
pub fn set_split_t(t: u32) {
crate::mb16::SPLIT_T.store(t, std::sync::atomic::Ordering::Relaxed);
}
pub fn set_subpel_dispatch(on: bool) {
crate::mb16::SP_DISPATCH.store(if on { 1 } else { 0 }, std::sync::atomic::Ordering::Relaxed);
}
pub fn set_subpel_pattern(p: u32) {
crate::mb16::SUBPEL_PAT.store(p, std::sync::atomic::Ordering::Relaxed);
}
impl Encoder {
pub fn new(cfg: EncoderConfig) -> Result<Self, EncodeError> {
if !matches!(
cfg.profile,
Profile::ConstrainedBaseline | Profile::Baseline | Profile::Main | Profile::High
) {
return Err(EncodeError::Unsupported("unsupported profile"));
}
if cfg.transform_8x8 && (!matches!(cfg.profile, Profile::High) || cfg.cabac) {
return Err(EncodeError::Unsupported("8x8 transform requires High profile + CAVLC"));
}
if cfg.bframes > 0 && !matches!(cfg.profile, Profile::Main) {
return Err(EncodeError::Unsupported("B-frames require Main profile"));
}
if cfg.chroma != ChromaFormat::Yuv420 {
return Err(EncodeError::Unsupported("only 4:2:0 chroma"));
}
if cfg.width == 0 || cfg.height == 0 || cfg.width % 2 != 0 || cfg.height % 2 != 0 {
return Err(EncodeError::Unsupported("dimensions must be positive and even"));
}
let sps = Sps::from_config(&cfg);
let pps = Pps::from_config(&cfg);
let rc = (cfg.bitrate > 0).then(|| RateControl::new(cfg.bitrate, cfg.framerate, cfg.qp));
Ok(Self {
cfg,
sps,
pps,
frame_index: 0,
next_frame_num: 0,
gop_index: 0,
refs: Vec::new(),
rc,
pending_qpo: None,
la_queue: Vec::new(),
})
}
pub(crate) fn set_pending_qpo(&mut self, qpo: Vec<i32>) {
self.pending_qpo = Some(qpo);
}
pub fn config(&self) -> &EncoderConfig {
&self.cfg
}
pub fn encode(&mut self, frame: &YuvFrame) -> Vec<u8> {
self.try_encode(frame).expect("frame matched config")
}
pub fn try_encode(&mut self, frame: &YuvFrame) -> Result<Vec<u8>, EncodeError> {
if !self.lookahead_active() {
return self.encode_direct(frame);
}
if frame.width != self.cfg.width || frame.height != self.cfg.height || !frame.is_valid() {
return Err(EncodeError::FrameMismatch);
}
self.la_queue.push(frame.clone());
if self.la_queue.len() >= self.cfg.gop_size.max(1) as usize {
self.emit_lookahead_gop()
} else {
Ok(Vec::new())
}
}
fn lookahead_active(&self) -> bool {
self.cfg.mbtree && self.cfg.bframes == 0 && self.cfg.bitrate == 0
}
fn emit_lookahead_gop(&mut self) -> Result<Vec<u8>, EncodeError> {
let frames = std::mem::take(&mut self.la_queue);
let offs = mbtree::gop_qp_offsets(&self.cfg, &frames, self.cfg.mbtree_strength);
let mut out = Vec::new();
for (i, f) in frames.iter().enumerate() {
if let Some(o) = offs.get(i) {
self.pending_qpo = Some(o.clone());
}
out.extend_from_slice(&self.encode_direct(f)?);
}
Ok(out)
}
pub fn flush(&mut self) -> Vec<u8> {
self.try_flush().expect("buffered frames matched the config when accepted")
}
pub fn try_flush(&mut self) -> Result<Vec<u8>, EncodeError> {
if self.la_queue.is_empty() {
return Ok(Vec::new());
}
self.emit_lookahead_gop()
}
fn encode_direct(&mut self, frame: &YuvFrame) -> Result<Vec<u8>, EncodeError> {
let _g = rusty_h264_common::prof::scope(rusty_h264_common::prof::Stage::Total);
if frame.width != self.cfg.width || frame.height != self.cfg.height || !frame.is_valid() {
return Err(EncodeError::FrameMismatch);
}
if self.cfg.bframes > 0 {
return Err(EncodeError::Unsupported("B-frames need encode_all (lookahead)"));
}
let is_idr = self.cfg.gop_size <= 1 || self.frame_index % self.cfg.gop_size == 0;
if is_idr {
self.gop_index = 0;
self.next_frame_num = 0;
self.refs.clear();
}
let frame_num = self.next_frame_num;
let poc_lsb = (2 * self.gop_index) % 16;
let qpo = self.pending_qpo.take().unwrap_or_default();
let complexity = if self.rc.is_some() {
lookahead::complexity(&self.cfg, frame, if is_idr { None } else { self.refs.first() })
} else {
0.0
};
let qp = match &self.rc {
Some(rc) => rc.pick_qp(is_idr, complexity),
None if is_idr => (self.cfg.qp as i32 + self.cfg.i_qp_offset).clamp(0, 51) as u8,
None => self.cfg.qp,
};
let mut out = Vec::new();
let mut w = BitWriter::with_capacity(self.cfg.width * self.cfg.height / 2 + 4096);
let (nal_type, mut reference) = if is_idr {
self.sps.to_nal().write_annex_b(&mut out);
self.pps.to_nal().write_annex_b(&mut out);
slice::write_idr_slice_header(&mut w, &self.cfg, qp);
let r = if self.cfg.cabac {
mb16::encode_slice_data_cabac_intra(&mut w, &self.cfg, frame, qp, &qpo)
} else {
mb16::encode_slice_data(&mut w, &self.cfg, frame, qp, false, &[], &qpo)
};
(NalUnitType::IdrSlice, r)
} else {
slice::write_p_slice_header(&mut w, &self.cfg, qp, frame_num, poc_lsb, self.refs.len());
let r = if self.cfg.cabac {
mb16::encode_slice_data_cabac_p(&mut w, &self.cfg, frame, qp, &self.refs, &qpo)
} else {
mb16::encode_slice_data(&mut w, &self.cfg, frame, qp, true, &self.refs, &qpo)
};
(NalUnitType::NonIdrSlice, r)
};
reference.poc = 2 * self.gop_index as i32;
reference.frame_num = frame_num;
let slice_bytes = w.into_bytes();
if let Some(rc) = &mut self.rc {
rc.update(is_idr, slice_bytes.len() * 8, qp, complexity);
}
{
let _g = rusty_h264_common::prof::scope(rusty_h264_common::prof::Stage::EncNal);
NalUnit::new(3, nal_type, slice_bytes).write_annex_b(&mut out);
}
self.refs.insert(0, reference);
self.refs.truncate(self.cfg.num_ref_frames.max(1) as usize);
self.frame_index += 1;
self.gop_index += 1;
self.next_frame_num = (self.next_frame_num + 1) % 16;
Ok(out)
}
pub fn encode_all(&self, frames: &[YuvFrame]) -> Result<Vec<Vec<u8>>, EncodeError> {
for f in frames {
if f.width != self.cfg.width || f.height != self.cfg.height || !f.is_valid() {
return Err(EncodeError::FrameMismatch);
}
}
if self.cfg.bframes > 0 {
let gop = self.cfg.gop_size.max(1) as usize;
let n_gops = frames.len().div_ceil(gop);
let (w, h) = (self.cfg.width, self.cfg.height);
let gop_sig: Vec<f64> = (0..n_gops)
.map(|g| gop_bi_residual(&frames[g * gop..((g + 1) * gop).min(frames.len())], w, h, 1))
.collect();
let gop_fav: Vec<bool> = if self.cfg.bframes_adaptive {
gop_sig.iter().map(|&s| bframes_favorable(s)).collect()
} else {
vec![true; n_gops]
};
let gop_iqp: Vec<i32> = gop_sig.iter().map(|&s| gop_iqp_offset(s, self.cfg.i_qp_offset)).collect();
let gop_bqp: Vec<i32> = gop_sig.iter().map(|&s| gop_bframe_qp_offset(s, self.cfg.bframe_qp_offset)).collect();
let bcount = if self.cfg.bframes_adaptive {
adaptive_bcount(frames, w, h, self.cfg.bframes as usize)
} else {
self.cfg.bframes as usize
};
if gop_fav.iter().any(|&f| f) {
return Ok(self.encode_all_bframes(frames, bcount, &gop_fav, &gop_iqp, &gop_bqp));
}
let mut pcfg = self.cfg.clone();
pcfg.bframes = 0;
return Encoder::new(pcfg)?.encode_all(frames);
}
if self.cfg.bitrate > 0 {
let mut enc = Encoder::new(self.cfg.clone())?;
let offs: Vec<Vec<i32>> = if self.cfg.mbtree {
let gop = self.cfg.gop_size.max(1) as usize;
frames
.chunks(gop)
.flat_map(|g| mbtree::gop_qp_offsets(&self.cfg, g, self.cfg.mbtree_strength))
.collect()
} else {
Vec::new()
};
return frames
.iter()
.enumerate()
.map(|(i, f)| {
if let Some(qpo) = offs.get(i) {
enc.pending_qpo = Some(qpo.clone());
}
enc.encode_direct(f)
})
.collect();
}
let gop = self.cfg.gop_size.max(1) as usize;
let gops: Vec<&[YuvFrame]> = frames.chunks(gop).collect();
if gops.is_empty() {
return Ok(Vec::new());
}
let n = std::env::var("RUSTY_THREADS")
.ok()
.and_then(|v| v.parse().ok())
.or_else(|| std::thread::available_parallelism().map(|n| n.get()).ok())
.unwrap_or(1)
.min(gops.len());
let mut out: Vec<Option<Vec<Vec<u8>>>> = (0..gops.len()).map(|_| None).collect();
let cfg = &self.cfg;
let gops_ref = &gops;
std::thread::scope(|s| {
let handles: Vec<_> = (0..n)
.map(|t| {
s.spawn(move || {
let mut local = Vec::new();
let mut i = t;
while i < gops_ref.len() {
let mut enc = Encoder::new(cfg.clone()).expect("config");
let offs = if cfg.mbtree {
mbtree::gop_qp_offsets(cfg, gops_ref[i], cfg.mbtree_strength)
} else {
Vec::new()
};
let aus: Vec<Vec<u8>> = gops_ref[i]
.iter()
.enumerate()
.map(|(fi, f)| {
if let Some(o) = offs.get(fi) {
enc.set_pending_qpo(o.clone());
}
enc.encode_direct(f).expect("frame matched config")
})
.collect();
local.push((i, aus));
i += n;
}
local
})
})
.collect();
for h in handles {
for (i, aus) in h.join().expect("encode worker panicked") {
out[i] = Some(aus);
}
}
});
Ok(out.into_iter().flatten().flatten().collect())
}
fn encode_all_bframes(&self, frames: &[YuvFrame], bcount: usize, gop_favorable: &[bool], gop_iqp: &[i32], gop_bqp: &[i32]) -> Vec<Vec<u8>> {
let n = frames.len();
if n == 0 {
return Vec::new();
}
let step = bcount.max(1) + 1; let gop = self.cfg.gop_size.max(1) as usize;
let mut cfg = self.cfg.clone();
cfg.num_ref_frames = cfg.num_ref_frames.max(2);
let sps = Sps::from_config(&cfg);
let pps = Pps::from_config(&cfg);
let mut is_anchor = vec![false; n];
for (d, a) in is_anchor.iter_mut().enumerate() {
*a = if gop_favorable.get(d / gop).copied().unwrap_or(true) {
d % gop == 0 || (d % gop) % step == 0 || (d + 1) % gop == 0
} else {
true
};
}
is_anchor[n - 1] = true;
let mbtree_off: Vec<Vec<i32>> = if cfg.mbtree {
let mut off = vec![Vec::new(); n];
let mut g = 0;
while g < n {
let gop_end = (g + gop).min(n);
let anchors: Vec<usize> = (g..gop_end).filter(|&d| is_anchor[d]).collect();
let aframes: Vec<YuvFrame> = anchors.iter().map(|&d| frames[d].clone()).collect();
let offs = mbtree::gop_qp_offsets(&cfg, &aframes, cfg.mbtree_strength);
for (i, &d) in anchors.iter().enumerate() {
off[d] = offs[i].clone();
}
g = gop_end;
}
off
} else {
Vec::new()
};
let mut order: Vec<usize> = Vec::with_capacity(n);
let mut prev: Option<usize> = None;
for d in 0..n {
if !is_anchor[d] {
continue;
}
order.push(d);
if let Some(p) = prev {
order.extend((p + 1)..d);
}
prev = Some(d);
}
let mut dpb: Vec<RefFrame> = Vec::new();
let mut aus: Vec<Vec<u8>> = Vec::with_capacity(n);
let mut frame_num: u32 = 0;
for &d in &order {
let is_idr = d % gop == 0;
if is_idr {
dpb.clear();
frame_num = 0;
}
let is_b = !is_anchor[d];
let gop_start = (d / gop) * gop;
let poc = ((d - gop_start) as i32) * 2; let iqp = gop_iqp.get(d / gop).copied().unwrap_or(cfg.i_qp_offset);
let bqp = gop_bqp.get(d / gop).copied().unwrap_or(cfg.bframe_qp_offset);
let qpo: &[i32] = mbtree_off.get(d).map(|v| v.as_slice()).unwrap_or(&[]);
let (au, recon) =
code_picture(&cfg, &sps, &pps, &frames[d], is_idr, is_b, poc, frame_num, &dpb, iqp, bqp, qpo);
aus.push(au);
if !is_b {
if let Some(r) = recon {
dpb.insert(0, r);
dpb.truncate(cfg.num_ref_frames as usize);
}
frame_num = (frame_num + 1) % 16;
}
}
aus
}
}
#[allow(clippy::too_many_arguments)]
fn code_picture(
cfg: &EncoderConfig,
sps: &Sps,
pps: &Pps,
frame: &YuvFrame,
is_idr: bool,
is_b: bool,
poc: i32,
frame_num: u32,
dpb: &[RefFrame],
i_qp_offset: i32,
b_qp_offset: i32,
qpo: &[i32],
) -> (Vec<u8>, Option<RefFrame>) {
let mut out = Vec::new();
let mut w = BitWriter::with_capacity(cfg.width * cfg.height / 2 + 4096);
let poc_lsb = (poc as u32) & 0xF; let qp = if is_b {
(cfg.qp as i32 + b_qp_offset).clamp(0, 51) as u8
} else if is_idr {
(cfg.qp as i32 + i_qp_offset).clamp(0, 51) as u8
} else {
cfg.qp
};
let (nal_type, nal_ref_idc, recon) = if is_idr {
sps.to_nal().write_annex_b(&mut out);
pps.to_nal().write_annex_b(&mut out);
slice::write_idr_slice_header(&mut w, cfg, qp);
let mut r = if cfg.cabac {
mb16::encode_slice_data_cabac_intra(&mut w, cfg, frame, qp, qpo)
} else {
mb16::encode_slice_data(&mut w, cfg, frame, qp, false, &[], qpo)
};
r.poc = poc;
r.frame_num = frame_num;
(NalUnitType::IdrSlice, 3u8, Some(r))
} else if is_b {
let l0 = dpb.iter().filter(|r| r.poc < poc).max_by_key(|r| r.poc);
let l1 = dpb.iter().filter(|r| r.poc > poc).min_by_key(|r| r.poc);
slice::write_b_slice_header(&mut w, cfg, qp, frame_num, poc_lsb, 1, 1);
match (l0, l1) {
(Some(l0), Some(l1)) if cfg.cabac => {
mb16::encode_slice_data_cabac_b(&mut w, cfg, frame, qp, poc, l0, l1, &[])
}
(Some(l0), Some(l1)) => mb16::encode_slice_data_b(&mut w, cfg, frame, qp, poc, l0, l1, &[]),
_ => {
let n = cfg.mb_width() * cfg.mb_height();
if cfg.cabac {
mb16::encode_all_skip_b_cabac(&mut w, cfg, qp, n);
} else {
w.write_ue(n as u32);
w.rbsp_trailing_bits();
}
}
}
(NalUnitType::NonIdrSlice, 0u8, None)
} else {
let p_dpb: &[RefFrame] = dpb;
slice::write_p_slice_header(&mut w, cfg, qp, frame_num, poc_lsb, p_dpb.len());
let mut r = if cfg.cabac {
mb16::encode_slice_data_cabac_p(&mut w, cfg, frame, qp, p_dpb, qpo)
} else {
mb16::encode_slice_data(&mut w, cfg, frame, qp, true, dpb, qpo)
};
r.poc = poc;
r.frame_num = frame_num;
(NalUnitType::NonIdrSlice, 3u8, Some(r))
};
let slice_bytes = w.into_bytes();
NalUnit::new(nal_ref_idc, nal_type, slice_bytes).write_annex_b(&mut out);
(out, recon)
}
const BI_THRESH: f64 = 4.0;
fn bframes_favorable(residual: f64) -> bool {
residual < BI_THRESH
}
fn gop_iqp_offset(residual: f64, base: i32) -> i32 {
if base == 0 {
return 0;
}
let bonus = (2.0 * ((BI_THRESH - residual) / BI_THRESH).clamp(0.0, 1.0)).round() as i32;
base - bonus
}
fn gop_bframe_qp_offset(residual: f64, base: i32) -> i32 {
const RAMP: f64 = 0.3; let boost = (4.0 * ((RAMP - residual) / RAMP).clamp(0.0, 1.0)).round() as i32;
base + boost
}
fn adaptive_bcount(frames: &[YuvFrame], w: usize, h: usize, max_b: usize) -> usize {
let cap = max_b.clamp(1, 3);
let g1 = gop_bi_residual(frames, w, h, 1);
let g2 = gop_bi_residual(frames, w, h, 2);
if !g1.is_finite() || !g2.is_finite() {
return 1;
}
let ratio = g2 / g1.max(1e-3);
let c = if ratio >= 1.4 { 1 } else if ratio >= 1.3 { 2 } else { 3 };
c.clamp(1, cap)
}
fn gop_bi_residual(frames: &[YuvFrame], w: usize, h: usize, gap: usize) -> f64 {
let n = frames.len();
if n < 2 * gap + 1 || w < 48 || h < 48 {
return f64::INFINITY;
}
let sad = |cur: &[u8], rf: &[u8], dx: isize, dy: isize| -> u64 {
let mut s = 0u64;
let mut y = 16;
while y < h - 16 {
let cbase = (y * w) as isize;
let rbase = ((y as isize + dy) * w as isize) + dx;
let mut x = 16isize;
while x < (w - 16) as isize {
let c = cur[(cbase + x) as usize] as i32;
let r = rf[(rbase + x) as usize] as i32;
s += (c - r).unsigned_abs() as u64;
x += 8;
}
y += 8;
}
s
};
let global_me = |cur: &[u8], rf: &[u8]| -> (isize, isize) {
let (mut best, mut bc) = ((0isize, 0isize), u64::MAX);
let mut dy = -12;
while dy <= 12 {
let mut dx = -12;
while dx <= 12 {
let c = sad(cur, rf, dx, dy);
if c < bc {
bc = c;
best = (dx, dy);
}
dx += 4;
}
dy += 4;
}
for dy in best.1 - 3..=best.1 + 3 {
for dx in best.0 - 3..=best.0 + 3 {
let c = sad(cur, rf, dx, dy);
if c < bc {
bc = c;
best = (dx, dy);
}
}
}
best
};
let mut n_samp = 0usize;
{
let mut y = 16;
while y < h - 16 {
let mut x = 16;
while x < w - 16 {
n_samp += 1;
x += 8;
}
y += 8;
}
}
let step = (n / 5).max(1);
let (mut total, mut cnt) = (0f64, 0usize);
let mut d = gap;
while d < n - gap {
let (cur, past, fut) = (&frames[d].y, &frames[d - gap].y, &frames[d + gap].y);
let (mpx, mpy) = global_me(cur, past);
let (mfx, mfy) = global_me(cur, fut);
let mut bi = 0u64;
let mut y = 16;
while y < h - 16 {
let mut x = 16isize;
while x < (w - 16) as isize {
let c = cur[y * w + x as usize] as i32;
let p = past[((y as isize + mpy) * w as isize + x + mpx) as usize] as i32;
let f = fut[((y as isize + mfy) * w as isize + x + mfx) as usize] as i32;
bi += (c - ((p + f + 1) >> 1)).unsigned_abs() as u64;
x += 8;
}
y += 8;
}
total += bi as f64 / n_samp as f64;
cnt += 1;
d += step;
}
if cnt > 0 {
total / cnt as f64
} else {
f64::INFINITY
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn rejects_unsupported_config() {
let mut cfg = EncoderConfig::new(16, 16);
cfg.profile = Profile::High;
cfg.transform_8x8 = true;
cfg.cabac = true;
assert!(matches!(Encoder::new(cfg), Err(EncodeError::Unsupported(_))));
}
#[test]
fn encodes_access_unit_with_sps_pps_idr() {
use rusty_h264_common::nal::split_annex_b;
let cfg = EncoderConfig::new(32, 32);
let mut enc = Encoder::new(cfg).unwrap();
let frame = YuvFrame::black(32, 32);
let au = enc.encode(&frame);
let nals = split_annex_b(&au);
assert_eq!(nals.len(), 3);
assert_eq!(NalUnitType::from_id(nals[0][0]), NalUnitType::Sps);
assert_eq!(NalUnitType::from_id(nals[1][0]), NalUnitType::Pps);
assert_eq!(NalUnitType::from_id(nals[2][0]), NalUnitType::IdrSlice);
}
#[test]
fn encode_all_matches_sequential_cqp() {
let (w, h) = (48usize, 32usize);
let mut cfg = EncoderConfig::new(w, h);
cfg.gop_size = 4; let frames: Vec<YuvFrame> = (0..10u8)
.map(|t| YuvFrame {
width: w,
height: h,
y: (0..w * h).map(|i| (i as u8).wrapping_add(t.wrapping_mul(7))).collect(),
u: vec![128u8.wrapping_add(t); (w / 2) * (h / 2)],
v: vec![128u8.wrapping_sub(t); (w / 2) * (h / 2)],
})
.collect();
let mut seq_enc = Encoder::new(cfg.clone()).unwrap();
let mut seq: Vec<u8> = frames.iter().flat_map(|f| seq_enc.encode(f)).collect();
seq.extend_from_slice(&seq_enc.flush()); let par: Vec<u8> = Encoder::new(cfg).unwrap().encode_all(&frames).unwrap().concat();
assert_eq!(seq, par, "GOP-parallel must equal sequential+flush at CQP");
}
#[test]
fn encode_all_matches_sequential_quality_preset() {
let (w, h) = (48usize, 32usize);
let mut cfg = EncoderConfig::new(w, h);
cfg.gop_size = 3; cfg.preset = crate::config::Preset::Quality;
let frames: Vec<YuvFrame> = (0..12u8)
.map(|t| YuvFrame {
width: w,
height: h,
y: (0..w * h)
.map(|i| {
let base = (i as u8).wrapping_add(t.wrapping_mul(3));
if t % 2 == 0 { base } else { base.wrapping_mul(37).wrapping_add(i as u8) }
})
.collect(),
u: vec![128u8.wrapping_add(t); (w / 2) * (h / 2)],
v: vec![128u8.wrapping_sub(t); (w / 2) * (h / 2)],
})
.collect();
let mut seq_enc = Encoder::new(cfg.clone()).unwrap();
let mut seq: Vec<u8> = frames.iter().flat_map(|f| seq_enc.encode(f)).collect();
seq.extend_from_slice(&seq_enc.flush());
let par: Vec<u8> = Encoder::new(cfg).unwrap().encode_all(&frames).unwrap().concat();
assert_eq!(seq, par, "quality-preset GOP-parallel must equal sequential+flush");
}
#[test]
fn rejects_mismatched_frame() {
let cfg = EncoderConfig::new(16, 16);
let mut enc = Encoder::new(cfg).unwrap();
let frame = YuvFrame::black(32, 16);
assert_eq!(enc.try_encode(&frame), Err(EncodeError::FrameMismatch));
}
}