codecs/vp6enc/mod.rs

   1 use nihav_core::codecs::*;
   2 use nihav_core::io::byteio::*;
   3 use super::vp6data::*;
   4 use super::vpcommon::*;
   5
   6 mod coder;
   7 use coder::*;
   8 mod dsp;
   9 use dsp::MVSearchMode;
  10 mod huff;
  11 use huff::*;
  12 mod mb;
  13 use mb::*;
  14 mod models;
  15 use models::*;
  16 mod ratectl;
  17 use ratectl::*;
  18 mod rdo;
  19
  20 const VERSION_VP61: u8 = VERSION_VP60 + 1;
  21
  22 enum VP6Writer<'a, 'b> {
  23     BoolCoder(BoolEncoder<'a, 'b>),
  24     Huffman(HuffEncoder<'a, 'b>),
  25 }
  26
  27 #[derive(Default)]
  28 pub struct VP56DCPred {
  29     dc_y:       Vec<i16>,
  30     dc_u:       Vec<i16>,
  31     dc_v:       Vec<i16>,
  32     ldc_y:      [i16; 2],
  33     ldc_u:      i16,
  34     ldc_v:      i16,
  35     ref_y:      Vec<u8>,
  36     ref_c:      Vec<u8>,
  37     ref_left:   u8,
  38     y_idx:      usize,
  39     c_idx:      usize,
  40
  41     last_dc:    [[i16; 4]; 3],
  42 }
  43
  44 const INVALID_REF: u8 = 42;
  45
  46 impl VP56DCPred {
  47     fn new() -> Self { Self::default() }
  48     fn resize(&mut self, mb_w: usize) {
  49         self.dc_y.resize(mb_w * 2 + 2, 0);
  50         self.dc_u.resize(mb_w     + 2, 0);
  51         self.dc_v.resize(mb_w     + 2, 0);
  52         self.ref_y.resize(mb_w * 2 + 2, INVALID_REF);
  53         self.ref_c.resize(mb_w     + 2, INVALID_REF);
  54         self.ref_c[0] = 0;
  55     }
  56     fn reset(&mut self) {
  57         self.update_row();
  58         for el in self.ref_y.iter_mut().skip(1) { *el = INVALID_REF; }
  59         for el in self.ref_c.iter_mut().skip(1) { *el = INVALID_REF; }
  60
  61         self.last_dc = [[0; 4]; 3];
  62         self.last_dc[0][1] = 0x80;
  63         self.last_dc[0][2] = 0x80;
  64     }
  65     fn update_row(&mut self) {
  66         self.y_idx = 1;
  67         self.c_idx = 1;
  68         self.ldc_y = [0; 2];
  69         self.ldc_u = 0;
  70         self.ldc_v = 0;
  71         self.ref_left = INVALID_REF;
  72     }
  73     fn next_mb(&mut self) {
  74         self.y_idx += 2;
  75         self.c_idx += 1;
  76     }
  77     fn predict_dc(&mut self, mb_type: VPMBType, blk_no: usize, coeffs: &mut [i16; 64], fwd: bool) {
  78         let is_luma = blk_no < 4;
  79         let (plane, dcs) = match blk_no {
  80                 4 => (1, &mut self.dc_u),
  81                 5 => (2, &mut self.dc_v),
  82                 _ => (0, &mut self.dc_y),
  83              };
  84         let (dc_ref, dc_idx) = if is_luma {
  85                 (&mut self.ref_y, self.y_idx + (blk_no & 1))
  86             } else {
  87                 (&mut self.ref_c, self.c_idx)
  88             };
  89         let ref_id = mb_type.get_ref_id();
  90         let mut dc_pred = 0;
  91         let mut count = 0;
  92         let has_left_blk = is_luma && ((blk_no & 1) == 1);
  93         if has_left_blk || self.ref_left == ref_id {
  94             dc_pred += match blk_no {
  95                     0 | 1 => self.ldc_y[0],
  96                     2 | 3 => self.ldc_y[1],
  97                     4     => self.ldc_u,
  98                     _     => self.ldc_v,
  99                 };
 100             count += 1;
 101         }
 102         if dc_ref[dc_idx] == ref_id {
 103             dc_pred += dcs[dc_idx];
 104             count += 1;
 105         }
 106         if count == 0 {
 107             dc_pred = self.last_dc[ref_id as usize][plane];
 108         } else if count == 2 {
 109             dc_pred /= 2;
 110         }
 111         if !fwd {
 112             coeffs[0] += dc_pred;
 113         }
 114
 115         let dc = coeffs[0];
 116         if blk_no != 4 { // update top block reference only for the second chroma component
 117             dc_ref[dc_idx] = ref_id;
 118         }
 119         match blk_no {
 120             0 | 1 => {
 121                 self.ldc_y[0] = dc;
 122             },
 123             2 | 3 => {
 124                 self.ldc_y[1] = dc;
 125             },
 126             4 => {
 127                 self.ldc_u = dc;
 128             },
 129             _ => {
 130                 self.ldc_v = dc;
 131                 self.ref_left = ref_id;
 132             },
 133         };
 134         dcs[dc_idx] = dc;
 135
 136         self.last_dc[ref_id as usize][plane] = dc;
 137         if fwd {
 138             coeffs[0] -= dc_pred;
 139         }
 140     }
 141 }
 142
 143 trait ZeroBlock {
 144     fn no_dc(&self) -> bool;
 145     fn no_ac(&self) -> bool;
 146 }
 147
 148 impl ZeroBlock for [i16; 64] {
 149     fn no_dc(&self) -> bool {
 150         self[0] == 0
 151     }
 152     fn no_ac(&self) -> bool {
 153         for &el in self[1..].iter() {
 154             if el != 0 {
 155                 return false;
 156             }
 157         }
 158         true
 159     }
 160 }
 161
 162 struct VP6Encoder {
 163     stream:     Option<NAStreamRef>,
 164     pkt:        Option<NAPacket>,
 165     key_int:    u8,
 166     frmcount:   u8,
 167     mb_w:       usize,
 168     mb_h:       usize,
 169     dc_pred:    VP56DCPred,
 170     top_ctx:    [Vec<bool>; 4],
 171     mc_buf:     NAVideoBufferRef<u8>,
 172     fenc:       FrameEncoder,
 173     ratectl:    RateControl,
 174
 175     flipped:    bool,
 176     huffman:    bool,
 177
 178     version:    u8,
 179     profile:    u8,
 180
 181     models:     VP56Models,
 182     stats:      VP56ModelsStat,
 183     pmodels:    VP56Models,
 184
 185     last_frame: NABufferType,
 186     gold_frame: NABufferType,
 187     last_gold:  bool,
 188     me_mode:    MVSearchMode,
 189     me_range:   i16,
 190
 191     force_q:    Option<usize>,
 192 }
 193
 194 impl VP6Encoder {
 195     fn new(flipped: bool) -> Self {
 196         let vt = alloc_video_buffer(NAVideoInfo::new(24, 24, false, VP_YUVA420_FORMAT), 4).unwrap();
 197         let mc_buf = vt.get_vbuf().unwrap();
 198         Self {
 199             stream:     None,
 200             pkt:        None,
 201             key_int:    10,
 202             frmcount:   0,
 203             mb_w:       0,
 204             mb_h:       0,
 205             dc_pred:    VP56DCPred::new(),
 206             top_ctx:    [Vec::new(), Vec::new(), Vec::new(), Vec::new()],
 207             fenc:       FrameEncoder::new(),
 208             ratectl:    RateControl::new(),
 209             mc_buf,
 210
 211             flipped,
 212             huffman:    false,
 213
 214             version:    VERSION_VP60,
 215             profile:    VP6_SIMPLE_PROFILE,
 216
 217             models:     VP56Models::new(),
 218             pmodels:    VP56Models::new(),
 219             stats:      VP56ModelsStat::new(),
 220
 221             last_frame: NABufferType::None,
 222             gold_frame: NABufferType::None,
 223             last_gold:  false,
 224             me_mode:    MVSearchMode::default(),
 225             me_range:   16,
 226
 227             force_q:    None,
 228         }
 229     }
 230     fn decide_encoding(&mut self) -> bool {
 231         false
 232     }
 233     fn estimate_blocks(&mut self, is_intra: bool) {
 234         for top_ctx in self.top_ctx.iter_mut() {
 235             for el in top_ctx.iter_mut() {
 236                 *el = false;
 237             }
 238         }
 239         let mut last_mbt = VPMBType::InterNoMV;
 240         let mut mb_idx = 0;
 241         for _mb_y in 0..self.mb_h {
 242             let mut left_dc = [false; 4];
 243             for mb_x in 0..self.mb_w {
 244                 let mb_type = self.fenc.mb_types[mb_idx];
 245                 if !is_intra {
 246                     estimate_mb_type(mb_type, last_mbt, ((self.fenc.num_mv[mb_idx] + 1) % 3) as usize, &mut self.stats);
 247                     last_mbt = mb_type;
 248                     match mb_type {
 249                         VPMBType::InterMV | VPMBType::GoldenMV => {
 250                             estimate_mv(self.fenc.coded_mv[mb_idx][3], &mut self.stats);
 251                         },
 252                         VPMBType::InterFourMV => {
 253                             for (&sub_type, &mv) in self.fenc.fmv_sub[mb_idx].iter().zip(self.fenc.coded_mv[mb_idx].iter()) {
 254                                 if sub_type == VPMBType::InterMV {
 255                                     estimate_mv(mv, &mut self.stats);
 256                                 }
 257                             }
 258                         },
 259                         _ => {},
 260                     };
 261                 }
 262                 let mb = self.fenc.get_mb(mb_idx);
 263                 for i in 0..4 {
 264                     let cur_idx = mb_x * 2 + (i & 1);
 265                     let mut dc_mode = 0;
 266                     if self.top_ctx[0][cur_idx] {
 267                         dc_mode += 1;
 268                     }
 269                     if left_dc[i >> 1] {
 270                         dc_mode += 1;
 271                     }
 272                     self.top_ctx[0][cur_idx] = mb.coeffs[i][0] != 0;
 273                     left_dc[i >> 1] = mb.coeffs[i][0] != 0;
 274                     estimate_block(&mb.coeffs[i], dc_mode, &mut self.stats.coeff_models[0], &mut self.stats.vp6models, &self.models.vp6models.zigzag);
 275                 }
 276
 277                 let mut dc_mode = 0;
 278                 if self.top_ctx[1][mb_x] {
 279                     dc_mode += 1;
 280                 }
 281                 if left_dc[2] {
 282                     dc_mode += 1;
 283                 }
 284                 self.top_ctx[1][mb_x] = mb.coeffs[4][0] != 0;
 285                 left_dc[2] = mb.coeffs[4][0] != 0;
 286                 estimate_block(&mb.coeffs[4], dc_mode, &mut self.stats.coeff_models[1], &mut self.stats.vp6models, &self.models.vp6models.zigzag);
 287
 288                 let mut dc_mode = 0;
 289                 if self.top_ctx[2][mb_x] {
 290                     dc_mode += 1;
 291                 }
 292                 if left_dc[3] {
 293                     dc_mode += 1;
 294                 }
 295                 self.top_ctx[2][mb_x] = mb.coeffs[5][0] != 0;
 296                 left_dc[3] = mb.coeffs[5][0] != 0;
 297                 estimate_block(&mb.coeffs[5], dc_mode, &mut self.stats.coeff_models[1], &mut self.stats.vp6models, &self.models.vp6models.zigzag);
 298
 299                 mb_idx += 1;
 300             }
 301         }
 302     }
 303     fn prepare_huff_models(&mut self) {
 304         for i in 0..2 {
 305             self.models.vp6huff.dc_token_tree[i].build_codes(&self.models.coeff_models[i].dc_value_probs);
 306         }
 307         for i in 0..2 {
 308             for mode in 0..3 {
 309                 for band in 0..6 {
 310                     self.models.vp6huff.ac_token_tree[i][mode][band].build_codes(&self.models.coeff_models[i].ac_val_probs[mode][band]);
 311                 }
 312             }
 313         }
 314         for i in 0..2 {
 315             self.models.vp6huff.zero_run_tree[i].build_codes_zero_run(&self.models.vp6models.zero_run_probs[i]);
 316         }
 317     }
 318     fn determine_coeff_runs_luma(&self, hstate: &mut HuffState, mb_pos: usize, blk: usize) {
 319         let mb = self.fenc.get_mb(mb_pos);
 320
 321         if !hstate.dc_zr_coded[0] {
 322             if mb.coeffs[blk].no_dc() {
 323                 hstate.dc_zero_run[0] = 0;
 324                 let mut blk_no = (blk + 1) & 3;
 325                 let mut mb_no = mb_pos + ((blk + 1) >> 2);
 326                 let mut cmb = mb;
 327                 let mut last_mb_no = mb_pos;
 328                 while (hstate.dc_zero_run[0] < MAX_EOB_RUN) && (mb_no < self.mb_w * self.mb_h) {
 329                     if mb_no != last_mb_no {
 330                         cmb = self.fenc.get_mb(mb_no);
 331                         last_mb_no = mb_no;
 332                     }
 333                     if !cmb.coeffs[blk_no].no_dc() {
 334                         break;
 335                     }
 336                     hstate.dc_zero_run[0] += 1;
 337                     blk_no += 1;
 338                     if blk_no == 4 {
 339                         blk_no = 0;
 340                         mb_no += 1;
 341                     }
 342                 }
 343             }
 344         }
 345         if !hstate.ac_zr_coded[0] {
 346             if mb.coeffs[blk].no_ac() {
 347                 hstate.ac_zero_run[0] = 0;
 348                 let mut blk_no = (blk + 1) & 3;
 349                 let mut mb_no = mb_pos + ((blk + 1) >> 2);
 350                 let mut cmb = mb;
 351                 let mut last_mb_no = mb_pos;
 352                 while (hstate.ac_zero_run[0] < MAX_EOB_RUN) && (mb_no < self.mb_w * self.mb_h) {
 353                     if mb_no != last_mb_no {
 354                         cmb = self.fenc.get_mb(mb_no);
 355                         last_mb_no = mb_no;
 356                     }
 357                     if !cmb.coeffs[blk_no].no_ac() {
 358                         break;
 359                     }
 360                     hstate.ac_zero_run[0] += 1;
 361                     blk_no += 1;
 362                     if blk_no == 4 {
 363                         blk_no = 0;
 364                         mb_no += 1;
 365                     }
 366                 }
 367             }
 368         }
 369     }
 370     fn determine_coeff_runs_chroma(&self, hstate: &mut HuffState, mb_pos: usize, plane: usize) {
 371         let mb = self.fenc.get_mb(mb_pos);
 372         let blk = plane + 3;
 373
 374         if !hstate.dc_zr_coded[1] {
 375             if mb.coeffs[blk].no_dc() {
 376                 hstate.dc_zero_run[1] = 0;
 377                 let mut blk_no = if blk == 4 { 5 } else { 4 };
 378                 let mut mb_no = mb_pos + if blk == 4 { 0 } else { 1 };
 379                 while (hstate.dc_zero_run[1] < MAX_EOB_RUN) && (mb_no < self.mb_w * self.mb_h) {
 380                     let mb = self.fenc.get_mb(mb_no);
 381                     if !mb.coeffs[blk_no].no_dc() {
 382                         break;
 383                     }
 384                     hstate.dc_zero_run[1] += 1;
 385                     blk_no += 1;
 386                     if blk_no == 6 {
 387                         blk_no = 4;
 388                         mb_no += 1;
 389                     }
 390                 }
 391             }
 392         }
 393         if !hstate.ac_zr_coded[1] {
 394             if mb.coeffs[blk].no_ac() {
 395                 hstate.ac_zero_run[1] = 0;
 396                 let mut blk_no = if blk == 4 { 5 } else { 4 };
 397                 let mut mb_no = mb_pos + if blk == 4 { 0 } else { 1 };
 398                 while (hstate.ac_zero_run[1] < MAX_EOB_RUN) && (mb_no < self.mb_w * self.mb_h) {
 399                     let mb = self.fenc.get_mb(mb_no);
 400                     if !mb.coeffs[blk_no].no_ac() {
 401                         break;
 402                     }
 403                     hstate.ac_zero_run[1] += 1;
 404                     blk_no += 1;
 405                     if blk_no == 6 {
 406                         blk_no = 4;
 407                         mb_no += 1;
 408                     }
 409                 }
 410             }
 411         }
 412     }
 413     fn encode_intra(&mut self, bw: &mut ByteWriter, quant: usize) -> EncoderResult<bool> {
 414         self.models.reset(false);
 415         self.models.reset_mbtype_models();
 416         self.stats.reset();
 417
 418         self.pmodels.reset(false);
 419         self.pmodels.reset_mbtype_models();
 420
 421         let multistream = self.huffman || self.version != VERSION_VP60;
 422
 423         self.fenc.prepare_intra_blocks();
 424         self.fenc.apply_dc_prediction(&mut self.dc_pred);
 425         self.estimate_blocks(true);
 426         self.stats.generate(&mut self.models, true);
 427
 428         // header
 429         bw.write_byte(((quant as u8) << 1) | (multistream as u8))?;
 430         bw.write_byte((self.version << 3) | (self.profile << 1))?;
 431         bw.write_u16be(0)?; // part 2 offset placeholder
 432
 433         let mut bc = BoolEncoder::new(bw);
 434
 435         bc.put_bits(self.mb_h as u32, 8)?;
 436         bc.put_bits(self.mb_w as u32, 8)?;
 437         bc.put_bits(self.mb_h as u32, 8)?; // display MB height
 438         bc.put_bits(self.mb_w as u32, 8)?; // display MB width
 439         bc.put_bits(0, 2)?; // scaline mode
 440         // todo other advanced profile bits
 441         bc.put_bits(self.huffman as u32, 1)?; // Huffman mode
 442
 443         encode_coeff_models(&mut bc, &mut self.models, &self.pmodels, true, false)?;
 444         self.pmodels = self.models.clone();
 445
 446         if multistream || (self.profile == VP6_SIMPLE_PROFILE) {
 447             bc.flush()?;
 448
 449             // patch coefficient offset
 450             let offset = bw.tell();
 451             if offset >= 65535 {
 452                 return Err(EncoderError::Bug);
 453             }
 454             bw.seek(SeekFrom::Start(2))?;
 455             bw.write_u16be(offset as u16)?;
 456             bw.seek(SeekFrom::End(0))?;
 457
 458             bc = BoolEncoder::new(bw);
 459         }
 460         let writer = if !self.huffman {
 461                 VP6Writer::BoolCoder(bc)
 462             } else {
 463                 VP6Writer::Huffman(HuffEncoder::new(bw))
 464             };
 465         self.encode_coeffs(writer)?;
 466         Ok(true)
 467     }
 468     fn encode_inter(&mut self, bw: &mut ByteWriter, quant: usize) -> EncoderResult<bool> {
 469         self.stats.reset();
 470
 471         let multistream = self.huffman || self.version != VERSION_VP60;
 472         let loop_filter = false;
 473
 474         self.fenc.prepare_intra_blocks();
 475         self.fenc.prepare_inter_blocks(false);
 476         if !self.last_gold {
 477             self.fenc.prepare_inter_blocks(true);
 478         }
 479         let lambda = if self.force_q.is_some() { 1.0 } else { self.ratectl.lambda };
 480         self.fenc.select_inter_blocks(self.last_frame.get_vbuf().unwrap(), self.mc_buf.clone(), !self.last_gold, lambda);
 481         // todo implement forced intra
 482         let (_force_intra, golden_frame) = self.fenc.decide_frame_type();
 483         self.fenc.apply_dc_prediction(&mut self.dc_pred);
 484         self.fenc.predict_mvs();
 485         self.estimate_blocks(false);
 486
 487         self.stats.generate(&mut self.models, false);
 488
 489         // header
 490         bw.write_byte(0x80 | ((quant as u8) << 1) | (multistream as u8))?;
 491         bw.write_u16be(0)?; // part 2 offset placeholder
 492
 493         let mut bc = BoolEncoder::new(bw);
 494
 495         bc.put_bits(golden_frame as u32, 1)?; // refresh golden frame
 496         if self.profile == VP6_ADVANCED_PROFILE {
 497             bc.put_bits(loop_filter as u32, 1)?; // use loop filter
 498             if loop_filter {
 499                 bc.put_bits(0, 1)?; // loop filter selector
 500             }
 501             if self.version == VERSION_VP62 {
 502                 bc.put_bits(0, 1)?; // auto select PM
 503             }
 504         }
 505         // todo other advanced profile bits
 506         bc.put_bits(self.huffman as u32, 1)?;
 507
 508         encode_mode_prob_models(&mut bc, &mut self.models, &self.pmodels, &self.stats.mbtype_models)?;
 509         encode_mv_models(&mut bc, &self.models.mv_models, &self.pmodels.mv_models)?;
 510         encode_coeff_models(&mut bc, &mut self.models, &self.pmodels, false, false)?;
 511         self.pmodels = self.models.clone();
 512
 513         let mut last_mbt = VPMBType::InterNoMV;
 514         for mb_idx in 0..self.mb_w * self.mb_h {
 515             let mb_type = self.fenc.mb_types[mb_idx];
 516             encode_mb_type(&mut bc, self.fenc.mb_types[mb_idx], last_mbt, ((self.fenc.num_mv[mb_idx] + 1) % 3) as usize, &self.models)?;
 517             last_mbt = mb_type;
 518             match mb_type {
 519                 VPMBType::InterMV | VPMBType::GoldenMV => {
 520                     encode_mv(&mut bc, self.fenc.coded_mv[mb_idx][3], &self.models)?;
 521                 },
 522                 VPMBType::InterFourMV => {
 523                     for &sub_type in self.fenc.fmv_sub[mb_idx].iter() {
 524                         let id = match sub_type {
 525                                 VPMBType::InterNoMV     => 0,
 526                                 VPMBType::InterMV       => 1,
 527                                 VPMBType::InterNearest  => 2,
 528                                 VPMBType::InterNear     => 3,
 529                                 _ => unreachable!(),
 530                             };
 531                         bc.put_bits(id, 2)?;
 532                     }
 533                     for (&sub_type, &mv) in self.fenc.fmv_sub[mb_idx].iter().zip(self.fenc.coded_mv[mb_idx].iter()) {
 534                         if sub_type == VPMBType::InterMV {
 535                             encode_mv(&mut bc, mv, &self.models)?;
 536                         }
 537                     }
 538                 },
 539                 _ => {},
 540             };
 541         }
 542
 543         if multistream || (self.profile == VP6_SIMPLE_PROFILE) {
 544             bc.flush()?;
 545
 546             // patch coefficient offset
 547             let offset = bw.tell();
 548             if offset >= 65535 {
 549                 return Err(EncoderError::Bug);
 550             }
 551             bw.seek(SeekFrom::Start(1))?;
 552             bw.write_u16be(offset as u16)?;
 553             bw.seek(SeekFrom::End(0))?;
 554
 555             bc = BoolEncoder::new(bw);
 556         }
 557         let writer = if !self.huffman {
 558                 VP6Writer::BoolCoder(bc)
 559             } else {
 560                 VP6Writer::Huffman(HuffEncoder::new(bw))
 561             };
 562         self.encode_coeffs(writer)?;
 563         Ok(golden_frame)
 564     }
 565     fn encode_coeffs(&mut self, mut writer: VP6Writer) -> EncoderResult<()> {
 566         if self.huffman {
 567             self.prepare_huff_models();
 568         }
 569
 570         let mut hstate = HuffState::new();
 571         for top_ctx in self.top_ctx.iter_mut() {
 572             for el in top_ctx.iter_mut() {
 573                 *el = false;
 574             }
 575         }
 576         let mut mb_pos = 0;
 577         for _mb_y in 0..self.mb_h {
 578             let mut left_dc = [false; 4];
 579             for mb_x in 0..self.mb_w {
 580                 let mb = self.fenc.get_mb(mb_pos);
 581                 for i in 0..4 {
 582                     let cur_idx = mb_x * 2 + (i & 1);
 583                     let mut dc_mode = 0;
 584                     if self.top_ctx[0][cur_idx] {
 585                         dc_mode += 1;
 586                     }
 587                     if left_dc[i >> 1] {
 588                         dc_mode += 1;
 589                     }
 590                     self.top_ctx[0][cur_idx] = mb.coeffs[i][0] != 0;
 591                     left_dc[i >> 1] = mb.coeffs[i][0] != 0;
 592                     if self.huffman {
 593                         self.determine_coeff_runs_luma(&mut hstate, mb_pos, i);
 594                     }
 595                     match writer {
 596                         VP6Writer::BoolCoder(ref mut bc) => encode_block(bc, &mb.coeffs[i], dc_mode, &self.models.coeff_models[0], &self.models.vp6models)?,
 597                         VP6Writer::Huffman(ref mut huff) => encode_block_huff(huff, &self.models.vp6models.zigzag, &mb.coeffs[i], 0, &mut hstate, &self.models.vp6huff)?,
 598                     };
 599                 }
 600
 601                 for plane in 1..3 {
 602                     let mut dc_mode = 0;
 603                     if self.top_ctx[plane][mb_x] {
 604                         dc_mode += 1;
 605                     }
 606                     if left_dc[plane + 1] {
 607                         dc_mode += 1;
 608                     }
 609                     self.top_ctx[plane][mb_x] = mb.coeffs[plane + 3][0] != 0;
 610                     left_dc[plane + 1] = mb.coeffs[plane + 3][0] != 0;
 611                     if self.huffman {
 612                         self.determine_coeff_runs_chroma(&mut hstate, mb_pos, plane);
 613                     }
 614                     match writer {
 615                         VP6Writer::BoolCoder(ref mut bc) => encode_block(bc, &mb.coeffs[plane + 3], dc_mode, &self.models.coeff_models[1], &self.models.vp6models)?,
 616                         VP6Writer::Huffman(ref mut huff) => encode_block_huff(huff, &self.models.vp6models.zigzag, &mb.coeffs[plane + 3], 1, &mut hstate, &self.models.vp6huff)?,
 617                     };
 618                 }
 619
 620                 mb_pos += 1;
 621             }
 622         }
 623
 624         match writer {
 625             VP6Writer::BoolCoder(bc) => bc.flush()?,
 626             VP6Writer::Huffman(huff) => huff.flush()?,
 627         };
 628
 629         Ok(())
 630     }
 631 }
 632
 633 impl NAEncoder for VP6Encoder {
 634     fn negotiate_format(&self, encinfo: &EncodeParameters) -> EncoderResult<EncodeParameters> {
 635         match encinfo.format {
 636             NACodecTypeInfo::None => {
 637                 let mut ofmt = EncodeParameters::default();
 638                 ofmt.format = NACodecTypeInfo::Video(NAVideoInfo::new(0, 0, self.flipped, YUV420_FORMAT));
 639                 Ok(ofmt)
 640             },
 641             NACodecTypeInfo::Audio(_) => Err(EncoderError::FormatError),
 642             NACodecTypeInfo::Video(vinfo) => {
 643                 let outinfo = NAVideoInfo::new((vinfo.width + 3) & !3, (vinfo.height + 3) & !3, self.flipped, YUV420_FORMAT);
 644                 let mut ofmt = *encinfo;
 645                 ofmt.format = NACodecTypeInfo::Video(outinfo);
 646                 Ok(ofmt)
 647             }
 648         }
 649     }
 650     fn init(&mut self, stream_id: u32, encinfo: EncodeParameters) -> EncoderResult<NAStreamRef> {
 651         match encinfo.format {
 652             NACodecTypeInfo::None => Err(EncoderError::FormatError),
 653             NACodecTypeInfo::Audio(_) => Err(EncoderError::FormatError),
 654             NACodecTypeInfo::Video(vinfo) => {
 655                 if vinfo.format != YUV420_FORMAT {
 656                     return Err(EncoderError::FormatError);
 657                 }
 658                 if ((vinfo.width | vinfo.height) & 15) != 0 {
 659                     return Err(EncoderError::FormatError);
 660                 }
 661                 if (vinfo.width | vinfo.height) >= (1 << 12) {
 662                     return Err(EncoderError::FormatError);
 663                 }
 664
 665                 let out_info = NAVideoInfo::new(vinfo.width, vinfo.height, self.flipped, vinfo.format);
 666                 let info = NACodecInfo::new(if self.flipped { "vp6" } else { "vp6f" }, NACodecTypeInfo::Video(out_info), None);
 667                 let mut stream = NAStream::new(StreamType::Video, stream_id, info, encinfo.tb_num, encinfo.tb_den, 0);
 668                 stream.set_num(stream_id as usize);
 669                 let stream = stream.into_ref();
 670
 671                 self.last_frame = alloc_video_buffer(out_info, 4)?;
 672                 self.gold_frame = alloc_video_buffer(out_info, 4)?;
 673
 674                 self.stream = Some(stream.clone());
 675
 676                 self.mb_w = (vinfo.width  + 15) >> 4;
 677                 self.mb_h = (vinfo.height + 15) >> 4;
 678                 self.fenc.resize(self.mb_w, self.mb_h);
 679                 self.ratectl.init(self.mb_w, self.mb_h, encinfo.bitrate, encinfo.tb_num, encinfo.tb_den);
 680
 681                 self.dc_pred.resize(self.mb_w);
 682                 self.top_ctx = [vec![false; self.mb_w * 2], vec![false; self.mb_w], vec![false; self.mb_w], vec![false; self.mb_w * 2]];
 683
 684                 self.version = VERSION_VP60;
 685                 self.profile = VP6_SIMPLE_PROFILE;
 686
 687                 Ok(stream)
 688             },
 689         }
 690     }
 691     fn encode(&mut self, frm: &NAFrame) -> EncoderResult<()> {
 692         let buf = frm.get_buffer();
 693         if let Some(ref vbuf) = buf.get_vbuf() {
 694             let mut dbuf = Vec::with_capacity(4);
 695             let mut gw   = GrowableMemoryWriter::new_write(&mut dbuf);
 696             let mut bw   = ByteWriter::new(&mut gw);
 697
 698             // todo integrate with rate control
 699             let is_intra = (self.frmcount == 0) || self.decide_encoding();
 700             let quant = if let Some(q) = self.force_q {
 701                     q
 702                 } else {
 703                     self.ratectl.guess_quant(is_intra, self.huffman)
 704                 };
 705
 706             self.fenc.read_mbs(vbuf);
 707             self.fenc.set_quant(quant);
 708             self.fenc.me_mode = self.me_mode;
 709             self.fenc.me_range = self.me_range;
 710             let golden_frame = if is_intra {
 711                     self.encode_intra(&mut bw, quant)?
 712                 } else {
 713                     self.fenc.estimate_mvs(self.last_frame.get_vbuf().unwrap(), self.mc_buf.clone(), false);
 714                     if !self.last_gold {
 715                         self.fenc.estimate_mvs(self.gold_frame.get_vbuf().unwrap(), self.mc_buf.clone(), true);
 716                     }
 717                     self.encode_inter(&mut bw, quant)?
 718                 };
 719             self.fenc.reconstruct_frame(&mut self.dc_pred, self.last_frame.get_vbuf().unwrap());
 720             self.last_gold = golden_frame;
 721             if golden_frame {
 722                 let mut dfrm = self.gold_frame.get_vbuf().unwrap();
 723                 let src = self.last_frame.get_vbuf().unwrap();
 724
 725                 let dst = dfrm.get_data_mut().unwrap();
 726                 dst.copy_from_slice(src.get_data());
 727             }
 728
 729             if self.force_q.is_none() {
 730                 self.ratectl.update(dbuf.len() * 8);
 731             }
 732
 733             self.pkt = Some(NAPacket::new(self.stream.clone().unwrap(), frm.ts, is_intra, dbuf));
 734             if self.key_int > 0 {
 735                 self.frmcount += 1;
 736             }
 737             if self.frmcount == self.key_int {
 738                 self.frmcount = 0;
 739             }
 740             Ok(())
 741         } else {
 742             Err(EncoderError::InvalidParameters)
 743         }
 744     }
 745     fn get_packet(&mut self) -> EncoderResult<Option<NAPacket>> {
 746         let mut npkt = None;
 747         std::mem::swap(&mut self.pkt, &mut npkt);
 748         Ok(npkt)
 749     }
 750     fn flush(&mut self) -> EncoderResult<()> {
 751         self.frmcount = 0;
 752         Ok(())
 753     }
 754 }
 755
 756 const HUFFMAN_OPTION: &str = "huffman";
 757 const QUANT_OPTION: &str = "quant";
 758 const VERSION_OPTION: &str = "version";
 759 const MV_SEARCH_OPTION: &str = "mv_mode";
 760 const MV_RANGE_OPTION: &str = "mv_range";
 761
 762 const ENCODER_OPTS: &[NAOptionDefinition] = &[
 763     NAOptionDefinition {
 764         name: KEYFRAME_OPTION, description: KEYFRAME_OPTION_DESC,
 765         opt_type: NAOptionDefinitionType::Int(Some(0), Some(128)) },
 766     NAOptionDefinition {
 767         name: HUFFMAN_OPTION, description: "use Huffman encoding",
 768         opt_type: NAOptionDefinitionType::Bool },
 769     NAOptionDefinition {
 770         name: VERSION_OPTION, description: "codec minor version",
 771         opt_type: NAOptionDefinitionType::String(Some(&["vp60", "vp61", "vp62"])) },
 772     NAOptionDefinition {
 773         name: QUANT_OPTION, description: "force fixed quantiser for encoding",
 774         opt_type: NAOptionDefinitionType::Int(Some(-1), Some(63)) },
 775     NAOptionDefinition {
 776         name: MV_SEARCH_OPTION, description: "motion search mode",
 777         opt_type: NAOptionDefinitionType::String(Some(&["full", "dia", "hex"])) },
 778     NAOptionDefinition {
 779         name: MV_RANGE_OPTION, description: "motion search range (in pixels)",
 780         opt_type: NAOptionDefinitionType::Int(Some(0), Some(30)) },
 781 ];
 782
 783 impl NAOptionHandler for VP6Encoder {
 784     fn get_supported_options(&self) -> &[NAOptionDefinition] { ENCODER_OPTS }
 785     fn set_options(&mut self, options: &[NAOption]) {
 786         for option in options.iter() {
 787             for opt_def in ENCODER_OPTS.iter() {
 788                 if opt_def.check(option).is_ok() {
 789                     match option.name {
 790                         KEYFRAME_OPTION => {
 791                             if let NAValue::Int(intval) = option.value {
 792                                 self.key_int = intval as u8;
 793                             }
 794                         },
 795                         HUFFMAN_OPTION => {
 796                             if let NAValue::Bool(bval) = option.value {
 797                                 self.huffman = bval;
 798                             }
 799                         },
 800                         VERSION_OPTION => {
 801                             if let NAValue::String(ref string) = option.value {
 802                                 self.version = match string.as_str() {
 803                                         "vp60" => VERSION_VP60,
 804                                         "vp61" => VERSION_VP61,
 805                                         "vp62" => VERSION_VP62,
 806                                         _ => unreachable!(),
 807                                     };
 808                             }
 809                         },
 810                         QUANT_OPTION => {
 811                             if let NAValue::Int(intval) = option.value {
 812                                 self.force_q = if intval < 0 { None } else { Some(intval as usize) };
 813                             }
 814                         },
 815                         MV_SEARCH_OPTION => {
 816                             if let NAValue::String(ref string) = option.value {
 817                                 if let Ok(mv_mode) = string.parse::<MVSearchMode>() {
 818                                     self.me_mode = mv_mode;
 819                                 }
 820                             }
 821                         },
 822                         MV_RANGE_OPTION => {
 823                             if let NAValue::Int(intval) = option.value {
 824                                 self.me_range = intval as i16;
 825                             }
 826                         },
 827                         _ => {},
 828                     };
 829                 }
 830             }
 831         }
 832     }
 833     fn query_option_value(&self, name: &str) -> Option<NAValue> {
 834         match name {
 835             KEYFRAME_OPTION => Some(NAValue::Int(i64::from(self.key_int))),
 836             HUFFMAN_OPTION => Some(NAValue::Bool(self.huffman)),
 837             VERSION_OPTION => {
 838                     let ver = match self.version {
 839                             VERSION_VP60 => "vp60",
 840                             VERSION_VP61 => "vp61",
 841                             VERSION_VP62 => "vp62",
 842                             _ => unreachable!(),
 843                         };
 844                     Some(NAValue::String(ver.to_string()))
 845                 },
 846             QUANT_OPTION => if let Some(q) = self.force_q {
 847                     Some(NAValue::Int(q as i64))
 848                 } else {
 849                     Some(NAValue::Int(-1))
 850                 },
 851             MV_SEARCH_OPTION => Some(NAValue::String(self.me_mode.to_string())),
 852             MV_RANGE_OPTION => Some(NAValue::Int(i64::from(self.me_range))),
 853             _ => None,
 854         }
 855     }
 856 }
 857
 858 pub fn get_encoder() -> Box<dyn NAEncoder + Send> {
 859     Box::new(VP6Encoder::new(true))
 860 }
 861
 862 pub fn get_encoder_flv() -> Box<dyn NAEncoder + Send> {
 863     Box::new(VP6Encoder::new(false))
 864 }
 865
 866 #[cfg(test)]
 867 mod test {
 868     use nihav_core::codecs::*;
 869     use nihav_core::demuxers::*;
 870     use nihav_core::muxers::*;
 871     use crate::*;
 872     use nihav_commonfmt::*;
 873     use nihav_codec_support::test::enc_video::*;
 874
 875     fn encode_test(out_name: &'static str, enc_options: &[NAOption], hash: &[u32; 4]) {
 876         let mut dmx_reg = RegisteredDemuxers::new();
 877         generic_register_all_demuxers(&mut dmx_reg);
 878         let mut dec_reg = RegisteredDecoders::new();
 879         duck_register_all_decoders(&mut dec_reg);
 880         let mut mux_reg = RegisteredMuxers::new();
 881         generic_register_all_muxers(&mut mux_reg);
 882         let mut enc_reg = RegisteredEncoders::new();
 883         duck_register_all_encoders(&mut enc_reg);
 884
 885         // sample: https://samples.mplayerhq.hu/V-codecs/VP4/ot171_vp40.avi
 886         let dec_config = DecoderTestParams {
 887                 demuxer:        "avi",
 888                 in_name:        "assets/Duck/ot171_vp40.avi",
 889                 stream_type:    StreamType::Video,
 890                 limit:          Some(1),
 891                 dmx_reg, dec_reg,
 892             };
 893         let enc_config = EncoderTestParams {
 894                 muxer:          "avi",
 895                 enc_name:       "vp6",
 896                 out_name,
 897                 mux_reg, enc_reg,
 898             };
 899         let dst_vinfo = NAVideoInfo {
 900                 width:   0,
 901                 height:  0,
 902                 format:  YUV420_FORMAT,
 903                 flipped: true,
 904                 bits:    12,
 905             };
 906         let enc_params = EncodeParameters {
 907                 format:  NACodecTypeInfo::Video(dst_vinfo),
 908                 quality: 0,
 909                 bitrate: 25000,
 910                 tb_num:  0,
 911                 tb_den:  0,
 912                 flags:   0,
 913             };
 914         //test_encoding_to_file(&dec_config, &enc_config, enc_params, enc_options);
 915         test_encoding_md5(&dec_config, &enc_config, enc_params, enc_options,
 916                           hash);
 917     }
 918     #[test]
 919     fn test_vp6_encoder_bc() {
 920         let enc_options = &[
 921                 NAOption { name: super::QUANT_OPTION, value: NAValue::Int(42) },
 922             ];
 923         encode_test("vp6-bool.avi", enc_options, &[0xb57f49e5, 0x6b48accd, 0xc28fadb3, 0xc89a30d2]);
 924     }
 925     #[test]
 926     fn test_vp6_encoder_rc() {
 927         let enc_options = &[
 928             ];
 929         encode_test("vp6-rc.avi", enc_options, &[0x790baca9, 0x663eafcf, 0x36d1bed8, 0xddf882de]);
 930     }
 931     #[test]
 932     fn test_vp6_encoder_huff() {
 933         let enc_options = &[
 934                 NAOption { name: super::HUFFMAN_OPTION, value: NAValue::Bool(true) },
 935                 NAOption { name: super::QUANT_OPTION, value: NAValue::Int(42) },
 936             ];
 937         encode_test("vp6-huff.avi", enc_options, &[0x6e9bb23d, 0xde296d92, 0x4c225bae, 0x3651e31f]);
 938     }
 939 }