]>
Commit | Line | Data |
---|---|---|
217de10b KS |
1 | use nihav_core::codecs::*; |
2 | use nihav_core::io::bitwriter::*; | |
3 | ||
4 | use super::binkviddata::*; | |
5 | ||
6 | mod bundle; | |
7 | use bundle::*; | |
8 | mod dsp; | |
9 | use dsp::*; | |
10 | mod mc; | |
11 | use mc::*; | |
12 | mod rc; | |
13 | use rc::*; | |
14 | ||
15 | const THRESHOLD: u32 = 64; | |
08e0a8bf | 16 | const MAX_DIST: u32 = u32::MAX; |
217de10b KS |
17 | |
18 | #[derive(Clone,Copy,Default,PartialEq)] | |
19 | enum DoublingMode { | |
20 | #[default] | |
21 | None, | |
22 | Height2X, | |
23 | HeightIlace, | |
24 | Scale2X, | |
25 | Width2X, | |
26 | Width2XIlace | |
27 | } | |
28 | ||
29 | impl std::string::ToString for DoublingMode { | |
30 | fn to_string(&self) -> String { | |
31 | match *self { | |
32 | DoublingMode::None => "none".to_string(), | |
33 | DoublingMode::Height2X => "height2x".to_string(), | |
34 | DoublingMode::HeightIlace => "height_il".to_string(), | |
35 | DoublingMode::Scale2X => "scale2x".to_string(), | |
36 | DoublingMode::Width2X => "width2x".to_string(), | |
37 | DoublingMode::Width2XIlace => "width2x_il".to_string(), | |
38 | } | |
39 | } | |
40 | } | |
41 | ||
217de10b KS |
42 | #[derive(Clone,Copy,Debug,PartialEq)] |
43 | enum BlockMode { | |
44 | Skip, | |
45 | Scaled, | |
46 | Run, | |
47 | Intra, | |
48 | Residue, | |
49 | Inter, | |
50 | Fill, | |
51 | Pattern, | |
52 | Motion, | |
53 | Raw, | |
54 | } | |
55 | ||
56 | const BLOCK_MODE_NAMES: &[&str] = &[ | |
57 | "skip", "run", "intra", "residue", "inter", "fill", "pattern", "motion", "raw", "scaled" | |
58 | ]; | |
59 | ||
60 | impl From<BlockMode> for usize { | |
61 | fn from(bmode: BlockMode) -> usize { | |
62 | match bmode { | |
63 | BlockMode::Skip => 0, | |
64 | BlockMode::Run => 1, | |
65 | BlockMode::Intra => 2, | |
66 | BlockMode::Residue => 3, | |
67 | BlockMode::Inter => 4, | |
68 | BlockMode::Fill => 5, | |
69 | BlockMode::Pattern => 6, | |
70 | BlockMode::Motion => 7, | |
71 | BlockMode::Raw => 8, | |
72 | BlockMode::Scaled => 9, | |
73 | } | |
74 | } | |
75 | } | |
76 | ||
77 | impl std::str::FromStr for BlockMode { | |
78 | type Err = (); | |
79 | fn from_str(s: &str) -> Result<Self, Self::Err> { | |
80 | match s { | |
81 | "skip" => Ok(BlockMode::Skip), | |
82 | "run" => Ok(BlockMode::Run), | |
83 | "intra" => Ok(BlockMode::Intra), | |
84 | "residue" => Ok(BlockMode::Residue), | |
85 | "inter" => Ok(BlockMode::Inter), | |
86 | "fill" => Ok(BlockMode::Fill), | |
87 | "pattern" => Ok(BlockMode::Pattern), | |
88 | "motion" => Ok(BlockMode::Motion), | |
89 | "raw" => Ok(BlockMode::Raw), | |
90 | "scaled" => Ok(BlockMode::Scaled), | |
91 | _ => Err(()), | |
92 | } | |
93 | } | |
94 | } | |
95 | ||
96 | impl std::string::ToString for BlockMode { | |
97 | fn to_string(&self) -> String { BLOCK_MODE_NAMES[usize::from(*self)].to_string() } | |
98 | } | |
99 | ||
100 | impl BlockMode { | |
101 | fn to_code_b(self) -> u8 { | |
102 | match self { | |
103 | BlockMode::Skip => 0, | |
104 | BlockMode::Run => 1, | |
105 | BlockMode::Intra => 2, | |
106 | BlockMode::Residue => 3, | |
107 | BlockMode::Inter => 4, | |
108 | BlockMode::Fill => 5, | |
109 | BlockMode::Pattern => 6, | |
110 | BlockMode::Motion => 7, | |
111 | BlockMode::Raw => 8, | |
112 | _ => unreachable!(), | |
113 | } | |
114 | } | |
115 | } | |
116 | ||
117 | fn close_enough(base: u8, thr: u8, pix: u8) -> bool { | |
118 | pix >= base.saturating_sub(thr) && pix <= base.saturating_add(thr) | |
119 | } | |
120 | ||
121 | fn write_run(tokens: &mut BlockTokens, pos: usize, len: usize, is_run: bool, clrs: &[u8; 64], is_b: bool) { | |
122 | tokens.other.push((is_run as u16, 1)); | |
123 | if is_b { | |
124 | tokens.other.push(((len - 1) as u16, BINKB_RUN_BITS[pos])); | |
125 | } else { | |
126 | unimplemented!(); | |
127 | } | |
128 | if is_run { | |
129 | tokens.colors.push(clrs[pos]); | |
130 | } else { | |
131 | tokens.colors.extend_from_slice(&clrs[pos..][..len]); | |
132 | } | |
133 | } | |
134 | ||
135 | fn find_run_pattern(cur_blk: &[u8; 64], tokens: &mut BlockTokens, tmp_tok: &mut BlockTokens, is_b: bool, thr: u8) -> u32 { | |
136 | tokens.clear(); | |
137 | let mut best_diff = MAX_DIST; | |
138 | let mut clrs = [0; 64]; | |
139 | for (id, pattern) in BINK_PATTERNS.iter().enumerate() { | |
140 | tmp_tok.clear(); | |
141 | tmp_tok.other.push((id as u16, 4)); | |
142 | ||
143 | for (dst, &idx) in clrs.iter_mut().zip(pattern.iter()) { | |
144 | *dst = cur_blk[usize::from(idx)]; | |
145 | } | |
146 | ||
147 | let mut cur_diff = 0; | |
148 | let mut last_val = 0; | |
149 | let mut last_pos = 0; | |
150 | let mut len = 0; | |
151 | let mut is_run = false; | |
152 | ||
153 | for (i, &val) in clrs.iter().enumerate() { | |
154 | if len > 0 && close_enough(last_val, thr, val) { | |
155 | if is_run || len == 1 { | |
156 | is_run = true; | |
157 | len += 1; | |
158 | } else { | |
159 | write_run(tmp_tok, last_pos, len - 1, is_run, &clrs, is_b); | |
160 | last_pos = i - 1; | |
161 | last_val = clrs[last_pos]; | |
162 | len = 2; | |
163 | } | |
164 | cur_diff += pix_dist(last_val, val); | |
165 | } else { | |
166 | if len > 0 { | |
167 | write_run(tmp_tok, last_pos, len, is_run, &clrs, is_b); | |
168 | } | |
169 | last_pos = i; | |
170 | last_val = val; | |
171 | len = 1; | |
172 | is_run = false; | |
173 | } | |
174 | } | |
175 | match len { | |
176 | 0 => {}, | |
177 | 1 => tmp_tok.colors.push(last_val), | |
178 | _ => write_run(tmp_tok, last_pos, len, is_run, &clrs, is_b), | |
179 | }; | |
180 | ||
181 | if cur_diff < best_diff { | |
182 | best_diff = cur_diff; | |
183 | std::mem::swap(tokens, tmp_tok); | |
184 | } | |
185 | } | |
186 | best_diff | |
187 | } | |
188 | ||
189 | fn find_pattern(cur_blk: &[u8; 64], tokens: &mut BlockTokens) -> u32 { | |
190 | let sum = cur_blk.iter().fold(0u16, | |
191 | |acc, &a| acc + u16::from(a)); | |
192 | let avg = ((sum + 32) >> 6) as u8; | |
193 | let mut sum_ba = 0u16; | |
194 | let mut sum_aa = 0u16; | |
195 | let mut cnt_ba = 0u16; | |
196 | let mut cnt_aa = 0u16; | |
197 | for &pix in cur_blk.iter() { | |
198 | if pix < avg { | |
199 | sum_ba += u16::from(pix); | |
200 | cnt_ba += 1; | |
201 | } else { | |
202 | sum_aa += u16::from(pix); | |
203 | cnt_aa += 1; | |
204 | } | |
205 | } | |
206 | if cnt_ba > 0 { // not flat | |
207 | let clr_ba = ((sum_ba + cnt_ba / 2) / cnt_ba) as u8; | |
208 | let clr_aa = ((sum_aa + cnt_aa / 2) / cnt_aa) as u8; | |
209 | tokens.clear(); | |
210 | tokens.colors.push(clr_ba); | |
211 | tokens.colors.push(clr_aa); | |
212 | let mut diff = 0; | |
213 | for row in cur_blk.chunks_exact(8) { | |
214 | let mut pat = 0; | |
215 | let mut mask = 1; | |
216 | for &p in row.iter() { | |
217 | if p < avg { | |
218 | diff += pix_dist(p, clr_ba); | |
219 | } else { | |
220 | pat |= mask; | |
221 | diff += pix_dist(p, clr_aa); | |
222 | } | |
223 | mask <<= 1; | |
224 | } | |
225 | tokens.pattern.push(pat); | |
226 | } | |
227 | diff | |
228 | } else { | |
229 | MAX_DIST | |
230 | } | |
231 | } | |
232 | ||
233 | struct BinkEncoder { | |
234 | stream: Option<NAStreamRef>, | |
235 | pkt: Option<NAPacket>, | |
236 | cur_frm: NAVideoBufferRef<u8>, | |
237 | last_frm: NAVideoBufferRef<u8>, | |
238 | scale_mode: DoublingMode, | |
239 | version: char, | |
240 | bundles: Bundles, | |
241 | nframes: u64, | |
242 | frame_no: u64, | |
243 | ||
244 | bst_tokens: BlockTokens, | |
245 | tmp_tokens: BlockTokens, | |
246 | tmp_tok2: BlockTokens, | |
247 | ||
248 | dsp: DSP, | |
249 | rc: RateControl, | |
250 | ||
251 | forbidden: [bool; 12], | |
252 | ||
253 | print_stats: bool, | |
254 | blk_stats: [usize; 12], | |
255 | tot_size: usize, | |
256 | iq_stats: [usize; 16], | |
257 | pq_stats: [usize; 16], | |
258 | } | |
259 | ||
260 | impl BinkEncoder { | |
261 | fn new() -> Self { | |
262 | let frm = alloc_video_buffer(NAVideoInfo::new(4, 4, false, YUV420_FORMAT), 0).unwrap(); | |
263 | let cur_frm = frm.get_vbuf().unwrap(); | |
264 | let last_frm = cur_frm.clone(); | |
265 | Self { | |
266 | stream: None, | |
267 | pkt: None, | |
268 | cur_frm, last_frm, | |
269 | scale_mode: DoublingMode::default(), | |
270 | version: 'b', | |
271 | bundles: Bundles::default(), | |
272 | nframes: 0, | |
273 | frame_no: 0, | |
274 | ||
275 | bst_tokens: BlockTokens::new(), | |
276 | tmp_tokens: BlockTokens::new(), | |
277 | tmp_tok2: BlockTokens::new(), | |
278 | ||
279 | dsp: DSP::new(), | |
280 | rc: RateControl::new(), | |
281 | ||
282 | forbidden: [false; 12], | |
283 | ||
284 | print_stats: false, | |
285 | blk_stats: [0; 12], | |
286 | tot_size: 0, | |
287 | iq_stats: [0; 16], | |
288 | pq_stats: [0; 16], | |
289 | } | |
290 | } | |
291 | fn encode_frame(&mut self, bw: &mut BitWriter, srcbuf: &NAVideoBuffer<u8>) -> EncoderResult<bool> { | |
292 | let frm = NASimpleVideoFrame::from_video_buf(&mut self.cur_frm).unwrap(); | |
293 | let src = srcbuf.get_data(); | |
294 | let last = self.last_frm.get_data(); | |
295 | ||
296 | let is_b = self.version == 'b'; | |
297 | let first_frame = self.frame_no == 1; | |
298 | let pat_run_thr = self.rc.pattern_run_threshold(); | |
299 | ||
300 | let mut cur_blk = [0; 64]; | |
301 | let mut mv_blk = [0; 64]; | |
302 | let mut is_intra = true; | |
303 | let mut cur_forbidden = self.forbidden; | |
304 | self.rc.modify_forbidden_btypes(&mut cur_forbidden); | |
305 | self.dsp.set_quant_ranges(self.rc.get_quant_ranges()); | |
306 | for plane in 0..3 { | |
307 | let loff = self.last_frm.get_offset(plane); | |
308 | let soff = srcbuf.get_offset(plane); | |
309 | let doff = frm.offset[plane]; | |
310 | let lstride = self.last_frm.get_stride(plane); | |
311 | let sstride = srcbuf.get_stride(plane); | |
312 | let dstride = frm.stride[plane]; | |
313 | ||
314 | let cur_w = (frm.width[plane] + 7) & !7; | |
315 | let cur_h = (frm.height[plane] + 7) & !7; | |
316 | let dst = &mut frm.data[doff..]; | |
317 | let src = &src[soff..]; | |
318 | let last = &last[loff..]; | |
319 | ||
320 | if is_b { | |
321 | // copy last frame as motion search is performed on partially updated frame | |
322 | for (dline, sline) in dst.chunks_mut(dstride).zip(last.chunks(lstride)).take(cur_h) { | |
323 | dline[..cur_w].copy_from_slice(&sline[..cur_w]); | |
324 | } | |
325 | } | |
326 | ||
327 | self.bundles.reset(); | |
328 | for (row, stripe) in src.chunks(sstride * 8).take(cur_h / 8).enumerate() { | |
329 | self.bundles.new_row(row); | |
330 | let y = row * 8; | |
331 | for x in (0..cur_w).step_by(8) { | |
332 | for (dst, src) in cur_blk.chunks_exact_mut(8).zip(stripe[x..].chunks(sstride)) { | |
333 | dst.copy_from_slice(&src[..8]); | |
334 | } | |
335 | ||
336 | let (skip_dist, skip_diff) = if !first_frame && !cur_forbidden[usize::from(BlockMode::Skip)] { | |
337 | let diff = calc_diff(&cur_blk, 8, &last[x + y * lstride..], lstride); | |
338 | (self.rc.metric(diff, 0), diff) | |
339 | } else { // no skip blocks for the first frame | |
340 | (MAX_DIST, MAX_DIST) | |
341 | }; | |
342 | let mut block_mode = BlockMode::Skip; | |
343 | let mut best_dist = skip_dist; | |
344 | self.bst_tokens.clear(); | |
345 | ||
346 | if best_dist > THRESHOLD && !first_frame && !cur_forbidden[usize::from(BlockMode::Motion)] { | |
347 | let (mv, diff) = if is_b { | |
348 | mv_search(dst, dstride, cur_w, cur_h, x, y, skip_diff, &cur_blk, &mut mv_blk) | |
349 | } else { | |
350 | mv_search(last, lstride, cur_w, cur_h, x, y, skip_diff, &cur_blk, &mut mv_blk) | |
351 | }; | |
352 | if is_b { | |
353 | get_block(dst, dstride, x, y, mv, &mut mv_blk); | |
354 | } else { | |
355 | get_block(last, lstride, x, y, mv, &mut mv_blk); | |
356 | } | |
357 | self.tmp_tokens.clear(); | |
358 | self.tmp_tokens.xoff.push(mv.x as i8); | |
359 | self.tmp_tokens.yoff.push(mv.y as i8); | |
360 | let mv_dist = self.rc.metric(diff, self.tmp_tokens.bits(is_b)); | |
361 | if mv_dist < best_dist { | |
362 | block_mode = BlockMode::Motion; | |
363 | best_dist = mv_dist; | |
364 | std::mem::swap(&mut self.bst_tokens, &mut self.tmp_tokens); | |
365 | } | |
366 | self.dsp.get_diff(&mv_blk, &cur_blk); | |
367 | if best_dist > THRESHOLD && !cur_forbidden[usize::from(BlockMode::Residue)] { | |
368 | self.tmp_tokens.clear(); | |
369 | self.tmp_tokens.xoff.push(mv.x as i8); | |
370 | self.tmp_tokens.yoff.push(mv.y as i8); | |
371 | let diff = self.dsp.try_residue(&mut self.tmp_tokens); | |
372 | let res_dist = self.rc.metric(diff, self.tmp_tokens.bits(is_b)); | |
373 | if res_dist < best_dist { | |
374 | best_dist = res_dist; | |
375 | block_mode = BlockMode::Residue; | |
376 | std::mem::swap(&mut self.bst_tokens, &mut self.tmp_tokens); | |
377 | } | |
378 | } | |
379 | if best_dist > THRESHOLD && !cur_forbidden[usize::from(BlockMode::Inter)] { | |
380 | self.tmp_tokens.clear(); | |
381 | self.tmp_tokens.xoff.push(mv.x as i8); | |
382 | self.tmp_tokens.yoff.push(mv.y as i8); | |
383 | let dct_p_dist = self.dsp.try_dct_inter(&mv_blk, &cur_blk, &mut self.tmp_tokens, &mut self.tmp_tok2, is_b, &self.rc, best_dist); | |
384 | if dct_p_dist < best_dist { | |
385 | best_dist = dct_p_dist; | |
386 | block_mode = BlockMode::Inter; | |
387 | std::mem::swap(&mut self.bst_tokens, &mut self.tmp_tokens); | |
388 | } | |
389 | } | |
390 | } | |
391 | if best_dist > THRESHOLD && !cur_forbidden[usize::from(BlockMode::Fill)] { | |
392 | let sum = cur_blk.iter().fold(0u16, | |
393 | |acc, &a| acc + u16::from(a)); | |
394 | let avg = ((sum + 32) >> 6) as u8; | |
395 | self.tmp_tokens.clear(); | |
396 | self.tmp_tokens.colors.push(avg); | |
397 | let diff = cur_blk.iter().fold(0u32, | |
398 | |acc, &a| acc + pix_dist(a, avg)); | |
399 | let fill_dist = self.rc.metric(diff, self.tmp_tokens.bits(is_b)); | |
400 | if fill_dist < best_dist { | |
401 | block_mode = BlockMode::Fill; | |
402 | best_dist = fill_dist; | |
403 | std::mem::swap(&mut self.bst_tokens, &mut self.tmp_tokens); | |
404 | } | |
405 | } | |
406 | if best_dist > THRESHOLD && !cur_forbidden[usize::from(BlockMode::Pattern)] { | |
407 | let diff = find_pattern(&cur_blk, &mut self.tmp_tokens); | |
408 | let pat_dist = self.rc.metric(diff, self.tmp_tokens.bits(is_b)); | |
409 | if pat_dist < best_dist { | |
410 | best_dist = pat_dist; | |
411 | block_mode = BlockMode::Pattern; | |
412 | std::mem::swap(&mut self.bst_tokens, &mut self.tmp_tokens); | |
413 | } | |
414 | } | |
415 | if best_dist > THRESHOLD && !cur_forbidden[usize::from(BlockMode::Run)] { | |
416 | let diff = find_run_pattern(&cur_blk, &mut self.tmp_tokens, &mut self.tmp_tok2, is_b, pat_run_thr); | |
417 | let run_dist = self.rc.metric(diff, self.tmp_tokens.bits(is_b)); | |
418 | if run_dist < best_dist { | |
419 | best_dist = run_dist; | |
420 | block_mode = BlockMode::Run; | |
421 | std::mem::swap(&mut self.bst_tokens, &mut self.tmp_tokens); | |
422 | } | |
423 | } | |
424 | if best_dist > THRESHOLD && !cur_forbidden[usize::from(BlockMode::Intra)] { | |
425 | let dct_i_dist = self.dsp.try_dct_intra(&cur_blk, &mut self.tmp_tokens, &mut self.tmp_tok2, is_b, &self.rc, best_dist); | |
426 | if dct_i_dist < best_dist { | |
427 | best_dist = dct_i_dist; | |
428 | block_mode = BlockMode::Intra; | |
429 | std::mem::swap(&mut self.bst_tokens, &mut self.tmp_tokens); | |
430 | } | |
431 | } | |
432 | if best_dist > THRESHOLD && !cur_forbidden[usize::from(BlockMode::Raw)] | |
433 | && (!is_b || self.bundles.can_fit_raw_block()) { | |
434 | self.tmp_tokens.clear(); | |
435 | self.tmp_tokens.colors.extend_from_slice(&cur_blk); | |
436 | let raw_dist = self.rc.metric(0, self.tmp_tokens.bits(is_b)); | |
437 | if raw_dist < best_dist { | |
438 | best_dist = raw_dist; | |
439 | block_mode = BlockMode::Raw; | |
440 | std::mem::swap(&mut self.bst_tokens, &mut self.tmp_tokens); | |
441 | } | |
442 | } | |
443 | let _ = best_dist; // silence a warning | |
444 | ||
445 | let bmode = if is_b { block_mode.to_code_b() } else { unimplemented!() }; | |
446 | self.bundles.add_block_type(bmode); | |
447 | self.bundles.add_tokens(&self.bst_tokens); | |
448 | ||
449 | self.blk_stats[usize::from(block_mode)] += 1; | |
450 | ||
451 | match block_mode { | |
452 | BlockMode::Skip if is_b => {}, | |
453 | BlockMode::Skip => { | |
454 | for (dline, sline) in dst[x + y * dstride..].chunks_mut(dstride) | |
455 | .zip(last[x + y * lstride..].chunks(lstride)).take(8) { | |
456 | dline[..8].copy_from_slice(&sline[..8]); | |
457 | } | |
458 | is_intra = false; | |
459 | } | |
460 | BlockMode::Fill => { | |
461 | let avg = self.bst_tokens.colors[0]; | |
462 | for dline in dst[x + y * dstride..].chunks_mut(dstride).take(8) { | |
463 | for el in dline[..8].iter_mut() { | |
464 | *el = avg; | |
465 | } | |
466 | } | |
467 | }, | |
468 | BlockMode::Pattern => { | |
469 | for (&pat, dline) in self.bst_tokens.pattern.iter().zip(dst[x + y * dstride..].chunks_mut(dstride)) { | |
470 | let mut pattern = pat as usize; | |
471 | for el in dline[..8].iter_mut() { | |
472 | *el = self.bst_tokens.colors[pattern & 1]; | |
473 | pattern >>= 1; | |
474 | } | |
475 | } | |
476 | }, | |
477 | BlockMode::Run => { | |
478 | let mut clrs = self.bst_tokens.colors.iter(); | |
479 | let mut data = self.bst_tokens.other.iter(); | |
480 | let &(idx, _) = data.next().unwrap(); | |
481 | let pattern = BINK_PATTERNS[usize::from(idx)]; | |
482 | let mut len = 0; | |
483 | let mut is_run = false; | |
484 | let mut run_val = 0; | |
485 | for (i, &idx) in pattern.iter().enumerate() { | |
486 | let dst_idx = (idx & 7) as usize + ((idx >> 3) as usize) * dstride; | |
487 | if len == 0 { | |
488 | if i < 63 { | |
489 | let &(flag, _nbits) = data.next().unwrap(); | |
490 | assert_eq!(_nbits, 1); | |
491 | is_run = flag != 0; | |
492 | let &(len1, _) = data.next().unwrap(); | |
493 | len = usize::from(len1) + 1; | |
494 | if is_run { | |
495 | run_val = *clrs.next().unwrap(); | |
496 | } | |
497 | } else { | |
498 | len = 1; | |
499 | is_run = false; | |
500 | } | |
501 | } | |
502 | dst[x + y * dstride + dst_idx] = if is_run { | |
503 | run_val | |
504 | } else { | |
505 | *clrs.next().unwrap() | |
506 | }; | |
507 | len -= 1; | |
508 | } | |
509 | }, | |
510 | BlockMode::Raw => { | |
511 | put_block(&mut dst[x + y * dstride..], dstride, &cur_blk); | |
512 | }, | |
513 | BlockMode::Motion => { | |
514 | put_block(&mut dst[x + y * dstride..], dstride, &mv_blk); | |
515 | is_intra = false; | |
516 | }, | |
517 | BlockMode::Residue => { | |
518 | self.dsp.recon_residue(&mut dst[x + y * dstride..], dstride, &mv_blk); | |
519 | is_intra = false; | |
520 | }, | |
521 | BlockMode::Inter => { | |
522 | self.dsp.recon_dct_p(&mut dst[x + y * dstride..], dstride); | |
523 | let q = if is_b { | |
524 | usize::from(self.bst_tokens.interq[0]) | |
525 | } else { | |
526 | let (qval, _) = self.bst_tokens.other[self.bst_tokens.other.len() - 1]; | |
527 | usize::from(qval) | |
528 | }; | |
529 | self.pq_stats[q] += 1; | |
530 | is_intra = false; | |
531 | }, | |
532 | BlockMode::Intra => { | |
533 | self.dsp.recon_dct_i(&mut dst[x + y * dstride..], dstride); | |
534 | let q = if is_b { | |
535 | usize::from(self.bst_tokens.intraq[0]) | |
536 | } else { | |
537 | let (qval, _) = self.bst_tokens.other[self.bst_tokens.other.len() - 1]; | |
538 | usize::from(qval) | |
539 | }; | |
540 | self.iq_stats[q] += 1; | |
541 | }, | |
542 | _ => unimplemented!(), | |
543 | }; | |
544 | } | |
545 | self.bundles.end_row(); | |
546 | } | |
547 | for row in 0..(cur_h / 8) { | |
548 | self.bundles.write(bw, row); | |
549 | } | |
550 | while (bw.tell() & 0x1F) != 0 { | |
551 | bw.write0(); | |
552 | } | |
553 | } | |
554 | Ok(is_intra) | |
555 | } | |
556 | fn encode_skip(&mut self, bw: &mut BitWriter) -> EncoderResult<()> { | |
557 | let src = self.last_frm.get_data(); | |
558 | let dst = self.cur_frm.get_data_mut().unwrap(); | |
559 | dst.copy_from_slice(src); | |
560 | ||
561 | for plane in 0..3 { | |
562 | let (width, height) = self.cur_frm.get_dimensions(plane); | |
563 | let tiles_w = (width + 7) >> 3; | |
564 | let tiles_h = (height + 7) >> 3; | |
565 | self.bundles.reset(); | |
566 | for row in 0..tiles_h { | |
567 | self.bundles.new_row(row); | |
568 | for _ in 0..tiles_w { | |
569 | self.bundles.add_block_type(0); // skip always has code 0 | |
570 | } | |
571 | self.bundles.end_row(); | |
572 | } | |
573 | for row in 0..tiles_h { | |
574 | self.bundles.write(bw, row); | |
575 | } | |
576 | } | |
577 | ||
578 | Ok(()) | |
579 | } | |
580 | } | |
581 | ||
582 | impl NAEncoder for BinkEncoder { | |
583 | fn negotiate_format(&self, encinfo: &EncodeParameters) -> EncoderResult<EncodeParameters> { | |
584 | match encinfo.format { | |
585 | NACodecTypeInfo::None => { | |
586 | Ok(EncodeParameters { | |
587 | format: NACodecTypeInfo::Video(NAVideoInfo::new(0, 0, true, YUV420_FORMAT)), | |
588 | ..Default::default() | |
589 | }) | |
590 | }, | |
591 | NACodecTypeInfo::Audio(_) => Err(EncoderError::FormatError), | |
592 | NACodecTypeInfo::Video(vinfo) => { | |
593 | let outinfo = NAVideoInfo::new(vinfo.width, vinfo.height, false, YUV420_FORMAT); | |
594 | let mut ofmt = *encinfo; | |
595 | ofmt.format = NACodecTypeInfo::Video(outinfo); | |
596 | Ok(ofmt) | |
597 | } | |
598 | } | |
599 | } | |
600 | fn get_capabilities(&self) -> u64 { ENC_CAPS_SKIPFRAME } | |
601 | fn init(&mut self, stream_id: u32, encinfo: EncodeParameters) -> EncoderResult<NAStreamRef> { | |
602 | match encinfo.format { | |
603 | NACodecTypeInfo::None => Err(EncoderError::FormatError), | |
604 | NACodecTypeInfo::Audio(_) => Err(EncoderError::FormatError), | |
605 | NACodecTypeInfo::Video(vinfo) => { | |
606 | if vinfo.format != YUV420_FORMAT { | |
607 | return Err(EncoderError::FormatError); | |
608 | } | |
609 | ||
610 | let mut edata = vec![b'B', b'I', b'K', self.version as u8, 0, 0, 0, 0]; | |
611 | match self.scale_mode { | |
612 | DoublingMode::None => {}, | |
613 | DoublingMode::Height2X => { | |
614 | edata[7] |= 0x10; | |
615 | }, | |
616 | DoublingMode::HeightIlace => { | |
617 | edata[7] |= 0x20; | |
618 | }, | |
619 | DoublingMode::Width2X => { | |
620 | edata[7] |= 0x30; | |
621 | }, | |
622 | DoublingMode::Scale2X => { | |
623 | edata[7] |= 0x40; | |
624 | }, | |
625 | DoublingMode::Width2XIlace => { | |
626 | edata[7] |= 0x50; | |
627 | }, | |
628 | }; | |
629 | ||
630 | if self.nframes == 0 { | |
631 | println!("Bink should set the number of frames in the stream"); | |
632 | return Err(EncoderError::FormatError); | |
633 | } | |
634 | ||
635 | let out_info = NAVideoInfo::new(vinfo.width, vinfo.height, false, YUV420_FORMAT); | |
636 | let info = NACodecInfo::new("bink-video", NACodecTypeInfo::Video(out_info), Some(edata)); | |
637 | let mut stream = NAStream::new(StreamType::Video, stream_id, info, encinfo.tb_num, encinfo.tb_den, self.nframes); | |
638 | stream.set_num(stream_id as usize); | |
639 | let stream = stream.into_ref(); | |
640 | ||
641 | self.stream = Some(stream.clone()); | |
642 | ||
643 | let frm = alloc_video_buffer(out_info, 4)?; | |
644 | self.cur_frm = frm.get_vbuf().unwrap(); | |
645 | let frm = alloc_video_buffer(out_info, 4)?; | |
646 | self.last_frm = frm.get_vbuf().unwrap(); | |
647 | ||
648 | let cdata = self.cur_frm.get_data_mut().unwrap(); | |
649 | for el in cdata.iter_mut() { | |
650 | *el = 0x00; | |
651 | } | |
652 | let cdata = self.last_frm.get_data_mut().unwrap(); | |
653 | for el in cdata.iter_mut() { | |
654 | *el = 0x00; | |
655 | } | |
656 | ||
657 | self.rc.init(encinfo.tb_num, encinfo.tb_den, encinfo.bitrate, encinfo.quality); | |
658 | ||
659 | Ok(stream) | |
660 | }, | |
661 | } | |
662 | } | |
663 | fn encode(&mut self, frm: &NAFrame) -> EncoderResult<()> { | |
664 | if self.frame_no >= self.nframes { | |
665 | return Ok(()); | |
666 | } | |
667 | self.frame_no += 1; | |
668 | let mut bw = BitWriter::new(Vec::with_capacity(42), BitWriterMode::LE); | |
669 | ||
670 | let is_intra = match frm.get_buffer() { | |
671 | NABufferType::Video(ref buf) => { | |
672 | self.encode_frame(&mut bw, buf)? | |
673 | }, | |
674 | NABufferType::None => { | |
675 | self.encode_skip(&mut bw)?; | |
676 | false | |
677 | }, | |
678 | _ => return Err(EncoderError::InvalidParameters), | |
679 | }; | |
680 | let dbuf = bw.end(); | |
681 | self.tot_size += dbuf.len(); | |
682 | self.rc.update_size(dbuf.len()); | |
683 | self.pkt = Some(NAPacket::new(self.stream.clone().unwrap(), frm.ts, is_intra, dbuf)); | |
684 | ||
685 | std::mem::swap(&mut self.cur_frm, &mut self.last_frm); | |
686 | Ok(()) | |
687 | } | |
688 | fn get_packet(&mut self) -> EncoderResult<Option<NAPacket>> { | |
689 | let mut npkt = None; | |
690 | std::mem::swap(&mut self.pkt, &mut npkt); | |
691 | Ok(npkt) | |
692 | } | |
693 | fn flush(&mut self) -> EncoderResult<()> { | |
694 | Ok(()) | |
695 | } | |
696 | } | |
697 | ||
698 | impl Drop for BinkEncoder { | |
699 | fn drop(&mut self) { | |
700 | if self.print_stats { | |
701 | println!("encoded {} frame(s)", self.frame_no); | |
702 | println!("block statistics:"); | |
703 | for (name, &count) in BLOCK_MODE_NAMES.iter().zip(self.blk_stats.iter()) { | |
704 | if count != 0 { | |
705 | println!(" {:8}: {:8}", name, count); | |
706 | } | |
707 | } | |
708 | if self.blk_stats[usize::from(BlockMode::Intra)] != 0 { | |
709 | print!("intra block quants:"); | |
710 | for &count in self.iq_stats.iter() { | |
711 | print!(" {}", count); | |
712 | } | |
713 | println!(); | |
714 | } | |
715 | if self.blk_stats[usize::from(BlockMode::Inter)] != 0 { | |
716 | print!("inter block quants:"); | |
717 | for &count in self.pq_stats.iter() { | |
718 | print!(" {}", count); | |
719 | } | |
720 | println!(); | |
721 | } | |
722 | if self.frame_no > 0 { | |
723 | println!("average frame size {} byte(s)", self.tot_size / (self.frame_no as usize)); | |
724 | if let Some(ref stream) = self.stream { | |
725 | let bitrate = (self.tot_size as u64) * 8 * u64::from(stream.tb_den) / u64::from(stream.tb_num) / self.frame_no; | |
726 | let br_fmt = if bitrate >= 10_000_000 { | |
727 | format!("{}mbps", bitrate / 1000000) | |
728 | } else if bitrate >= 10_000 { | |
729 | format!("{}kbps", bitrate / 1000) | |
730 | } else { | |
731 | format!("{}bps", bitrate) | |
732 | }; | |
733 | println!("average bitrate {}", br_fmt); | |
734 | } | |
735 | } | |
736 | } | |
737 | } | |
738 | } | |
739 | ||
740 | const ENCODER_OPTS: &[NAOptionDefinition] = &[ | |
741 | NAOptionDefinition { | |
742 | name: "nframes", description: "duration in frames", | |
743 | opt_type: NAOptionDefinitionType::Int(Some(0), None) }, | |
744 | NAOptionDefinition { | |
745 | name: "version", description: "codec version", | |
746 | opt_type: NAOptionDefinitionType::String(Some(&["b", "f", "g", "h", "i", "k"])) }, | |
747 | NAOptionDefinition { | |
748 | name: "scale_mode", description: "output scaling mode", | |
749 | opt_type: NAOptionDefinitionType::String(Some(&["none", "height2x", "height_il", | |
750 | "width2x", "width2x_il", "scale2x"])) }, | |
751 | NAOptionDefinition { | |
752 | name: "forbidden", description: "block coding modes to omit (e.g. inter+residue+run)", | |
753 | opt_type: NAOptionDefinitionType::String(None) }, | |
754 | NAOptionDefinition { | |
755 | name: "print_stats", description: "print internal encoding statistics at the end", | |
756 | opt_type: NAOptionDefinitionType::Bool }, | |
757 | ]; | |
758 | ||
759 | impl NAOptionHandler for BinkEncoder { | |
760 | fn get_supported_options(&self) -> &[NAOptionDefinition] { ENCODER_OPTS } | |
761 | fn set_options(&mut self, options: &[NAOption]) { | |
762 | for option in options.iter() { | |
763 | for opt_def in ENCODER_OPTS.iter() { | |
764 | if opt_def.check(option).is_ok() { | |
765 | match option.name { | |
766 | "version" => { | |
767 | if let NAValue::String(ref strval) = option.value { | |
768 | match strval.as_str() { | |
769 | "b" => self.version = 'b', | |
770 | _ => { | |
771 | println!("versions beside 'b' are not supported"); | |
772 | }, | |
773 | }; | |
774 | } | |
775 | }, | |
776 | "scale_mode" => { | |
777 | if let NAValue::String(ref strval) = option.value { | |
778 | match strval.as_str() { | |
779 | "none" => self.scale_mode = DoublingMode::None, | |
780 | "height2x" => self.scale_mode = DoublingMode::Height2X, | |
781 | "height_il" => self.scale_mode = DoublingMode::HeightIlace, | |
782 | "scale2x" => self.scale_mode = DoublingMode::Scale2X, | |
783 | "width2x" => self.scale_mode = DoublingMode::Width2X, | |
784 | "width2x_il" => self.scale_mode = DoublingMode::Width2XIlace, | |
785 | _ => {}, | |
786 | }; | |
787 | } | |
788 | }, | |
789 | "nframes" => { | |
790 | if let NAValue::Int(ival) = option.value { | |
791 | self.nframes = ival as u64; | |
792 | } | |
793 | }, | |
794 | "forbidden" => { | |
795 | if let NAValue::String(ref strval) = option.value { | |
796 | for el in self.forbidden.iter_mut() { | |
797 | *el = false; | |
798 | } | |
799 | for name in strval.split('+') { | |
800 | if let Ok(bmode) = name.parse::<BlockMode>() { | |
801 | self.forbidden[usize::from(bmode)] = true; | |
802 | } | |
803 | } | |
804 | } | |
805 | }, | |
806 | "print_stats" => { | |
807 | if let NAValue::Bool(bval) = option.value { | |
808 | self.print_stats = bval; | |
809 | } | |
810 | }, | |
811 | _ => {}, | |
812 | }; | |
813 | } | |
814 | } | |
815 | } | |
816 | } | |
817 | fn query_option_value(&self, name: &str) -> Option<NAValue> { | |
818 | match name { | |
819 | "version" => Some(NAValue::String(self.version.to_string())), | |
820 | "scale_mode" => Some(NAValue::String(self.scale_mode.to_string())), | |
821 | "nframes" => Some(NAValue::Int(self.nframes as i64)), | |
822 | "forbidden" => { | |
823 | let mut result = String::new(); | |
824 | for (name, &flag) in BLOCK_MODE_NAMES.iter().zip(self.forbidden.iter()) { | |
825 | if flag { | |
826 | if !result.is_empty() { | |
827 | result.push('+'); | |
828 | } | |
829 | result += name; | |
830 | } | |
831 | } | |
832 | Some(NAValue::String(result)) | |
833 | }, | |
834 | "print_stats" => Some(NAValue::Bool(self.print_stats)), | |
835 | _ => None, | |
836 | } | |
837 | } | |
838 | } | |
839 | ||
840 | pub fn get_encoder() -> Box<dyn NAEncoder + Send> { | |
841 | Box::new(BinkEncoder::new()) | |
842 | } | |
843 | ||
844 | #[cfg(test)] | |
845 | mod test { | |
846 | use nihav_core::codecs::*; | |
847 | use nihav_core::demuxers::*; | |
848 | use nihav_core::muxers::*; | |
849 | use crate::*; | |
850 | use nihav_codec_support::test::enc_video::*; | |
851 | use nihav_commonfmt::*; | |
852 | ||
853 | fn test_bink_encoder(out_name: &'static str, enc_options: &[NAOption], br: u32, quality: u8, hash: &[u32; 4]) { | |
854 | let mut dmx_reg = RegisteredDemuxers::new(); | |
855 | generic_register_all_demuxers(&mut dmx_reg); | |
856 | let mut dec_reg = RegisteredDecoders::new(); | |
857 | generic_register_all_decoders(&mut dec_reg); | |
858 | let mut mux_reg = RegisteredMuxers::new(); | |
859 | rad_register_all_muxers(&mut mux_reg); | |
860 | let mut enc_reg = RegisteredEncoders::new(); | |
861 | rad_register_all_encoders(&mut enc_reg); | |
862 | ||
863 | // sample from private collection | |
864 | let dec_config = DecoderTestParams { | |
865 | demuxer: "yuv4mpeg", | |
866 | in_name: "assets/day3b.y4m", | |
867 | stream_type: StreamType::Video, | |
868 | limit: None, | |
869 | dmx_reg, dec_reg, | |
870 | }; | |
871 | let enc_config = EncoderTestParams { | |
872 | muxer: "bink", | |
873 | enc_name: "bink-video", | |
874 | out_name, | |
875 | mux_reg, enc_reg, | |
876 | }; | |
877 | let dst_vinfo = NAVideoInfo { | |
878 | width: 0, | |
879 | height: 0, | |
880 | format: YUV420_FORMAT, | |
881 | flipped: false, | |
882 | bits: 8, | |
883 | }; | |
884 | let enc_params = EncodeParameters { | |
885 | format: NACodecTypeInfo::Video(dst_vinfo), | |
886 | quality, | |
887 | bitrate: br * 1000, | |
888 | tb_num: 0, | |
889 | tb_den: 0, | |
890 | flags: 0, | |
891 | }; | |
892 | let _ = hash; | |
893 | //test_encoding_to_file(&dec_config, &enc_config, enc_params, enc_options); | |
894 | test_encoding_md5(&dec_config, &enc_config, enc_params, enc_options, hash); | |
895 | } | |
896 | #[test] | |
897 | fn test_binkb_quality() { | |
898 | let enc_options = &[ | |
899 | NAOption { name: "nframes", value: NAValue::Int(7) }, | |
900 | //NAOption { name: "print_stats", value: NAValue::Bool(true) }, | |
901 | ]; | |
902 | test_bink_encoder("bink-b-q50.bik", enc_options, 0, 50, &[0xd83936aa, 0xec3f55d4, 0x25e5c1fb, 0x0f3454ce]); | |
903 | let enc_options = &[ | |
904 | NAOption { name: "nframes", value: NAValue::Int(7) }, | |
905 | //NAOption { name: "print_stats", value: NAValue::Bool(true) }, | |
906 | ]; | |
907 | test_bink_encoder("bink-b-q75.bik", enc_options, 0, 75, &[0x45ccd3d4, 0xf09bd106, 0xc88751db, 0xca5294d7]); | |
908 | let enc_options = &[ | |
909 | NAOption { name: "nframes", value: NAValue::Int(7) }, | |
910 | //NAOption { name: "print_stats", value: NAValue::Bool(true) }, | |
911 | ]; | |
912 | test_bink_encoder("bink-b-q99.bik", enc_options, 0, 99, &[0xb516554e, 0xca025167, 0xd6c3dc06, 0x00e6ba25]); | |
913 | } | |
914 | #[test] | |
915 | fn test_binkb_features() { | |
916 | let enc_options = &[ | |
917 | NAOption { name: "nframes", value: NAValue::Int(7) }, | |
918 | NAOption { name: "forbidden", value: NAValue::String("intra+inter".to_string()) }, | |
919 | //NAOption { name: "print_stats", value: NAValue::Bool(true) }, | |
920 | ]; | |
921 | test_bink_encoder("bink-b-nodct.bik", enc_options, 0, 0, &[0x5e098760, 0x31c8982a, 0x90ce8441, 0x859d3cc6]); | |
922 | let enc_options = &[ | |
923 | NAOption { name: "nframes", value: NAValue::Int(7) }, | |
924 | NAOption { name: "forbidden", value: NAValue::String("skip+fill+run+pattern+residue+raw".to_string()) }, | |
925 | //NAOption { name: "print_stats", value: NAValue::Bool(true) }, | |
926 | ]; | |
927 | test_bink_encoder("bink-b-dct.bik", enc_options, 0, 0, &[0xed2fc7d2, 0x8a7a05ef, 0xd0b4ae2c, 0x622a4ef0]); | |
928 | } | |
929 | } |