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