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