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