9f87054ac17c37f50109a5a2430f5a61a8e1000b
[nihav.git] / nihav-duck / src / 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 }