From 19cfcd2f207b472c03de4c14db85d0cac1d93fa9 Mon Sep 17 00:00:00 2001 From: Kostya Shishkov Date: Wed, 2 Mar 2022 18:31:00 +0100 Subject: [PATCH] vp6enc: split out future common parts to share them with VP7 encoder --- nihav-duck/src/codecs/mod.rs | 3 + nihav-duck/src/codecs/vp6enc/coder.rs | 185 ++-------------------- nihav-duck/src/codecs/vp6enc/dsp.rs | 172 +------------------- nihav-duck/src/codecs/vp6enc/models.rs | 78 +-------- nihav-duck/src/codecs/vpenc/coder.rs | 181 +++++++++++++++++++++ nihav-duck/src/codecs/vpenc/mod.rs | 5 + nihav-duck/src/codecs/vpenc/models.rs | 75 +++++++++ nihav-duck/src/codecs/vpenc/motion_est.rs | 182 +++++++++++++++++++++ 8 files changed, 467 insertions(+), 414 deletions(-) create mode 100644 nihav-duck/src/codecs/vpenc/coder.rs create mode 100644 nihav-duck/src/codecs/vpenc/mod.rs create mode 100644 nihav-duck/src/codecs/vpenc/models.rs create mode 100644 nihav-duck/src/codecs/vpenc/motion_est.rs diff --git a/nihav-duck/src/codecs/mod.rs b/nihav-duck/src/codecs/mod.rs index 6086150..e2ba63e 100644 --- a/nihav-duck/src/codecs/mod.rs +++ b/nihav-duck/src/codecs/mod.rs @@ -124,6 +124,9 @@ pub fn duck_register_all_decoders(rd: &mut RegisteredDecoders) { } } +#[cfg(feature="encoder_vp6")] +#[macro_use] +mod vpenc; #[cfg(feature="encoder_vp6")] #[allow(clippy::needless_range_loop)] mod vp6enc; diff --git a/nihav-duck/src/codecs/vp6enc/coder.rs b/nihav-duck/src/codecs/vp6enc/coder.rs index 74b490b..cb0904f 100644 --- a/nihav-duck/src/codecs/vp6enc/coder.rs +++ b/nihav-duck/src/codecs/vp6enc/coder.rs @@ -1,38 +1,12 @@ -use nihav_core::io::byteio::*; use nihav_core::codecs::{EncoderResult, EncoderError}; use nihav_codec_support::codecs::MV; +use crate::codecs::vpenc::coder::*; +use crate::codecs::vpenc::models::*; use super::super::vpcommon::*; use super::super::vp6data::*; use super::models::*; -struct EncSeq { - bit: bool, - idx: u8, -} - -pub struct TokenSeq { - val: T, - seq: &'static [EncSeq], -} - -macro_rules! bit_entry { - (T; $idx:expr) => {EncSeq {bit: true, idx: $idx }}; - (F; $idx:expr) => {EncSeq {bit: false, idx: $idx }}; -} - -macro_rules! bit_seq { - ($val: expr; $( $bit:tt),* ; $( $idx:expr),* ) => { - TokenSeq { - val: $val, - seq: - &[ - $( - bit_entry!($bit; $idx), - )* - ] - } - }; -} +pub use crate::codecs::vpenc::coder::{BoolEncoder, Estimator}; pub const MODE_TREE: &[TokenSeq] = &[ bit_seq!(VPMBType::Intra; T, F, F; 0, 2, 5), @@ -138,126 +112,15 @@ const ZERO_RUN_TREE: &[TokenSeq] = &[ bit_seq!(9; T, T; 0, 4), ]; -pub struct BoolEncoder<'a, 'b> { - bw: &'a mut ByteWriter<'b>, - val: u32, - range: u32, - bits: u8, - saved: u8, - run: usize, +pub trait EncoderTrait { + fn write_cat(&mut self, cat: i8, tree: &[TokenSeq], tok_probs: &[u8], val_probs: &[u8; 11]) -> EncoderResult<()>; + fn write_large_coef(&mut self, val: i16, cat: usize) -> EncoderResult<()>; + fn write_dc(&mut self, val: i16, tok_probs: &[u8; 5], val_probs: &[u8; 11]) -> EncoderResult<()>; + fn write_ac(&mut self, val: i16, tree: &[TokenSeq], probs: &[u8; 11]) -> EncoderResult<()>; + fn write_zero_run(&mut self, val: usize, probs: &[u8; 14]) -> EncoderResult<()>; } -impl<'a, 'b> BoolEncoder<'a, 'b> { - pub fn new(bw: &'a mut ByteWriter<'b>) -> Self { - Self { - bw, - val: 0, - range: 255, - bits: 0, - saved: 0, - run: 0, - } - } - pub fn put_bool(&mut self, bit: bool, prob: u8) -> EncoderResult<()> { - let split = 1 + (((self.range - 1) * u32::from(prob)) >> 8); - if bit { - self.range -= split; - self.val += split; - } else { - self.range = split; - } - - if self.range < 128 { - self.renorm()?; - } - Ok(()) - } - fn flush_run(&mut self, overflow: bool) -> EncoderResult<()> { - if self.run > 0 { - self.bw.write_byte(self.saved + (overflow as u8))?; - if !overflow { - for _ in 1..self.run { - self.bw.write_byte(0xFF)?; - } - } else { - for _ in 1..self.run { - self.bw.write_byte(0)?; - } - } - self.run = 0; - } - Ok(()) - } - fn renorm(&mut self) -> EncoderResult<()> { - let bits = (self.range.leading_zeros() & 7) as u8; - self.range <<= bits; - if self.bits + bits < 23 { - self.bits += bits; - self.val <<= bits; - } else { - for _ in 0..bits { - if (self.bits == 23) && ((self.val >> 31) != 0) { - self.flush_run(true)?; - } - self.val <<= 1; - self.bits += 1; - if self.bits == 24 { - let tbyte = (self.val >> 24) as u8; - let nbyte = (self.val >> 16) as u8; - if tbyte < 0xFF { - self.flush_run(false)?; - if nbyte < 0xFE { - self.bw.write_byte(tbyte)?; - } else { - self.saved = tbyte; - self.run = 1; - } - } else { - self.run += 1; - } - self.val &= 0xFFFFFF; - self.bits -= 8; - } - } - } - Ok(()) - } - pub fn flush(mut self) -> EncoderResult<()> { - self.flush_run(false)?; - self.val <<= 24 - self.bits; - self.bw.write_u32be(self.val)?; - Ok(()) - } - - pub fn put_bits(&mut self, val: u32, len: u8) -> EncoderResult<()> { - let mut mask = 1 << (len - 1); - while mask != 0 { - self.put_bool((val & mask) != 0, 128)?; - mask >>= 1; - } - Ok(()) - } - fn put_probability(&mut self, prob: u8) -> EncoderResult<()> { - self.put_bits(u32::from(prob >> 1), 7) - } - fn encode_probability(&mut self, new: u8, old: u8, prob: u8) -> EncoderResult<()> { - self.put_bool(new != old, prob)?; - if new != old { - self.put_probability(new)?; - } - Ok(()) - } - pub fn write_el(&mut self, el: T, tree: &[TokenSeq], probs: &[u8]) -> EncoderResult<()> { - for entry in tree.iter() { - if entry.val == el { - for seq in entry.seq.iter() { - self.put_bool(seq.bit, probs[seq.idx as usize])?; - } - return Ok(()); - } - } - Err(EncoderError::Bug) - } +impl<'a, 'b> EncoderTrait for BoolEncoder<'a, 'b> { fn write_cat(&mut self, cat: i8, tree: &[TokenSeq], tok_probs: &[u8], val_probs: &[u8; 11]) -> EncoderResult<()> { for entry in tree.iter() { if entry.val == cat { @@ -679,20 +542,14 @@ pub fn encode_mv(bc: &mut BoolEncoder, mv: MV, model: &VP56Models) -> EncoderRes Ok(()) } -struct Estimator {} +pub trait VP6EstimatorTrait { + fn write_cat(&self, cat: i8, tree: &[TokenSeq], probs: &mut [ProbCounter; 11]); + fn write_dc(&self, val: i16, probs: &mut [ProbCounter; 11]); + fn write_ac(&self, val: i16, tree: &[TokenSeq], probs: &mut [ProbCounter; 11]); + fn write_zero_run(&self, val: usize, probs: &mut [ProbCounter; 14]); +} -impl Estimator { - fn new() -> Self { Self{} } - fn write_el(&self, el: T, tree: &[TokenSeq], probs: &mut [ProbCounter]) { - for entry in tree.iter() { - if entry.val == el { - for seq in entry.seq.iter() { - probs[seq.idx as usize].add(seq.bit); - } - return; - } - } - } +impl VP6EstimatorTrait for Estimator { fn write_cat(&self, cat: i8, tree: &[TokenSeq], probs: &mut [ProbCounter; 11]) { for entry in tree.iter() { if entry.val == cat { @@ -718,14 +575,6 @@ impl Estimator { } } } - fn est_nits(bit: bool, prob: u8) -> u32 { - if !bit { - u32::from(PROB_BITS[prob as usize]) - } else { - u32::from(PROB_BITS[256 - (prob as usize)]) - } - } - fn nits_to_bits(nits: u32) -> u32 { (nits + 7) >> 3 } } pub fn estimate_block(blk: &[i16; 64], _dc_mode: usize, model: &mut VP56CoeffModelStat, vp6model: &mut VP6ModelsStat, scan: &[usize; 64]) { diff --git a/nihav-duck/src/codecs/vp6enc/dsp.rs b/nihav-duck/src/codecs/vp6enc/dsp.rs index 41b72cf..b83cd11 100644 --- a/nihav-duck/src/codecs/vp6enc/dsp.rs +++ b/nihav-duck/src/codecs/vp6enc/dsp.rs @@ -4,45 +4,8 @@ use super::super::vpcommon::*; use super::super::vp6dsp::*; use super::super::vp6data::*; use super::ResidueMB; - -use std::str::FromStr; - -#[derive(Debug,Clone,Copy,PartialEq)] -pub enum MVSearchMode { - Full, - Diamond, - Hexagon, -} - -impl Default for MVSearchMode { - fn default() -> Self { MVSearchMode::Hexagon } -} - -pub struct ParseError{} - -impl FromStr for MVSearchMode { - type Err = ParseError; - - #[allow(clippy::single_match)] - fn from_str(s: &str) -> Result { - match s { - "full" => Ok(MVSearchMode::Full), - "dia" => Ok(MVSearchMode::Diamond), - "hex" => Ok(MVSearchMode::Hexagon), - _ => Err(ParseError{}), - } - } -} - -impl std::fmt::Display for MVSearchMode { - fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { - match *self { - MVSearchMode::Full => write!(f, "full"), - MVSearchMode::Diamond => write!(f, "dia"), - MVSearchMode::Hexagon => write!(f, "hex"), - } - } -} +use crate::codecs::vpenc::motion_est::*; +pub use crate::codecs::vpenc::motion_est::MVSearchMode; const C1S7: i32 = 64277; @@ -177,137 +140,6 @@ impl MVSearch for FullMVSearch { } } -const DIA_PATTERN: [MV; 9] = [ - ZERO_MV, - MV {x: -2, y: 0}, - MV {x: -1, y: 1}, - MV {x: 0, y: 2}, - MV {x: 1, y: 1}, - MV {x: 2, y: 0}, - MV {x: 1, y: -1}, - MV {x: 0, y: -2}, - MV {x: -1, y: -1} -]; - -const HEX_PATTERN: [MV; 7] = [ - ZERO_MV, - MV {x: -2, y: 0}, - MV {x: -1, y: 2}, - MV {x: 1, y: 2}, - MV {x: 2, y: 0}, - MV {x: 1, y: -2}, - MV {x: -1, y: -2} -]; - -const REFINEMENT: [MV; 4] = [ - MV {x: -1, y: 0}, - MV {x: 0, y: 1}, - MV {x: 1, y: 0}, - MV {x: 0, y: -1} -]; - -macro_rules! search_template { - ($self: expr, $mv_est: expr, $cur_blk: expr, $mb_x: expr, $mb_y: expr, $sad_func: ident) => ({ - let mut best_dist = MAX_DIST; - let mut best_mv; - - let mut min_dist; - let mut min_idx; - - $self.reset(); - loop { - let mut cur_best_dist = best_dist; - for (dist, &point) in $self.dist.iter_mut().zip($self.point.iter()) { - if *dist == MAX_DIST { - *dist = $mv_est.$sad_func($cur_blk, $mb_x, $mb_y, point.from_pixels(), cur_best_dist); - cur_best_dist = cur_best_dist.min(*dist); - if *dist <= DIST_THRESH { - break; - } - } - } - min_dist = $self.dist[0]; - min_idx = 0; - for (i, &dist) in $self.dist.iter().enumerate().skip(1) { - if dist < min_dist { - min_dist = dist; - min_idx = i; - if dist <= DIST_THRESH { - break; - } - } - } - if min_dist <= DIST_THRESH || min_idx == 0 || best_dist == min_dist || $self.point[min_idx].x.abs() >= $mv_est.mv_range || $self.point[min_idx].y.abs() >= $mv_est.mv_range { - break; - } - best_dist = min_dist; - $self.update($self.steps[min_idx]); - } - best_dist = min_dist; - best_mv = $self.point[min_idx]; - if best_dist <= DIST_THRESH { - return (best_mv.from_pixels(), best_dist); - } - for &step in REFINEMENT.iter() { - let mv = best_mv + step; - let dist = $mv_est.$sad_func($cur_blk, $mb_x, $mb_y, mv.from_pixels(), MAX_DIST); - if best_dist > dist { - best_dist = dist; - best_mv = mv; - } - } - best_mv = best_mv.from_pixels(); - if best_dist <= DIST_THRESH { - return (best_mv, best_dist); - } - - // subpel refinement - $self.set_new_point(best_mv, best_dist); - loop { - let mut cur_best_dist = best_dist; - for (dist, &point) in $self.dist.iter_mut().zip($self.point.iter()) { - if *dist == MAX_DIST { - *dist = $mv_est.$sad_func($cur_blk, $mb_x, $mb_y, point, cur_best_dist); - cur_best_dist = cur_best_dist.min(*dist); - if *dist <= DIST_THRESH { - break; - } - } - } - min_dist = $self.dist[0]; - min_idx = 0; - for (i, &dist) in $self.dist.iter().enumerate().skip(1) { - if dist < min_dist { - min_dist = dist; - min_idx = i; - if dist <= DIST_THRESH { - break; - } - } - } - if min_dist <= DIST_THRESH || min_idx == 0 || best_dist == min_dist || $self.point[min_idx].x.abs() >= $mv_est.mv_range * 4 || $self.point[min_idx].y.abs() >= $mv_est.mv_range * 4 { - break; - } - best_dist = min_dist; - $self.update($self.steps[min_idx]); - } - best_dist = min_dist; - best_mv = $self.point[min_idx]; - if best_dist <= DIST_THRESH { - return (best_mv, best_dist); - } - for &step in REFINEMENT.iter() { - let mv = best_mv + step; - let dist = $mv_est.$sad_func($cur_blk, $mb_x, $mb_y, mv, MAX_DIST); - if best_dist > dist { - best_dist = dist; - best_mv = mv; - } - } - (best_mv, best_dist) - }) -} - macro_rules! pattern_search { ($struct_name: ident, $patterns: expr) => { pub struct $struct_name { diff --git a/nihav-duck/src/codecs/vp6enc/models.rs b/nihav-duck/src/codecs/vp6enc/models.rs index c6344a3..7a0da4a 100644 --- a/nihav-duck/src/codecs/vp6enc/models.rs +++ b/nihav-duck/src/codecs/vp6enc/models.rs @@ -1,5 +1,7 @@ use nihav_codec_support::codecs::ZIGZAG; use super::super::vp6data::*; +use crate::codecs::vpenc::models::*; +pub use crate::codecs::vpenc::models::PROB_BITS; #[derive(Clone,Copy,Default)] pub struct VP56MVModel { @@ -233,82 +235,6 @@ pub fn reset_scan(model: &mut VP6Models, interlaced: bool) { model.zigzag.copy_from_slice(&ZIGZAG); } -#[derive(Clone,Copy,Default)] -pub struct ProbCounter { - zeroes: u32, - total: u32, -} - -// bits to code zero probability multiplied by eight -pub const PROB_BITS: [u8; 256] = [ - 0, 64, 56, 51, 48, 45, 43, 42, - 40, 39, 37, 36, 35, 34, 34, 33, - 32, 31, 31, 30, 29, 29, 28, 28, - 27, 27, 26, 26, 26, 25, 25, 24, - 24, 24, 23, 23, 23, 22, 22, 22, - 21, 21, 21, 21, 20, 20, 20, 20, - 19, 19, 19, 19, 18, 18, 18, 18, - 18, 17, 17, 17, 17, 17, 16, 16, - 16, 16, 16, 15, 15, 15, 15, 15, - 15, 14, 14, 14, 14, 14, 14, 14, - 13, 13, 13, 13, 13, 13, 13, 12, - 12, 12, 12, 12, 12, 12, 12, 11, - 11, 11, 11, 11, 11, 11, 11, 11, - 10, 10, 10, 10, 10, 10, 10, 10, - 10, 9, 9, 9, 9, 9, 9, 9, - 9, 9, 9, 8, 8, 8, 8, 8, - 8, 8, 8, 8, 8, 8, 7, 7, - 7, 7, 7, 7, 7, 7, 7, 7, - 7, 7, 6, 6, 6, 6, 6, 6, - 6, 6, 6, 6, 6, 6, 6, 5, - 5, 5, 5, 5, 5, 5, 5, 5, - 5, 5, 5, 5, 5, 5, 4, 4, - 4, 4, 4, 4, 4, 4, 4, 4, - 4, 4, 4, 4, 4, 4, 3, 3, - 3, 3, 3, 3, 3, 3, 3, 3, - 3, 3, 3, 3, 3, 3, 3, 2, - 2, 2, 2, 2, 2, 2, 2, 2, - 2, 2, 2, 2, 2, 2, 2, 2, - 2, 1, 1, 1, 1, 1, 1, 1, - 1, 1, 1, 1, 1, 1, 1, 1, - 1, 1, 1, 1, 1, 1, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0 -]; - -impl ProbCounter { - pub fn add(&mut self, b: bool) { - if !b { - self.zeroes += 1; - } - self.total += 1; - } - pub fn to_prob(self) -> u8 { - if self.total > 0 { - (((self.zeroes << 8) / self.total).min(254) & !1).max(1) as u8 - } else { - 128 - } - } - pub fn to_prob_worthy(&self, old_prob: u8) -> u8 { - if self.total > 0 { - let new_prob = self.to_prob(); - let new_bits = Self::est_bits(new_prob, self.zeroes, self.total); - let old_bits = Self::est_bits(old_prob, self.zeroes, self.total); - - if new_bits + 7 < old_bits { - new_prob - } else { - old_prob - } - } else { - old_prob - } - } - fn est_bits(prob: u8, zeroes: u32, total: u32) -> u32 { - (u32::from(PROB_BITS[prob as usize]) * zeroes + u32::from(PROB_BITS[256 - (prob as usize)]) * (total - zeroes) + 7) >> 3 - } -} - #[derive(Clone,Copy,Default)] pub struct VP56MVModelStat { pub nz_prob: ProbCounter, diff --git a/nihav-duck/src/codecs/vpenc/coder.rs b/nihav-duck/src/codecs/vpenc/coder.rs new file mode 100644 index 0000000..60b1fd3 --- /dev/null +++ b/nihav-duck/src/codecs/vpenc/coder.rs @@ -0,0 +1,181 @@ +use nihav_core::io::byteio::*; +use nihav_core::codecs::{EncoderResult, EncoderError}; +use super::models::*; + +pub struct EncSeq { + pub bit: bool, + pub idx: u8, +} + +pub struct TokenSeq { + pub val: T, + pub seq: &'static [EncSeq], +} + +#[macro_export] +macro_rules! bit_entry { + (T; $idx:expr) => {EncSeq {bit: true, idx: $idx }}; + (F; $idx:expr) => {EncSeq {bit: false, idx: $idx }}; +} + +#[macro_export] +macro_rules! bit_seq { + ($val: expr; $( $bit:tt),* ; $( $idx:expr),* ) => { + TokenSeq { + val: $val, + seq: + &[ + $( + bit_entry!($bit; $idx), + )* + ] + } + }; +} + +pub struct BoolEncoder<'a, 'b> { + bw: &'a mut ByteWriter<'b>, + val: u32, + range: u32, + bits: u8, + saved: u8, + run: usize, +} + +impl<'a, 'b> BoolEncoder<'a, 'b> { + pub fn new(bw: &'a mut ByteWriter<'b>) -> Self { + Self { + bw, + val: 0, + range: 255, + bits: 0, + saved: 0, + run: 0, + } + } + pub fn put_bool(&mut self, bit: bool, prob: u8) -> EncoderResult<()> { + let split = 1 + (((self.range - 1) * u32::from(prob)) >> 8); + if bit { + self.range -= split; + self.val += split; + } else { + self.range = split; + } + + if self.range < 128 { + self.renorm()?; + } + Ok(()) + } + fn flush_run(&mut self, overflow: bool) -> EncoderResult<()> { + if self.run > 0 { + self.bw.write_byte(self.saved + (overflow as u8))?; + if !overflow { + for _ in 1..self.run { + self.bw.write_byte(0xFF)?; + } + } else { + for _ in 1..self.run { + self.bw.write_byte(0)?; + } + } + self.run = 0; + } + Ok(()) + } + fn renorm(&mut self) -> EncoderResult<()> { + let bits = (self.range.leading_zeros() & 7) as u8; + self.range <<= bits; + if self.bits + bits < 23 { + self.bits += bits; + self.val <<= bits; + } else { + for _ in 0..bits { + if (self.bits == 23) && ((self.val >> 31) != 0) { + self.flush_run(true)?; + } + self.val <<= 1; + self.bits += 1; + if self.bits == 24 { + let tbyte = (self.val >> 24) as u8; + let nbyte = (self.val >> 16) as u8; + if tbyte < 0xFF { + self.flush_run(false)?; + if nbyte < 0xFE { + self.bw.write_byte(tbyte)?; + } else { + self.saved = tbyte; + self.run = 1; + } + } else { + self.run += 1; + } + self.val &= 0xFFFFFF; + self.bits -= 8; + } + } + } + Ok(()) + } + pub fn flush(mut self) -> EncoderResult<()> { + self.flush_run(false)?; + self.val <<= 24 - self.bits; + self.bw.write_u32be(self.val)?; + Ok(()) + } + + pub fn put_bits(&mut self, val: u32, len: u8) -> EncoderResult<()> { + let mut mask = 1 << (len - 1); + while mask != 0 { + self.put_bool((val & mask) != 0, 128)?; + mask >>= 1; + } + Ok(()) + } + pub fn put_probability(&mut self, prob: u8) -> EncoderResult<()> { + self.put_bits(u32::from(prob >> 1), 7) + } + pub fn encode_probability(&mut self, new: u8, old: u8, prob: u8) -> EncoderResult<()> { + self.put_bool(new != old, prob)?; + if new != old { + self.put_probability(new)?; + } + Ok(()) + } + pub fn write_el(&mut self, el: T, tree: &[TokenSeq], probs: &[u8]) -> EncoderResult<()> { + for entry in tree.iter() { + if entry.val == el { + for seq in entry.seq.iter() { + self.put_bool(seq.bit, probs[seq.idx as usize])?; + } + return Ok(()); + } + } + Err(EncoderError::Bug) + } +} + +pub struct Estimator {} + +#[allow(dead_code)] +impl Estimator { + pub fn new() -> Self { Self{} } + pub fn write_el(&self, el: T, tree: &[TokenSeq], probs: &mut [ProbCounter]) { + for entry in tree.iter() { + if entry.val == el { + for seq in entry.seq.iter() { + probs[seq.idx as usize].add(seq.bit); + } + return; + } + } + } + pub fn est_nits(bit: bool, prob: u8) -> u32 { + if !bit { + u32::from(PROB_BITS[prob as usize]) + } else { + u32::from(PROB_BITS[256 - (prob as usize)]) + } + } + pub fn nits_to_bits(nits: u32) -> u32 { (nits + 7) >> 3 } +} diff --git a/nihav-duck/src/codecs/vpenc/mod.rs b/nihav-duck/src/codecs/vpenc/mod.rs new file mode 100644 index 0000000..410d6e4 --- /dev/null +++ b/nihav-duck/src/codecs/vpenc/mod.rs @@ -0,0 +1,5 @@ +#[macro_use] +pub mod coder; +pub mod models; +#[macro_use] +pub mod motion_est; diff --git a/nihav-duck/src/codecs/vpenc/models.rs b/nihav-duck/src/codecs/vpenc/models.rs new file mode 100644 index 0000000..0495c89 --- /dev/null +++ b/nihav-duck/src/codecs/vpenc/models.rs @@ -0,0 +1,75 @@ +#[derive(Clone,Copy,Default)] +pub struct ProbCounter { + zeroes: u32, + total: u32, +} + +// bits to code zero probability multiplied by eight +pub const PROB_BITS: [u8; 256] = [ + 0, 64, 56, 51, 48, 45, 43, 42, + 40, 39, 37, 36, 35, 34, 34, 33, + 32, 31, 31, 30, 29, 29, 28, 28, + 27, 27, 26, 26, 26, 25, 25, 24, + 24, 24, 23, 23, 23, 22, 22, 22, + 21, 21, 21, 21, 20, 20, 20, 20, + 19, 19, 19, 19, 18, 18, 18, 18, + 18, 17, 17, 17, 17, 17, 16, 16, + 16, 16, 16, 15, 15, 15, 15, 15, + 15, 14, 14, 14, 14, 14, 14, 14, + 13, 13, 13, 13, 13, 13, 13, 12, + 12, 12, 12, 12, 12, 12, 12, 11, + 11, 11, 11, 11, 11, 11, 11, 11, + 10, 10, 10, 10, 10, 10, 10, 10, + 10, 9, 9, 9, 9, 9, 9, 9, + 9, 9, 9, 8, 8, 8, 8, 8, + 8, 8, 8, 8, 8, 8, 7, 7, + 7, 7, 7, 7, 7, 7, 7, 7, + 7, 7, 6, 6, 6, 6, 6, 6, + 6, 6, 6, 6, 6, 6, 6, 5, + 5, 5, 5, 5, 5, 5, 5, 5, + 5, 5, 5, 5, 5, 5, 4, 4, + 4, 4, 4, 4, 4, 4, 4, 4, + 4, 4, 4, 4, 4, 4, 3, 3, + 3, 3, 3, 3, 3, 3, 3, 3, + 3, 3, 3, 3, 3, 3, 3, 2, + 2, 2, 2, 2, 2, 2, 2, 2, + 2, 2, 2, 2, 2, 2, 2, 2, + 2, 1, 1, 1, 1, 1, 1, 1, + 1, 1, 1, 1, 1, 1, 1, 1, + 1, 1, 1, 1, 1, 1, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0 +]; + +impl ProbCounter { + pub fn add(&mut self, b: bool) { + if !b { + self.zeroes += 1; + } + self.total += 1; + } + pub fn to_prob(self) -> u8 { + if self.total > 0 { + (((self.zeroes << 8) / self.total).min(254) & !1).max(1) as u8 + } else { + 128 + } + } + pub fn to_prob_worthy(&self, old_prob: u8) -> u8 { + if self.total > 0 { + let new_prob = self.to_prob(); + let new_bits = Self::est_bits(new_prob, self.zeroes, self.total); + let old_bits = Self::est_bits(old_prob, self.zeroes, self.total); + + if new_bits + 7 < old_bits { + new_prob + } else { + old_prob + } + } else { + old_prob + } + } + fn est_bits(prob: u8, zeroes: u32, total: u32) -> u32 { + (u32::from(PROB_BITS[prob as usize]) * zeroes + u32::from(PROB_BITS[256 - (prob as usize)]) * (total - zeroes) + 7) >> 3 + } +} diff --git a/nihav-duck/src/codecs/vpenc/motion_est.rs b/nihav-duck/src/codecs/vpenc/motion_est.rs new file mode 100644 index 0000000..128633f --- /dev/null +++ b/nihav-duck/src/codecs/vpenc/motion_est.rs @@ -0,0 +1,182 @@ +use nihav_codec_support::codecs::{MV, ZERO_MV}; + +use std::str::FromStr; + +#[derive(Debug,Clone,Copy,PartialEq)] +pub enum MVSearchMode { + Full, + Diamond, + Hexagon, +} + +impl Default for MVSearchMode { + fn default() -> Self { MVSearchMode::Hexagon } +} + +pub struct ParseError{} + +impl FromStr for MVSearchMode { + type Err = ParseError; + + #[allow(clippy::single_match)] + fn from_str(s: &str) -> Result { + match s { + "full" => Ok(MVSearchMode::Full), + "dia" => Ok(MVSearchMode::Diamond), + "hex" => Ok(MVSearchMode::Hexagon), + _ => Err(ParseError{}), + } + } +} + +impl std::fmt::Display for MVSearchMode { + fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + match *self { + MVSearchMode::Full => write!(f, "full"), + MVSearchMode::Diamond => write!(f, "dia"), + MVSearchMode::Hexagon => write!(f, "hex"), + } + } +} + +trait FromPixels { + fn from_pixels(self) -> Self; +} + +impl FromPixels for MV { + fn from_pixels(self) -> MV { + MV { x: self.x * 4, y: self.y * 4 } + } +} + +pub const DIA_PATTERN: [MV; 9] = [ + ZERO_MV, + MV {x: -2, y: 0}, + MV {x: -1, y: 1}, + MV {x: 0, y: 2}, + MV {x: 1, y: 1}, + MV {x: 2, y: 0}, + MV {x: 1, y: -1}, + MV {x: 0, y: -2}, + MV {x: -1, y: -1} +]; + +pub const HEX_PATTERN: [MV; 7] = [ + ZERO_MV, + MV {x: -2, y: 0}, + MV {x: -1, y: 2}, + MV {x: 1, y: 2}, + MV {x: 2, y: 0}, + MV {x: 1, y: -2}, + MV {x: -1, y: -2} +]; + +pub const REFINEMENT: [MV; 4] = [ + MV {x: -1, y: 0}, + MV {x: 0, y: 1}, + MV {x: 1, y: 0}, + MV {x: 0, y: -1} +]; + +#[macro_export] +macro_rules! search_template { + ($self: expr, $mv_est: expr, $cur_blk: expr, $mb_x: expr, $mb_y: expr, $sad_func: ident) => ({ + let mut best_dist = MAX_DIST; + let mut best_mv; + + let mut min_dist; + let mut min_idx; + + $self.reset(); + loop { + let mut cur_best_dist = best_dist; + for (dist, &point) in $self.dist.iter_mut().zip($self.point.iter()) { + if *dist == MAX_DIST { + *dist = $mv_est.$sad_func($cur_blk, $mb_x, $mb_y, point.from_pixels(), cur_best_dist); + cur_best_dist = cur_best_dist.min(*dist); + if *dist <= DIST_THRESH { + break; + } + } + } + min_dist = $self.dist[0]; + min_idx = 0; + for (i, &dist) in $self.dist.iter().enumerate().skip(1) { + if dist < min_dist { + min_dist = dist; + min_idx = i; + if dist <= DIST_THRESH { + break; + } + } + } + if min_dist <= DIST_THRESH || min_idx == 0 || best_dist == min_dist || $self.point[min_idx].x.abs() >= $mv_est.mv_range || $self.point[min_idx].y.abs() >= $mv_est.mv_range { + break; + } + best_dist = min_dist; + $self.update($self.steps[min_idx]); + } + best_dist = min_dist; + best_mv = $self.point[min_idx]; + if best_dist <= DIST_THRESH { + return (best_mv.from_pixels(), best_dist); + } + for &step in REFINEMENT.iter() { + let mv = best_mv + step; + let dist = $mv_est.$sad_func($cur_blk, $mb_x, $mb_y, mv.from_pixels(), MAX_DIST); + if best_dist > dist { + best_dist = dist; + best_mv = mv; + } + } + best_mv = best_mv.from_pixels(); + if best_dist <= DIST_THRESH { + return (best_mv, best_dist); + } + + // subpel refinement + $self.set_new_point(best_mv, best_dist); + loop { + let mut cur_best_dist = best_dist; + for (dist, &point) in $self.dist.iter_mut().zip($self.point.iter()) { + if *dist == MAX_DIST { + *dist = $mv_est.$sad_func($cur_blk, $mb_x, $mb_y, point, cur_best_dist); + cur_best_dist = cur_best_dist.min(*dist); + if *dist <= DIST_THRESH { + break; + } + } + } + min_dist = $self.dist[0]; + min_idx = 0; + for (i, &dist) in $self.dist.iter().enumerate().skip(1) { + if dist < min_dist { + min_dist = dist; + min_idx = i; + if dist <= DIST_THRESH { + break; + } + } + } + if min_dist <= DIST_THRESH || min_idx == 0 || best_dist == min_dist || $self.point[min_idx].x.abs() >= $mv_est.mv_range * 4 || $self.point[min_idx].y.abs() >= $mv_est.mv_range * 4 { + break; + } + best_dist = min_dist; + $self.update($self.steps[min_idx]); + } + best_dist = min_dist; + best_mv = $self.point[min_idx]; + if best_dist <= DIST_THRESH { + return (best_mv, best_dist); + } + for &step in REFINEMENT.iter() { + let mv = best_mv + step; + let dist = $mv_est.$sad_func($cur_blk, $mb_x, $mb_y, mv, MAX_DIST); + if best_dist > dist { + best_dist = dist; + best_mv = mv; + } + } + (best_mv, best_dist) + }) +} -- 2.30.2