]>
Commit | Line | Data |
---|---|---|
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 | } |