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