| 1 | use nihav_core::frame::FrameType; |
| 2 | use nihav_codec_support::codecs::MV; |
| 3 | use super::types::*; |
| 4 | use super::super::rv40data::*; |
| 5 | |
| 6 | pub struct BitsEstimator { |
| 7 | ftype: FrameType, |
| 8 | pred_mbt: MBType, |
| 9 | cur_mbt: MBType, |
| 10 | } |
| 11 | |
| 12 | impl BitsEstimator { |
| 13 | pub fn new() -> Self { |
| 14 | Self { |
| 15 | ftype: FrameType::I, |
| 16 | pred_mbt: MBType::Invalid, |
| 17 | cur_mbt: MBType::Invalid, |
| 18 | } |
| 19 | } |
| 20 | pub fn set_frame_type(&mut self, ftype: FrameType) { |
| 21 | self.ftype = ftype; |
| 22 | } |
| 23 | pub fn set_quant(&mut self, _q: usize) { |
| 24 | } |
| 25 | pub fn set_pred_mb_type(&mut self, most_prob_type: MBType) { |
| 26 | self.pred_mbt = most_prob_type; |
| 27 | } |
| 28 | pub fn set_mb_type(&mut self, mbt: MBType) { |
| 29 | self.cur_mbt = mbt; |
| 30 | } |
| 31 | pub fn estimate_mb_hdr(&self, mvs: &[MV]) -> u32 { |
| 32 | if self.ftype == FrameType::I { |
| 33 | return 1; |
| 34 | } |
| 35 | let hdr_cw_bits = if self.ftype == FrameType::P { |
| 36 | RV40_PTYPE_BITS[self.pred_mbt.to_code()][self.cur_mbt.to_code()] |
| 37 | } else { |
| 38 | RV40_BTYPE_BITS[self.pred_mbt.to_code()][self.cur_mbt.to_code()] |
| 39 | }; |
| 40 | let mv_bits = mvs.iter().fold(0u32, |acc, &mv| acc + Self::mv_cost(mv)); |
| 41 | u32::from(hdr_cw_bits) + mv_bits |
| 42 | } |
| 43 | fn block_no_to_type(&self, blk_no: usize) -> usize { |
| 44 | match blk_no { |
| 45 | 0..=15 => { |
| 46 | match self.cur_mbt { |
| 47 | MBType::Intra16 | MBType::P16x16Mix => 2, |
| 48 | MBType::Intra => 1, |
| 49 | _ => 0, |
| 50 | } |
| 51 | }, |
| 52 | 24 => 3, |
| 53 | _ if self.cur_mbt.is_intra() => 4, |
| 54 | _ => 5, |
| 55 | } |
| 56 | } |
| 57 | pub fn block_bits(&self, blk: &Block, blk_no: usize) -> u32 { |
| 58 | let btype = self.block_no_to_type(blk_no); |
| 59 | |
| 60 | const EXPECTED_BLOCK_BITS: [[u8; 17]; 6] = [ |
| 61 | [ 0, 7, 12, 17, 22, 26, 31, 35, 39, 45, 51, 56, 61, 66, 85, 103, 117], |
| 62 | [ 0, 7, 13, 19, 26, 30, 36, 43, 49, 57, 65, 74, 87, 99, 115, 131, 147], |
| 63 | [ 0, 7, 14, 20, 25, 30, 35, 40, 45, 50, 56, 62, 69, 76, 84, 93, 113], |
| 64 | [ 2, 9, 13, 20, 25, 29, 33, 38, 43, 48, 54, 62, 71, 82, 98, 116, 141], |
| 65 | [ 0, 5, 12, 18, 24, 30, 35, 42, 48, 53, 62, 69, 78, 87, 97, 106, 121], |
| 66 | [ 0, 6, 12, 17, 22, 27, 33, 40, 47, 53, 60, 66, 73, 80, 85, 85, 103] |
| 67 | ]; |
| 68 | EXPECTED_BLOCK_BITS[btype][blk.count_nz()].into() |
| 69 | } |
| 70 | pub fn mv_cost(mv: MV) -> u32 { |
| 71 | let xval = mv.x.abs() * 2 + 1; |
| 72 | let yval = mv.y.abs() * 2 + 1; |
| 73 | (15 - xval.leading_zeros()) * 2 + (15 - yval.leading_zeros()) * 2 + 2 |
| 74 | } |
| 75 | pub fn decide_set(hist: &[usize; 17]) -> usize { |
| 76 | let max_val = hist[16]; |
| 77 | let threshold = max_val - max_val / 4; |
| 78 | if hist[3] > threshold { |
| 79 | 2 |
| 80 | } else if hist[6] > threshold { |
| 81 | 1 |
| 82 | } else { |
| 83 | 0 |
| 84 | } |
| 85 | } |
| 86 | } |