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 impl VP6Huff {
  59     pub fn build_codes(&mut self, probs: &[u8; 11]) {
  60         let mut weights = [0u8; 12];
  61
  62         weights[11] = prob2weight( probs[0],  probs[ 1]);
  63         weights[ 0] = prob2weight( probs[0], !probs[ 1]);
  64         weights[ 1] = prob2weight(!probs[0],  probs[ 2]);
  65         let lvroot  = prob2weight(!probs[0], !probs[ 2]);
  66         let tworoot = prob2weight( lvroot,    probs[ 3]);
  67         let hlroot  = prob2weight( lvroot,   !probs[ 3]);
  68         weights[ 2] = prob2weight( tworoot,   probs[ 4]);
  69         let root34  = prob2weight( tworoot,  !probs[ 4]);
  70         weights[ 3] = prob2weight( root34,    probs[ 5]);
  71         weights[ 4] = prob2weight( root34,   !probs[ 5]);
  72         let c1root  = prob2weight( hlroot,    probs[ 6]);
  73         let c34root = prob2weight( hlroot,   !probs[ 6]);
  74         weights[ 5] = prob2weight( c1root,    probs[ 7]);
  75         weights[ 6] = prob2weight( c1root,   !probs[ 7]);
  76         let c3root  = prob2weight( c34root,   probs[ 8]);
  77         let c4root  = prob2weight( c34root,  !probs[ 8]);
  78         weights[ 7] = prob2weight( c3root,    probs[ 9]);
  79         weights[ 8] = prob2weight( c3root,   !probs[ 9]);
  80         weights[ 9] = prob2weight( c4root,    probs[10]);
  81         weights[10] = prob2weight( c4root,   !probs[10]);
  82
  83         self.build(&weights);
  84     }
  85     pub fn build_codes_zero_run(&mut self, probs: &[u8; 14]) {
  86         let mut weights = [0u8; 9];
  87
  88         let root   = prob2weight( probs[0],  probs[1]);
  89         weights[0] = prob2weight( root,      probs[2]);
  90         weights[1] = prob2weight( root,     !probs[2]);
  91
  92         let root   = prob2weight( probs[0], !probs[1]);
  93         weights[2] = prob2weight( root,      probs[3]);
  94         weights[3] = prob2weight( root,     !probs[3]);
  95
  96         let root   = prob2weight(!probs[0],  probs[4]);
  97         weights[8] = prob2weight(!probs[0], !probs[4]);
  98         let root1  = prob2weight( root,      probs[5]);
  99         let root2  = prob2weight( root,     !probs[5]);
 100         weights[4] = prob2weight( root1,     probs[6]);
 101         weights[5] = prob2weight( root1,    !probs[6]);
 102         weights[6] = prob2weight( root2,     probs[7]);
 103         weights[7] = prob2weight( root2,    !probs[7]);
 104
 105         self.build(&weights);
 106     }
 107     fn build(&mut self, weights: &[u8]) {
 108         let mut nodes = [Node::default(); MAX_HUFF_ELEMS * 2];
 109         let mut nlen = 0;
 110
 111         for w in weights.iter().rev() {
 112             let weight = u16::from(*w);
 113             let mut pos = nlen;
 114             for i in 0..nlen {
 115                 if nodes[i].weight > weight {
 116                     pos = i;
 117                     break;
 118                 }
 119             }
 120             for j in (pos..nlen).rev() {
 121                 nodes[j + 1] = nodes[j];
 122             }
 123             nodes[pos].weight = weight;
 124             nodes[pos].sym    = (weights.len() - nlen - 1) as i8;
 125             nlen += 1;
 126         }
 127
 128         let mut low = 0;
 129         for _ in 0..nlen-1 {
 130             let nnode = Node {
 131                     weight: nodes[low + 0].weight + nodes[low + 1].weight,
 132                     sym:    -1,
 133                     ch0:    low + 0,
 134                     ch1:    low + 1,
 135                 };
 136             low += 2;
 137             let mut pos = low;
 138             while (pos < nlen) && (nodes[pos].weight < nnode.weight) {
 139                 pos += 1;
 140             }
 141             for j in (pos..nlen).rev() {
 142                 nodes[j + 1] = nodes[j];
 143             }
 144             nodes[pos] = nnode;
 145             nlen += 1;
 146         }
 147         self.get_codes(&nodes, nlen - 1, 0, 0);
 148         for i in nlen..self.codes.len() {
 149             self.codes[i]   = self.codes[0];
 150             self.bits[i]    = self.bits[0];
 151         }
 152     }
 153     fn get_codes(&mut self, nodes: &[Node], pos: usize, code: u16, len: u8) {
 154         if nodes[pos].sym >= 0 {
 155             self.codes[nodes[pos].sym as usize] = code;
 156             self.bits [nodes[pos].sym as usize] = len;
 157         } else {
 158             self.get_codes(nodes, nodes[pos].ch0, (code << 1) | 0, len + 1);
 159             self.get_codes(nodes, nodes[pos].ch1, (code << 1) | 1, len + 1);
 160         }
 161     }
 162 }
 163
 164 #[derive(Clone,Copy,Default)]
 165 pub struct VP6HuffModels {
 166     pub dc_token_tree:      [VP6Huff; 2],
 167     pub ac_token_tree:      [[[VP6Huff; 6]; 3]; 2],
 168     pub zero_run_tree:      [VP6Huff; 2],
 169 }
 170
 171 impl VP6Models {
 172     fn new() -> Self {
 173         Self {
 174             scan_order:         [0; 64],
 175             scan:               [0; 64],
 176             zigzag:             [0; 64],
 177             zero_run_probs:     [[0; 14]; 2],
 178         }
 179     }
 180 }
 181
 182 #[derive(Clone)]
 183 pub struct VP56Models {
 184     pub mv_models:          [VP56MVModel; 2],
 185     pub mbtype_models:      [[VP56MBTypeModel; 10]; 3],
 186     pub coeff_models:       [VP56CoeffModel; 2],
 187     pub prob_xmitted:       [[u8; 20]; 3],
 188     pub vp6models:          VP6Models,
 189     pub vp6huff:            VP6HuffModels,
 190 }
 191
 192 impl VP56Models {
 193     pub fn new() -> Self {
 194         Self {
 195             mv_models:      [VP56MVModel::default(); 2],
 196             mbtype_models:  [[VP56MBTypeModel::default(); 10]; 3],
 197             coeff_models:   [VP56CoeffModel::default(); 2],
 198             prob_xmitted:   [[0; 20]; 3],
 199             vp6models:      VP6Models::new(),
 200             vp6huff:        VP6HuffModels::default(),
 201         }
 202     }
 203     pub fn reset(&mut self, interlaced: bool) {
 204         for (i, mdl) in self.mv_models.iter_mut().enumerate() {
 205             mdl.nz_prob         = NZ_PROBS[i];
 206             mdl.sign_prob       = 128;
 207             mdl.raw_probs.copy_from_slice(&RAW_PROBS[i]);
 208             mdl.tree_probs.copy_from_slice(&TREE_PROBS[i]);
 209         }
 210
 211         for mdl in self.coeff_models.iter_mut() {
 212             mdl.dc_value_probs = [128; 11];
 213             mdl.ac_val_probs = [[[128; 11]; 6]; 3];
 214         }
 215         self.vp6models.zero_run_probs.copy_from_slice(&ZERO_RUN_PROBS);
 216         reset_scan(&mut self.vp6models, interlaced);
 217     }
 218     pub fn reset_mbtype_models(&mut self) {
 219         const DEFAULT_XMITTED_PROBS: [[u8; 20]; 3] = [
 220             [ 42,  69, 2, 1, 7, 1, 42, 44, 22, 6, 3, 1, 2, 0, 5, 1, 1, 0, 0, 0 ],
 221             [  8, 229, 1, 1, 8, 0,  0,  0,  0, 0, 2, 1, 1, 0, 0, 0, 1, 1, 0, 0 ],
 222             [ 35, 122, 1, 1, 6, 1, 34, 46,  0, 0, 2, 1, 1, 0, 1, 0, 1, 1, 0, 0 ]
 223         ];
 224         self.prob_xmitted.copy_from_slice(&DEFAULT_XMITTED_PROBS);
 225     }
 226 }
 227
 228 pub fn reset_scan(model: &mut VP6Models, interlaced: bool) {
 229     if !interlaced {
 230         model.scan_order.copy_from_slice(&VP6_DEFAULT_SCAN_ORDER);
 231     } else {
 232         model.scan_order.copy_from_slice(&VP6_INTERLACED_SCAN_ORDER);
 233     }
 234     for i in 0..64 { model.scan[i] = i; }
 235     model.zigzag.copy_from_slice(&ZIGZAG);
 236 }
 237
 238 #[derive(Clone,Copy,Default)]
 239 pub struct VP56MVModelStat {
 240     pub nz_prob:        ProbCounter,
 241     pub sign_prob:      ProbCounter,
 242     pub raw_probs:      [ProbCounter; 8],
 243     pub tree_probs:     [ProbCounter; 7],
 244 }
 245
 246 #[derive(Clone,Copy,Default)]
 247 pub struct VP56CoeffModelStat {
 248     pub dc_token_probs: [[[ProbCounter; 5]; 6]; 6],
 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 }