Commit | Line | Data |
---|---|---|
3952bfd9 KS |
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 | } |