nihav-duck/src/codecs/vp6enc/models.rs

   1 use nihav_codec_support::codecs::ZIGZAG;
   2 use super::super::vp6data::*;
   3 use crate::codecs::vpenc::models::*;
   4 pub use crate::codecs::vpenc::models::PROB_BITS;
   5
   6 #[derive(Clone,Copy,Default)]
   7 pub struct VP56MVModel {
   8     pub nz_prob:        u8,
   9     pub sign_prob:      u8,
  10     pub raw_probs:      [u8; 8],
  11     pub tree_probs:     [u8; 7],
  12 }
  13
  14 #[derive(Clone,Copy,Default)]
  15 pub struct VP56MBTypeModel {
  16     pub probs:          [u8; 10],
  17 }
  18
  19 #[derive(Clone,Copy,Default)]
  20 pub struct VP56CoeffModel {
  21     pub dc_token_probs: [[[u8; 5]; 6]; 6],
  22     pub dc_value_probs: [u8; 11],
  23     pub ac_val_probs:   [[[u8; 11]; 6]; 3],
  24 }
  25
  26 #[derive(Clone)]
  27 pub struct VP6Models {
  28     pub scan_order:         [usize; 64],
  29     pub scan:               [usize; 64],
  30     pub zigzag:             [usize; 64],
  31     pub zero_run_probs:     [[u8; 14]; 2],
  32 }
  33
  34 const MAX_HUFF_ELEMS: usize = 12;
  35 #[derive(Clone,Copy,Default)]
  36 pub struct VP6Huff {
  37     pub codes:  [u16; MAX_HUFF_ELEMS],
  38     pub bits:   [u8; MAX_HUFF_ELEMS],
  39 }
  40
  41 #[derive(Clone,Copy,Default)]
  42 struct Node {
  43     weight:     u16,
  44     sym:        i8,
  45     ch0:        usize,
  46     ch1:        usize,
  47 }
  48
  49 fn prob2weight(a: u8, b: u8) -> u8 {
  50     let w = ((u16::from(a) * u16::from(b)) >> 8) as u8;
  51     if w == 0 {
  52         1
  53     } else {
  54         w
  55     }
  56 }
  57
  58 #[allow(clippy::identity_op)]
  59 impl VP6Huff {
  60     pub fn build_codes(&mut self, probs: &[u8; 11]) {
  61         let mut weights = [0u8; 12];
  62
  63         weights[11] = prob2weight( probs[0],  probs[ 1]);
  64         weights[ 0] = prob2weight( probs[0], !probs[ 1]);
  65         weights[ 1] = prob2weight(!probs[0],  probs[ 2]);
  66         let lvroot  = prob2weight(!probs[0], !probs[ 2]);
  67         let tworoot = prob2weight( lvroot,    probs[ 3]);
  68         let hlroot  = prob2weight( lvroot,   !probs[ 3]);
  69         weights[ 2] = prob2weight( tworoot,   probs[ 4]);
  70         let root34  = prob2weight( tworoot,  !probs[ 4]);
  71         weights[ 3] = prob2weight( root34,    probs[ 5]);
  72         weights[ 4] = prob2weight( root34,   !probs[ 5]);
  73         let c1root  = prob2weight( hlroot,    probs[ 6]);
  74         let c34root = prob2weight( hlroot,   !probs[ 6]);
  75         weights[ 5] = prob2weight( c1root,    probs[ 7]);
  76         weights[ 6] = prob2weight( c1root,   !probs[ 7]);
  77         let c3root  = prob2weight( c34root,   probs[ 8]);
  78         let c4root  = prob2weight( c34root,  !probs[ 8]);
  79         weights[ 7] = prob2weight( c3root,    probs[ 9]);
  80         weights[ 8] = prob2weight( c3root,   !probs[ 9]);
  81         weights[ 9] = prob2weight( c4root,    probs[10]);
  82         weights[10] = prob2weight( c4root,   !probs[10]);
  83
  84         self.build(&weights);
  85     }
  86     pub fn build_codes_zero_run(&mut self, probs: &[u8; 14]) {
  87         let mut weights = [0u8; 9];
  88
  89         let root   = prob2weight( probs[0],  probs[1]);
  90         weights[0] = prob2weight( root,      probs[2]);
  91         weights[1] = prob2weight( root,     !probs[2]);
  92
  93         let root   = prob2weight( probs[0], !probs[1]);
  94         weights[2] = prob2weight( root,      probs[3]);
  95         weights[3] = prob2weight( root,     !probs[3]);
  96
  97         let root   = prob2weight(!probs[0],  probs[4]);
  98         weights[8] = prob2weight(!probs[0], !probs[4]);
  99         let root1  = prob2weight( root,      probs[5]);
 100         let root2  = prob2weight( root,     !probs[5]);
 101         weights[4] = prob2weight( root1,     probs[6]);
 102         weights[5] = prob2weight( root1,    !probs[6]);
 103         weights[6] = prob2weight( root2,     probs[7]);
 104         weights[7] = prob2weight( root2,    !probs[7]);
 105
 106         self.build(&weights);
 107     }
 108     fn build(&mut self, weights: &[u8]) {
 109         let mut nodes = [Node::default(); MAX_HUFF_ELEMS * 2];
 110         let mut nlen = 0;
 111
 112         for w in weights.iter().rev() {
 113             let weight = u16::from(*w);
 114             let mut pos = nlen;
 115             for (i, node) in nodes[..nlen].iter().enumerate() {
 116                 if node.weight > weight {
 117                     pos = i;
 118                     break;
 119                 }
 120             }
 121             for j in (pos..nlen).rev() {
 122                 nodes[j + 1] = nodes[j];
 123             }
 124             nodes[pos].weight = weight;
 125             nodes[pos].sym    = (weights.len() - nlen - 1) as i8;
 126             nlen += 1;
 127         }
 128
 129         let mut low = 0;
 130         for _ in 0..nlen-1 {
 131             let nnode = Node {
 132                     weight: nodes[low + 0].weight + nodes[low + 1].weight,
 133                     sym:    -1,
 134                     ch0:    low + 0,
 135                     ch1:    low + 1,
 136                 };
 137             low += 2;
 138             let mut pos = low;
 139             while (pos < nlen) && (nodes[pos].weight < nnode.weight) {
 140                 pos += 1;
 141             }
 142             for j in (pos..nlen).rev() {
 143                 nodes[j + 1] = nodes[j];
 144             }
 145             nodes[pos] = nnode;
 146             nlen += 1;
 147         }
 148         self.get_codes(&nodes, nlen - 1, 0, 0);
 149         for i in nlen..self.codes.len() {
 150             self.codes[i]   = self.codes[0];
 151             self.bits[i]    = self.bits[0];
 152         }
 153     }
 154     fn get_codes(&mut self, nodes: &[Node], pos: usize, code: u16, len: u8) {
 155         if nodes[pos].sym >= 0 {
 156             self.codes[nodes[pos].sym as usize] = code;
 157             self.bits [nodes[pos].sym as usize] = len;
 158         } else {
 159             self.get_codes(nodes, nodes[pos].ch0, (code << 1) | 0, len + 1);
 160             self.get_codes(nodes, nodes[pos].ch1, (code << 1) | 1, len + 1);
 161         }
 162     }
 163 }
 164
 165 #[derive(Clone,Copy,Default)]
 166 pub struct VP6HuffModels {
 167     pub dc_token_tree:      [VP6Huff; 2],
 168     pub ac_token_tree:      [[[VP6Huff; 6]; 3]; 2],
 169     pub zero_run_tree:      [VP6Huff; 2],
 170 }
 171
 172 impl VP6Models {
 173     fn new() -> Self {
 174         Self {
 175             scan_order:         [0; 64],
 176             scan:               [0; 64],
 177             zigzag:             [0; 64],
 178             zero_run_probs:     [[0; 14]; 2],
 179         }
 180     }
 181 }
 182
 183 #[derive(Clone)]
 184 pub struct VP56Models {
 185     pub mv_models:          [VP56MVModel; 2],
 186     pub mbtype_models:      [[VP56MBTypeModel; 10]; 3],
 187     pub coeff_models:       [VP56CoeffModel; 2],
 188     pub prob_xmitted:       [[u8; 20]; 3],
 189     pub vp6models:          VP6Models,
 190     pub vp6huff:            VP6HuffModels,
 191 }
 192
 193 impl VP56Models {
 194     pub fn new() -> Self {
 195         Self {
 196             mv_models:      [VP56MVModel::default(); 2],
 197             mbtype_models:  [[VP56MBTypeModel::default(); 10]; 3],
 198             coeff_models:   [VP56CoeffModel::default(); 2],
 199             prob_xmitted:   [[0; 20]; 3],
 200             vp6models:      VP6Models::new(),
 201             vp6huff:        VP6HuffModels::default(),
 202         }
 203     }
 204     pub fn reset(&mut self, interlaced: bool) {
 205         for (i, mdl) in self.mv_models.iter_mut().enumerate() {
 206             mdl.nz_prob         = NZ_PROBS[i];
 207             mdl.sign_prob       = 128;
 208             mdl.raw_probs.copy_from_slice(&RAW_PROBS[i]);
 209             mdl.tree_probs.copy_from_slice(&TREE_PROBS[i]);
 210         }
 211
 212         for mdl in self.coeff_models.iter_mut() {
 213             mdl.dc_value_probs = [128; 11];
 214             mdl.ac_val_probs = [[[128; 11]; 6]; 3];
 215         }
 216         self.vp6models.zero_run_probs.copy_from_slice(&ZERO_RUN_PROBS);
 217         reset_scan(&mut self.vp6models, interlaced);
 218     }
 219     pub fn reset_mbtype_models(&mut self) {
 220         const DEFAULT_XMITTED_PROBS: [[u8; 20]; 3] = [
 221             [ 42,  69, 2, 1, 7, 1, 42, 44, 22, 6, 3, 1, 2, 0, 5, 1, 1, 0, 0, 0 ],
 222             [  8, 229, 1, 1, 8, 0,  0,  0,  0, 0, 2, 1, 1, 0, 0, 0, 1, 1, 0, 0 ],
 223             [ 35, 122, 1, 1, 6, 1, 34, 46,  0, 0, 2, 1, 1, 0, 1, 0, 1, 1, 0, 0 ]
 224         ];
 225         self.prob_xmitted.copy_from_slice(&DEFAULT_XMITTED_PROBS);
 226     }
 227 }
 228
 229 pub fn reset_scan(model: &mut VP6Models, interlaced: bool) {
 230     if !interlaced {
 231         model.scan_order.copy_from_slice(&VP6_DEFAULT_SCAN_ORDER);
 232     } else {
 233         model.scan_order.copy_from_slice(&VP6_INTERLACED_SCAN_ORDER);
 234     }
 235     for i in 0..64 { model.scan[i] = i; }
 236     model.zigzag.copy_from_slice(&ZIGZAG);
 237 }
 238
 239 #[derive(Clone,Copy,Default)]
 240 pub struct VP56MVModelStat {
 241     pub nz_prob:        ProbCounter,
 242     pub sign_prob:      ProbCounter,
 243     pub raw_probs:      [ProbCounter; 8],
 244     pub tree_probs:     [ProbCounter; 7],
 245 }
 246
 247 #[derive(Clone,Copy,Default)]
 248 pub struct VP56CoeffModelStat {
 249     pub dc_value_probs: [ProbCounter; 11],
 250     pub ac_val_probs:   [[[ProbCounter; 11]; 6]; 3],
 251 }
 252
 253 #[derive(Default)]
 254 pub struct VP6ModelsStat {
 255     pub zero_run_probs:     [[ProbCounter; 14]; 2],
 256 }
 257
 258 pub struct VP56ModelsStat {
 259     pub mv_models:          [VP56MVModelStat; 2],
 260     pub mbtype_models:      [[[usize; 10]; 10]; 3],
 261     pub coeff_models:       [VP56CoeffModelStat; 2],
 262     pub vp6models:          VP6ModelsStat,
 263 }
 264
 265 impl VP56ModelsStat {
 266     pub fn new() -> Self {
 267         Self {
 268             mv_models:      [VP56MVModelStat::default(); 2],
 269             mbtype_models:  [[[0; 10]; 10]; 3],
 270             coeff_models:   [VP56CoeffModelStat::default(); 2],
 271             vp6models:      VP6ModelsStat::default(),
 272         }
 273     }
 274     pub fn reset(&mut self) {
 275         self.mv_models      = [VP56MVModelStat::default(); 2];
 276         self.mbtype_models  = [[[0; 10]; 10]; 3];
 277         self.coeff_models   = [VP56CoeffModelStat::default(); 2];
 278         self.vp6models      = VP6ModelsStat::default();
 279     }
 280     pub fn generate(&self, dst: &mut VP56Models, is_intra: bool) {
 281         if !is_intra {
 282             for (dmv, smv) in dst.mv_models.iter_mut().zip(self.mv_models.iter()) {
 283                 dmv.nz_prob = smv.nz_prob.to_prob_worthy(dmv.nz_prob);
 284                 dmv.sign_prob = smv.sign_prob.to_prob_worthy(dmv.sign_prob);
 285                 for (dp, sp) in dmv.raw_probs.iter_mut().zip(smv.raw_probs.iter()) {
 286                     *dp = sp.to_prob_worthy(*dp);
 287                 }
 288                 for (dp, sp) in dmv.tree_probs.iter_mut().zip(smv.tree_probs.iter()) {
 289                     *dp = sp.to_prob_worthy(*dp);
 290                 }
 291             }
 292             for (xmit, mdl) in dst.prob_xmitted.iter_mut().zip(self.mbtype_models.iter()) {
 293                 Self::generate_prob_xmitted(xmit, mdl);
 294             }
 295         }
 296         for (dmv, smv) in dst.coeff_models.iter_mut().zip(self.coeff_models.iter()) {
 297             for (dp, sp) in dmv.dc_value_probs.iter_mut().zip(smv.dc_value_probs.iter()) {
 298                 *dp = sp.to_prob_worthy(*dp);
 299             }
 300             for (dp, sp) in dmv.ac_val_probs.iter_mut().zip(smv.ac_val_probs.iter()) {
 301                 for (dp, sp) in dp.iter_mut().zip(sp.iter()) {
 302                     for (dp, sp) in dp.iter_mut().zip(sp.iter()) {
 303                         *dp = sp.to_prob_worthy(*dp);
 304                     }
 305                 }
 306             }
 307         }
 308         for (dp, sp) in dst.vp6models.zero_run_probs.iter_mut().zip(self.vp6models.zero_run_probs.iter()) {
 309             for (dp, sp) in dp.iter_mut().zip(sp.iter()) {
 310                 *dp = sp.to_prob_worthy(*dp);
 311             }
 312         }
 313     }
 314     /*
 315         VPMBType::InterNoMV     => 0,
 316         VPMBType::Intra         => 1,
 317         VPMBType::InterMV       => 2,
 318         VPMBType::InterNearest  => 3,
 319         VPMBType::InterNear     => 4,
 320         VPMBType::GoldenNoMV    => 5,
 321         VPMBType::GoldenMV      => 6,
 322         VPMBType::InterFourMV   => 7,
 323         VPMBType::GoldenNearest => 8,
 324         VPMBType::GoldenNear    => 9,
 325     */
 326     fn generate_prob_xmitted(probs: &mut [u8; 20], mbtype: &[[usize; 10]; 10]) {
 327         let mut sums = [0; 20];
 328         let mut total = 0;
 329         for (last, row) in mbtype.iter().enumerate() {
 330             for (cur, &count) in row.iter().enumerate() {
 331                 if last == cur {
 332                     sums[cur * 2 + 1] = count;
 333                 } else {
 334                     sums[cur * 2] += count;
 335                 }
 336                 total += count;
 337             }
 338         }
 339         if total != 0 {
 340             let mut sum = 0;
 341             for (dprob, &sprob) in probs.iter_mut().zip(sums.iter()) {
 342                 if sprob != 0 {
 343                     *dprob = ((sprob * 256 + total - 1) / total).min(255) as u8;
 344                     sum += u16::from(*dprob);
 345                 } else {
 346                     *dprob = 0;
 347                 }
 348             }
 349             while sum > 256 {
 350                 for prob in probs.iter_mut() {
 351                     if *prob > 1 {
 352                         *prob -= 1;
 353                         sum -= 1;
 354                         if sum == 256 {
 355                             break;
 356                         }
 357                     }
 358                 }
 359             }
 360         } else {
 361             *probs = [0; 20];
 362         }
 363     }
 364 }