| 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 | } |