fix clippy warnings
[nihav.git] / nihav-realmedia / src / codecs / rv40enc / mod.rs
CommitLineData
4965a5e5
KS
1use std::collections::VecDeque;
2use std::str::FromStr;
3
4use nihav_core::codecs::*;
5use nihav_core::io::byteio::*;
6use nihav_core::io::bitwriter::*;
7
8mod bitstream;
9use bitstream::*;
10
11mod dsp;
12use dsp::loop_filter_frame;
13
14mod estimator;
15use estimator::*;
16
17mod mb_coding;
18use mb_coding::*;
19
20mod motion_est;
21use motion_est::*;
22
23mod ratectl;
24use ratectl::*;
25pub use ratectl::RateDistMetric;
26
27mod types;
28pub use types::*;
29
30const DEBUG_BIT_FRAMENO: u8 = 0;
31const DEBUG_BIT_SLICE_SIZE: u8 = 1;
32const DEBUG_BIT_PSNR: u8 = 2;
33const DEBUG_BIT_RATECTL: u8 = 3;
34const DEBUG_FLAG_BITS: &[(&str, u8)] = &[
35 ("frameno", DEBUG_BIT_FRAMENO),
36 ("slicesize", DEBUG_BIT_SLICE_SIZE),
37 ("psnr", DEBUG_BIT_PSNR),
38 ("rc", DEBUG_BIT_RATECTL),
39];
40
41#[derive(Clone,Copy,Default)]
42struct DebugFlags {
43 flags: u32,
44}
45
46impl DebugFlags {
47 fn new() -> Self { Self::default() }
48 fn is_set(self, bit: u8) -> bool { (self.flags & (1 << bit)) != 0 }
49 fn parse(&mut self, args: &str) {
50 self.flags = 0;
51 for arg in args.split('+') {
52 for &(name, bit) in DEBUG_FLAG_BITS.iter() {
53 if name == arg {
54 self.flags += 1 << bit;
55 }
56 }
57 }
58 }
59}
60
61impl std::fmt::Display for DebugFlags {
62 fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
63 let mut flags = String::new();
64 let mut first = true;
65 for &(name, bit) in DEBUG_FLAG_BITS.iter() {
66 if self.is_set(bit) {
67 if !first {
68 flags.push('+');
69 }
70 flags.push_str(name);
71 first = false;
72 }
73 }
74 write!(f, "{}", flags)
75 }
76}
77
78struct StaticFrameOrder {
79 groups: Vec<(FrameType, usize)>,
80 start: bool,
81 cur_grp: usize,
82 cur_frm: usize,
83}
84
85impl StaticFrameOrder {
86 /*fn new() -> Self {
87 Self {
88 groups: vec![(FrameType::I, 0)],
89 start: true,
90 cur_grp: 0,
91 cur_frm: 0,
92 }
93 }*/
94 fn get_max_grp_len(&self) -> usize {
95 let mut g_len = 1;
96 for &(_, num_b) in self.groups.iter() {
97 g_len = g_len.max(1 + num_b);
98 }
99 g_len
100 }
101 fn peek_next_frame(&self) -> (FrameType, usize) {
102 if !self.start {
103 let grp = &self.groups[self.cur_grp];
104 if self.cur_frm == 0 {
105 (grp.0, grp.1)
106 } else {
107 (FrameType::B, 0)
108 }
109 } else {
110 (FrameType::I, 0)
111 }
112 }
113 fn next_frame(&mut self) -> FrameType {
114 if !self.start {
115 let grp = &self.groups[self.cur_grp];
116 let frm_type = if self.cur_frm == 0 {
117 grp.0
118 } else {
119 FrameType::B
120 };
121 self.cur_frm += 1;
122 if self.cur_frm > grp.1 {
123 self.cur_frm = 0;
124 self.cur_grp += 1;
125 if self.cur_grp >= self.groups.len() {
126 self.cur_grp = 0;
127 }
128 }
129 frm_type
130 } else {
131 self.start = false;
132 self.cur_grp = if self.groups.len() > 1 { 1 } else { 0 };
133 self.cur_frm = 0;
134 FrameType::I
135 }
136 }
137}
138
139impl std::fmt::Display for StaticFrameOrder {
140 fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
141 let mut seq = String::with_capacity(self.groups.len() * 2);
142 for &(ftype, brun) in self.groups.iter() {
143 seq.push(match ftype {
144 FrameType::I => 'I',
145 _ => 'P',
146 });
147 for _ in 0..brun {
148 seq.push('B');
149 }
150 }
151 write!(f, "{}", seq)
152 }
153}
154
155struct DynamicFrameOrder {
156 cur_ft: FrameType,
157 next_ft: FrameType,
158 p_count: usize,
159}
160
161const NUM_GOP_KF: usize = 8;
162
163impl DynamicFrameOrder {
164 fn new() -> Self {
165 Self {
166 cur_ft: FrameType::I,
167 next_ft: FrameType::Other,
168 p_count: 0,
169 }
170 }
171 fn peek_next_frame(&self) -> (FrameType, usize) {
172 (self.cur_ft, if self.cur_ft == FrameType::Other || self.next_ft == FrameType::B { 1 } else { 0 })
173 }
174 fn next_frame(&mut self) -> FrameType {
175 if self.cur_ft == FrameType::P {
176 self.p_count += 1;
177 if self.p_count >= NUM_GOP_KF {
178 self.cur_ft = FrameType::I;
179 self.p_count = 0;
180 }
181 }
182 let next = self.cur_ft;
183 self.cur_ft = self.next_ft;
184 self.next_ft = if self.cur_ft != FrameType::B { FrameType::Other } else { FrameType::P };
185 next
186 }
187 fn update(&mut self, ftype: FrameType) {
188 if self.cur_ft == FrameType::Other {
189 self.cur_ft = ftype;
190 if self.cur_ft == FrameType::B {
191 self.cur_ft = FrameType::P;
192 self.next_ft = FrameType::B;
193 } else {
194 self.next_ft = FrameType::Other;
195 }
196 }
197 }
198}
199
200impl std::fmt::Display for DynamicFrameOrder {
201 fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
202 write!(f, "dynamic")
203 }
204}
205
206enum FrameOrder {
207 Static(StaticFrameOrder),
208 Dynamic(DynamicFrameOrder),
209}
210
211impl FrameOrder {
212 fn new() -> Self {
213 FrameOrder::Dynamic(DynamicFrameOrder::new())
214 }
215 fn get_max_grp_len(&self) -> usize {
216 match self {
217 FrameOrder::Static(ref order) => order.get_max_grp_len(),
218 FrameOrder::Dynamic(ref _order) => 2,
219 }
220 }
221 fn peek_next_frame(&self) -> (FrameType, usize) {
222 match self {
223 FrameOrder::Static(ref order) => order.peek_next_frame(),
224 FrameOrder::Dynamic(ref order) => order.peek_next_frame(),
225 }
226 }
227 fn next_frame(&mut self) -> FrameType {
228 match self {
229 FrameOrder::Static(ref mut order) => order.next_frame(),
230 FrameOrder::Dynamic(ref mut order) => order.next_frame(),
231 }
232 }
233 fn update(&mut self, ftype: FrameType) {
234 if let FrameOrder::Dynamic(ref mut order) = self {
235 order.update(ftype);
236 }
237 }
238 fn is_dynamic(&self) -> bool { matches!(self, FrameOrder::Dynamic(_)) }
239}
240
241impl std::fmt::Display for FrameOrder {
242 fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
243 match *self {
244 FrameOrder::Static(ref order) => order.fmt(f),
245 FrameOrder::Dynamic(ref order) => order.fmt(f),
246 }
247 }
248}
249
250#[derive(Clone,Copy,Debug)]
251enum ParseError {
252 TooShort,
253 TooLong,
254 InvalidValue,
255 InvalidCombination,
256}
257
258impl FromStr for FrameOrder {
259 type Err = ParseError;
260 fn from_str(s: &str) -> Result<Self, Self::Err> {
261 if s == "dynamic" {
262 return Ok(FrameOrder::Dynamic(DynamicFrameOrder::new()));
263 }
264 let mut ftypes = Vec::new();
265 for ch in s.bytes() {
266 match ch {
267 b'I' | b'i' => ftypes.push(FrameType::I),
268 b'P' | b'p' => ftypes.push(FrameType::P),
269 b'B' | b'b' => ftypes.push(FrameType::B),
270 b' ' | b',' => {},
271 _ => return Err(ParseError::InvalidValue),
272 };
273 if ftypes.len() > 16 {
274 return Err(ParseError::TooLong);
275 }
276 }
277 if ftypes.is_empty() {
278 return Err(ParseError::TooShort);
279 }
280 if ftypes[0] != FrameType::I {
281 return Err(ParseError::InvalidCombination);
282 }
283
284 let mut groups = Vec::new();
285 let mut cur_ftype = ftypes[0];
286 let mut cur_run = 0;
287 for &ft in ftypes[1..].iter() {
288 match ft {
289 FrameType::I | FrameType::P => {
290 groups.push((cur_ftype, cur_run));
291 cur_ftype = ft;
292 cur_run = 0;
293 },
294 _ => {
295 cur_run += 1;
296 if cur_run > 4 {
297 return Err(ParseError::InvalidCombination);
298 }
299 },
300 };
301 }
302 groups.push((cur_ftype, cur_run));
303
304 Ok(FrameOrder::Static(StaticFrameOrder{
305 groups,
306 start: true,
307 cur_grp: 0,
308 cur_frm: 0,
309 }))
310 }
311}
312
313struct RV40Encoder {
314 stream: Option<NAStreamRef>,
315 vinfo: NAVideoInfo,
316 width: usize,
317 height: usize,
318 mb_w: usize,
319 mb_h: usize,
320
321 slice_bits: u32,
322 deblock: bool,
323 force_set: Option<usize>,
324
325 fce: FrameComplexityEstimate,
326 mbd: MacroblockDecider,
327 order: FrameOrder,
328 brc: BitRateControl,
329 rdm: RateDistMetric,
330 be: BitsEstimator,
331 me: MotionEstimator,
332 cset: CodeSets,
333
334 sstate: SliceState,
335 mbstate: MBState,
336 mbs: VecDeque<Macroblock>,
337 dblk: Vec<DeblockInfo>,
338
339 qframes: Vec<NAFrame>,
340 frm_pool: NAVideoBufferPool<u8>,
341 ref_p: NAVideoBufferRef<u8>,
342 ref_n: NAVideoBufferRef<u8>,
343 pkts: VecDeque<NAPacket>,
344
345 p_pts: u64,
346 n_pts: u64,
347 last_k_ts: u64,
348 last_b_ts: u64,
349
350 needs_alloc: bool,
351 max_grp_bufs: usize,
352
353 debug_log: DebugFlags,
354
355 refine_b: bool,
356 i4_in_b: bool,
357}
358
359impl RV40Encoder {
360 fn new() -> Self {
361 let vinfo = NAVideoInfo::new(24, 24, false, YUV420_FORMAT);
362 let vt = alloc_video_buffer(vinfo, 4).unwrap();
363 let ref_p = vt.get_vbuf().unwrap();
364 let vt = alloc_video_buffer(vinfo, 4).unwrap();
365 let ref_n = vt.get_vbuf().unwrap();
366 Self {
367 stream: None,
368 vinfo,
369 width: 0,
370 height: 0,
371 mb_w: 0,
372 mb_h: 0,
373
374 slice_bits: 10000,
375 deblock: true,
376 force_set: None,
377
378 fce: FrameComplexityEstimate::new(),
379 mbd: MacroblockDecider::new(),
380 order: FrameOrder::new(),
381 brc: BitRateControl::new(),
382 rdm: RateDistMetric::new(),
383 be: BitsEstimator::new(),
384 me: MotionEstimator::new(),
385 cset: CodeSets::new(),
386
387 sstate: SliceState::new(),
388 mbstate: MBState::new(),
389 mbs: VecDeque::new(),
390 dblk: Vec::new(),
391
392 qframes: Vec::new(),
393 frm_pool: NAVideoBufferPool::new(0),
394 pkts: VecDeque::new(),
395 ref_p, ref_n,
396
397 p_pts: 0,
398 n_pts: 0,
399 last_k_ts: 0,
400 last_b_ts: 0,
401
402 needs_alloc: true,
403 max_grp_bufs: 0,
404
405 debug_log: DebugFlags::new(),
406
407 refine_b: false,
408 i4_in_b: false,
409 }
410 }
411 fn encode_frame(&mut self, frm: NAFrame, frameno: usize) -> EncoderResult<NAPacket> {
412 let ftype = self.order.next_frame();
413 let buf = frm.get_buffer();
414
415 let tinfo = frm.get_time_information();
416 let pts = NATimeInfo::ts_to_time(tinfo.pts.unwrap_or(0), 1000, tinfo.tb_num, tinfo.tb_den);
417 let fpts = (pts & 0x1FFF) as u32;
418
419 let ts_diff = if ftype == FrameType::B {
420 pts.saturating_sub(self.last_k_ts.min(self.last_b_ts)) as u32
421 } else {
422 let diff = pts.saturating_sub(self.last_k_ts) as u32;
423 diff / ((frameno + 1) as u32)
424 };
425
426 if self.debug_log.is_set(DEBUG_BIT_FRAMENO) {
427 println!("encode frame type {} pts {}", ftype, pts);
428 }
429 let is_ref_frame = matches!(ftype, FrameType::I | FrameType::P);
430
431 let tr_d = (self.n_pts - self.p_pts) as u32;
432 let tr_b = (pts - self.p_pts) as u32;
433 if !is_ref_frame {
434 self.mbd.set_b_distance(tr_b, tr_d);
435 }
436
437 let mut rvbuf = if let Some(nfrm) = self.frm_pool.get_free() {
438 nfrm
439 } else {
440 return Err(EncoderError::AllocError);
441 };
442 let mut recon_frm = NASimpleVideoFrame::from_video_buf(&mut rvbuf).unwrap();
443
444 self.be.set_frame_type(ftype);
445 if let Some(ref vbuf) = buf.get_vbuf() {
446 let src = vbuf.get_data();
447
448 if self.brc.rate_ctl_in_use() || self.order.is_dynamic() {
449 self.fce.set_current(vbuf);
450 }
451
452 let complexity = if self.brc.rate_ctl_in_use() {
453 self.fce.get_complexity(ftype)
454 } else { 0 };
455
456 self.mbd.q = self.brc.get_quant(ftype, complexity);
457 self.brc.init_metric(ftype, &mut self.rdm);
458 self.be.set_quant(self.mbd.q);
459 if self.debug_log.is_set(DEBUG_BIT_RATECTL) {
460 println!(" expected frame size {}", self.brc.get_target_size(ftype));
461 println!(" target quantiser {} lambda {} thresholds {} / {}", self.brc.get_last_quant(ftype), self.rdm.lambda, self.rdm.good_enough, self.rdm.p_split_thr);
462 }
463
464 let mut nslices = 0;
465 let mut dvec = Vec::new();
466 let mut mb_idx = 0;
467 let mut slice_starts = Vec::new();
468 let num_mbs = self.mb_w * self.mb_h;
469 while mb_idx < num_mbs {
470 slice_starts.push(dvec.len());
471 let mut bw = BitWriter::new(dvec, BitWriterMode::BE);
472 let slice_start_mb = mb_idx;
473
474 self.mbstate.reset();
475
476 let mut est_bits = 0;
477 while est_bits < self.slice_bits && mb_idx < num_mbs {
478 let mb_x = mb_idx % self.mb_w;
479 let mb_y = mb_idx / self.mb_w;
480 self.sstate.has_t = mb_idx >= slice_start_mb + self.mb_w;
481 self.sstate.has_l = (mb_idx > slice_start_mb) && (mb_x > 0);
482 self.sstate.has_tl = (mb_idx > slice_start_mb + self.mb_w) && (mb_x > 0);
483 self.sstate.has_tr = (mb_idx >= slice_start_mb + self.mb_w - 1) && (mb_x + 1 < self.mb_w);
484 self.sstate.mb_x = mb_x;
485 self.sstate.mb_y = mb_y;
486
487 let offsets = [
488 vbuf.get_offset(0) + mb_x * 16 + mb_y * 16 * vbuf.get_stride(0),
489 vbuf.get_offset(1) + mb_x * 8 + mb_y * 8 * vbuf.get_stride(1),
490 vbuf.get_offset(2) + mb_x * 8 + mb_y * 8 * vbuf.get_stride(2),
491 ];
492 let strides = [vbuf.get_stride(0), vbuf.get_stride(1), vbuf.get_stride(2)];
493 self.mbd.load_mb(src, offsets, strides, &self.sstate);
494
495 self.be.set_pred_mb_type(self.mbstate.get_pred_mbtype(&self.sstate, ftype == FrameType::B));
496 if ftype == FrameType::B {
497 self.mbd.try_b_coding(&self.ref_p, &self.ref_n, &mut self.be, &mut self.me, &self.rdm, &self.mbstate, self.refine_b);
498 }
499 if ftype == FrameType::P {
500 self.mbd.try_p_coding(&self.ref_n, &mut self.be, &mut self.me, &self.rdm, &self.mbstate);
501 }
502 self.mbd.try_intra_16_pred(&mut self.be, &self.rdm);
503 if ftype != FrameType::B || self.i4_in_b {
504 self.mbd.try_intra_4x4_pred(&mut self.be, &self.rdm, &mut self.mbstate);
505 }
506
507 let mb = self.mbd.get_macroblock();
508 est_bits += self.mbd.get_est_bits();
509 self.mbd.recon_mb(&mut recon_frm);
510 self.mbstate.update(&mb.mb_type, mb_x, mb_y);
511
512 if self.deblock {
513 self.dblk[mb_idx].q = self.mbd.q as u8;
514 if ftype == FrameType::I {
515 self.dblk[mb_idx].is_strong = true;
516 self.dblk[mb_idx].cbp_y = 0xFFFF;
517 self.dblk[mb_idx].cbp_c = 0xFF;
518 } else {
519 self.dblk[mb_idx].is_strong = mb.mb_type.is_intra() || mb.mb_type.is_16();
520 let mut cbp = 0u16;
521 let mut mask = 1;
522 for blk in mb.coeffs[..16].iter() {
523 if !blk.is_empty() {
524 cbp |= mask;
525 }
526 mask <<= 1;
527 }
528 self.dblk[mb_idx].cbp_y = cbp;
529 let mut cbp = 0u8;
530 let mut mask = 1;
531 for blk in mb.coeffs[16..24].iter() {
532 if !blk.is_empty() {
533 cbp |= mask;
534 }
535 mask <<= 1;
536 }
537 self.dblk[mb_idx].cbp_c = cbp;
538 }
539 self.mbstate.fill_deblock(&mut self.dblk[mb_idx], &self.sstate);
540 }
541
542 self.mbs.push_back(mb);
543
544 mb_idx += 1;
545 }
546
547 let set_idx = if let Some(idx) = self.force_set {
548 idx
549 } else {
550 let mut hist = [0usize; 17];
551 for mb in self.mbs.iter() {
552 let blocks = if mb.mb_type.is_16() { &mb.coeffs } else { &mb.coeffs[..24] };
553
554 for blk in blocks.iter() {
555 let nz = blk.count_nz();
556 for el in hist[nz..].iter_mut() {
557 *el += 1;
558 }
559 }
560 }
561 BitsEstimator::decide_set(&hist)
562 };
563
564 let start_bits = bw.tell();
565 write_slice_header(&mut bw, ftype, self.mbd.q, set_idx, self.deblock, fpts);
566 if ftype == FrameType::I {
567 write_slice_dimensions(&mut bw, self.width, self.height);
568 } else {
569 bw.write1(); // keep dimensions flag
570 }
571 write_slice_mb_idx(&mut bw, slice_start_mb, num_mbs);
572
573 mb_idx = slice_start_mb;
574 let mut skip_count = 0;
575 self.cset.init(self.mbd.q, set_idx);
576 while let Some(mb) = self.mbs.pop_front() {
577 if bw.tell() > start_bits + (self.slice_bits as usize) {
578 break;
579 }
580 let mb_x = mb_idx % self.mb_w;
581 let mb_y = mb_idx / self.mb_w;
582 self.sstate.has_t = mb_idx >= slice_start_mb + self.mb_w;
583 self.sstate.has_l = (mb_idx > slice_start_mb) && (mb_x > 0);
584 self.sstate.has_tl = (mb_idx > slice_start_mb + self.mb_w) && (mb_x > 0);
585 self.sstate.has_tr = (mb_idx >= slice_start_mb + self.mb_w - 1) && (mb_x + 1 < self.mb_w);
586 self.sstate.mb_x = mb_x;
587 self.sstate.mb_y = mb_y;
588 if mb.mb_type.is_skip() {
589 skip_count += 1;
590 } else {
591 if skip_count > 0 {
592 write_skip_count(&mut bw, skip_count);
593 skip_count = 0;
594 } else if ftype != FrameType::I {
595 bw.write1(); // zero skip count
596 }
597
598 write_mb_header(&mut bw, ftype, &self.sstate, &self.mbstate);
599 self.cset.set_params(&mb.mb_type);
600 self.cset.write_coeffs(&mut bw, &mb.coeffs);
601 }
602 mb_idx += 1;
603 }
604 self.mbs.clear();
605 if skip_count > 0 {
606 write_skip_count(&mut bw, skip_count);
607 }
608 while (bw.tell() & 7) != 0 {
609 bw.write0();
610 }
611 if self.debug_log.is_set(DEBUG_BIT_SLICE_SIZE) {
612 println!(" slice {}..{} wrote {} bits / estimated {} bits", slice_start_mb, mb_idx, bw.tell(), est_bits);
613 }
614 dvec = bw.end();
615 nslices += 1;
616 }
617 for _ in 0..(nslices * 8 + 1) {
618 dvec.insert(0, 0);
619 }
620 dvec[0] = (nslices - 1) as u8;
621 for (i, &off) in slice_starts.iter().enumerate() {
622 dvec[i * 8 + 4] = 1;
623 write_u32be(&mut dvec[i * 8 + 5..], off as u32)?;
624 }
625 if self.debug_log.is_set(DEBUG_BIT_RATECTL) {
626 println!(" got frame size {}", dvec.len());
627 }
628
629 if is_ref_frame && self.deblock {
630 loop_filter_frame(&mut recon_frm, &self.dblk, self.mb_w, self.mb_h);
631 }
632
633 if self.debug_log.is_set(DEBUG_BIT_PSNR) {
e6aaad5c 634 let psnr = calc_psnr(vbuf, &rvbuf);
4965a5e5
KS
635 println!(" encoded frame PSNR {} size {}", psnr, dvec.len());
636 }
637
638 if is_ref_frame {
639 std::mem::swap(&mut self.ref_p, &mut self.ref_n);
640 self.ref_n = rvbuf;
641
642 self.p_pts = self.n_pts;
643 self.n_pts = pts;
644
645 self.mbstate.swap_mvs();
646 }
647
648 if is_ref_frame {
649 if self.last_k_ts > self.last_b_ts {
650 self.last_b_ts = self.last_k_ts;
651 }
652 self.last_k_ts = pts;
653 self.fce.update_ref();
654 } else {
655 self.last_b_ts = pts;
656 }
657
658 self.brc.update_stats(ftype, dvec.len(), ts_diff);
659
660 Ok(NAPacket::new(self.stream.clone().unwrap(), frm.ts, ftype == FrameType::I, dvec))
661 } else {
662 Err(EncoderError::InvalidParameters)
663 }
664 }
665}
666
667fn calc_psnr(pic1: &NAVideoBuffer<u8>, pic2: &NAVideoBuffer<u8>) -> f64 {
668 let data1 = pic1.get_data();
669 let data2 = pic2.get_data();
670 let mut sum = 0u64;
671 let mut size = 0;
672 for comp in 0..3 {
673 let (w, h) = pic1.get_dimensions(comp);
674 size += w * h;
675 for (line1, line2) in data1[pic1.get_offset(comp)..].chunks(pic1.get_stride(comp)).zip(
676 data2[pic2.get_offset(comp)..].chunks(pic2.get_stride(comp))).take(h) {
677 for (&pix1, &pix2) in line1[..w].iter().zip(line2.iter()) {
e6aaad5c 678 let diff = (i32::from(pix1) - i32::from(pix2)).unsigned_abs();
4965a5e5
KS
679 sum += u64::from(diff * diff);
680 }
681 }
682 }
683 if size > 0 {
684 48.13080360867910341240 - 10.0 * ((sum as f64) / (size as f64)).log10()
685 } else {
686 std::f64::INFINITY
687 }
688}
689
690impl NAEncoder for RV40Encoder {
691 fn negotiate_format(&self, encinfo: &EncodeParameters) -> EncoderResult<EncodeParameters> {
692 match encinfo.format {
693 NACodecTypeInfo::None => {
694 Ok(EncodeParameters {
695 format: NACodecTypeInfo::Video(NAVideoInfo::new(0, 0, false, YUV420_FORMAT)),
696 ..Default::default() })
697 },
698 NACodecTypeInfo::Audio(_) => Err(EncoderError::FormatError),
699 NACodecTypeInfo::Video(vinfo) => {
700 let outinfo = NAVideoInfo::new((vinfo.width + 15) & !15, (vinfo.height + 15) & !15, false, YUV420_FORMAT);
701 let mut ofmt = *encinfo;
702 ofmt.format = NACodecTypeInfo::Video(outinfo);
703 Ok(ofmt)
704 }
705 }
706 }
707 fn get_capabilities(&self) -> u64 { 0 }
708 fn init(&mut self, stream_id: u32, encinfo: EncodeParameters) -> EncoderResult<NAStreamRef> {
709 match encinfo.format {
710 NACodecTypeInfo::None => Err(EncoderError::FormatError),
711 NACodecTypeInfo::Audio(_) => Err(EncoderError::FormatError),
712 NACodecTypeInfo::Video(vinfo) => {
713 if vinfo.format != YUV420_FORMAT {
714 return Err(EncoderError::FormatError);
715 }
716 if ((vinfo.width | vinfo.height) & 15) != 0 {
717 return Err(EncoderError::FormatError);
718 }
719 if (vinfo.width | vinfo.height) >= (1 << 12) {
720 return Err(EncoderError::FormatError);
721 }
722
723 // 32-bit flags (VBR, bframes, slices, something else) and 32-bit version
724 let edata = vec![0x01, 0x08, 0x10, 0x20, 0x40, 0x00, 0x80, 0x00];
725 let out_info = NAVideoInfo::new(vinfo.width, vinfo.height, false, vinfo.format);
726 let info = NACodecInfo::new("realvideo4", NACodecTypeInfo::Video(out_info), Some(edata));
727 let mut stream = NAStream::new(StreamType::Video, stream_id, info, encinfo.tb_num, encinfo.tb_den, 0);
728 stream.set_num(stream_id as usize);
729 let stream = stream.into_ref();
730
731 self.stream = Some(stream.clone());
732
733 self.width = vinfo.width;
734 self.height = vinfo.height;
735 self.mb_w = (vinfo.width + 15) >> 4;
736 self.mb_h = (vinfo.height + 15) >> 4;
737
738 if self.mb_w * self.mb_h > 9216 {
739 return Err(EncoderError::FormatError);
740 }
741
742 if (1..=100u8).contains(&encinfo.quality) {
743 self.brc.set_force_quality(Some(encinfo.quality));
744 } else {
745 self.brc.set_force_quality(None);
746 }
747 self.brc.set_bitrate(encinfo.bitrate);
748
749 self.vinfo = out_info;
750 let max_frames = self.order.get_max_grp_len();
751 self.frm_pool.set_dec_bufs(max_frames + 3);
752 self.max_grp_bufs = max_frames;
753 self.needs_alloc = true;
754
755 self.fce.resize(self.width, self.height);
756 self.mbstate.resize(self.mb_w, self.mb_h);
757 self.mbd.resize(self.mb_w);
758 self.dblk.resize(self.mb_w * self.mb_h, DeblockInfo::default());
759
760 Ok(stream)
761 },
762 }
763 }
764 fn encode(&mut self, frm: &NAFrame) -> EncoderResult<()> {
765 if self.needs_alloc {
766 self.frm_pool.prealloc_video(self.vinfo, 4)?;
767 self.ref_n = self.frm_pool.get_free().unwrap();
768 self.ref_p = self.frm_pool.get_free().unwrap();
769 self.needs_alloc = false;
770 }
771 if let Some(ref vbuf) = frm.get_buffer().get_vbuf() {
772 if let Some(dbuf) = self.frm_pool.get_copy(vbuf) {
773 let newfrm = NAFrame::new(frm.ts, frm.frame_type, frm.key, frm.get_info(), NABufferType::Video(dbuf));
774 self.qframes.push(newfrm);
775
776 loop {
777 let (mut ftype, mut frame_pos) = self.order.peek_next_frame();
778 if frame_pos >= self.qframes.len() {
779 break;
780 }
781
782 if ftype == FrameType::Other {
783 if self.qframes.len() < 2 {
784 return Err(EncoderError::Bug);
785 }
786 if let (Some(ref frm1), Some(ref frm2)) = (self.qframes[0].get_buffer().get_vbuf(), self.qframes[1].get_buffer().get_vbuf()) {
787 let is_b = self.fce.decide_b_frame(frm1, frm2);
788 ftype = if is_b {
789 frame_pos = 1;
790 FrameType::B
791 } else {
792 frame_pos = 0;
793 FrameType::P
794 };
795 } else {
796 return Err(EncoderError::Bug);
797 }
798 self.order.update(ftype);
799 }
800
801 let frm = self.qframes.remove(frame_pos);
802 let pkt = self.encode_frame(frm, frame_pos)?;
803 self.pkts.push_back(pkt);
804 }
805 Ok(())
806 } else {
807 Err(EncoderError::AllocError)
808 }
809 } else {
810 Err(EncoderError::FormatError)
811 }
812 }
813 fn get_packet(&mut self) -> EncoderResult<Option<NAPacket>> {
814 Ok(self.pkts.pop_front())
815 }
816 fn flush(&mut self) -> EncoderResult<()> {
817 Ok(())
818 }
819}
820
821const DEBUG_LOG_OPTION: &str = "debug";
822const SLICE_SIZE_OPTION: &str = "slice_size";
823const FRAME_ORDER_OPTION: &str = "frame_order";
824const DEBLOCK_OPTION: &str = "loop_filt";
825const QUANT_OPTION: &str = "quant";
826const QUALITY_OPTION: &str = "quality";
827const SET_OPTION: &str = "coding_set";
828const SEARCH_MODE_OPTION: &str = "me_mode";
829const SEARCH_RANGE_OPTION: &str = "me_range";
830const SEARCH_THR_OPTION: &str = "me_thr";
831const B_REFINE_OPTION: &str = "refine_b";
832const I4_IN_B_OPTION: &str = "i4_in_b";
833const B_OFFSET_OPTION: &str = "b_offset";
834
835const ENCODER_OPTS: &[NAOptionDefinition] = &[
836 NAOptionDefinition {
837 name: DEBUG_LOG_OPTION, description: "debug flags",
838 opt_type: NAOptionDefinitionType::String(None) },
839 NAOptionDefinition {
840 name: SLICE_SIZE_OPTION, description: "soft slice size limit in bits",
841 opt_type: NAOptionDefinitionType::Int(Some(4096), Some(100000)) },
842 NAOptionDefinition {
843 name: FRAME_ORDER_OPTION, description: "frame order (e.g. IBBPBB)",
844 opt_type: NAOptionDefinitionType::String(None) },
845 NAOptionDefinition {
846 name: DEBLOCK_OPTION, description: "in-loop filter",
847 opt_type: NAOptionDefinitionType::Bool },
848 NAOptionDefinition {
849 name: QUANT_OPTION, description: "force quantiser (-1 = none)",
850 opt_type: NAOptionDefinitionType::Int(Some(-1), Some(31)) },
851 NAOptionDefinition {
852 name: QUALITY_OPTION, description: "force quality (-1 = none)",
853 opt_type: NAOptionDefinitionType::Int(Some(-1), Some(100)) },
854 NAOptionDefinition {
855 name: SET_OPTION, description: "force coding set (-1 = none)",
856 opt_type: NAOptionDefinitionType::Int(Some(-1), Some(3)) },
857 NAOptionDefinition {
858 name: SEARCH_MODE_OPTION, description: "motion search mode",
859 opt_type: NAOptionDefinitionType::String(Some(MVSearchMode::get_possible_modes())) },
860 NAOptionDefinition {
861 name: SEARCH_RANGE_OPTION, description: "motion search range",
862 opt_type: NAOptionDefinitionType::Int(Some(0), Some(256)) },
863 NAOptionDefinition {
864 name: SEARCH_THR_OPTION, description: "motion search cut-off threshold",
865 opt_type: NAOptionDefinitionType::Int(Some(0), Some(1048576)) },
866 NAOptionDefinition {
867 name: B_REFINE_OPTION, description: "better ME for B-frames",
868 opt_type: NAOptionDefinitionType::Bool },
869 NAOptionDefinition {
870 name: I4_IN_B_OPTION, description: "allow intra 4x4 coding in B-frames",
871 opt_type: NAOptionDefinitionType::Bool },
872 NAOptionDefinition {
873 name: B_OFFSET_OPTION, description: "B-frame quantiser offset",
874 opt_type: NAOptionDefinitionType::Int(Some(0), Some(16)) },
875];
876
877impl NAOptionHandler for RV40Encoder {
878 fn get_supported_options(&self) -> &[NAOptionDefinition] { ENCODER_OPTS }
879 fn set_options(&mut self, options: &[NAOption]) {
880 for option in options.iter() {
881 for opt_def in ENCODER_OPTS.iter() {
882 if opt_def.check(option).is_ok() {
883 match option.name {
884 DEBUG_LOG_OPTION => {
885 if let NAValue::String(ref strval) = option.value {
886 self.debug_log.parse(strval);
887 }
888 },
889 SLICE_SIZE_OPTION => {
890 if let NAValue::Int(intval) = option.value {
891 self.slice_bits = intval as u32;
892 }
893 },
894 FRAME_ORDER_OPTION => {
895 if let NAValue::String(ref strval) = option.value {
896 if let Ok(norder) = strval.parse::<FrameOrder>() {
897 self.order = norder;
898 let max_frames = self.order.get_max_grp_len();
899 if max_frames > self.max_grp_bufs {
900 self.frm_pool.set_dec_bufs(max_frames + 3);
901 self.needs_alloc = true;
902 self.max_grp_bufs = max_frames;
903 }
904 } else {
905 println!("Invalid order sequence");
906 }
907 }
908 },
909 DEBLOCK_OPTION => {
910 if let NAValue::Bool(val) = option.value {
911 self.deblock = val;
912 }
913 },
914 QUANT_OPTION => {
915 if let NAValue::Int(val) = option.value {
916 if val != -1 {
917 self.brc.set_force_quant(Some(val as usize));
918 } else {
919 self.brc.set_force_quant(None);
920 }
921 }
922 },
923 QUALITY_OPTION => {
924 if let NAValue::Int(val) = option.value {
925 if val != -1 {
926 self.brc.set_force_quality(Some(val as u8));
927 } else {
928 self.brc.set_force_quality(None);
929 }
930 }
931 },
932 SET_OPTION => {
933 if let NAValue::Int(val) = option.value {
934 self.force_set = if val != -1 { Some(val as usize) } else { None };
935 }
936 },
937 SEARCH_MODE_OPTION => {
938 if let NAValue::String(ref strval) = option.value {
939 if let Ok(mmode) = strval.parse::<MVSearchMode>() {
940 self.me.set_mode(mmode);
941 } else {
942 println!("Invalid mode");
943 }
944 }
945 },
946 SEARCH_RANGE_OPTION => {
947 if let NAValue::Int(val) = option.value {
948 self.me.range = val as i16;
949 }
950 },
951 SEARCH_THR_OPTION => {
952 if let NAValue::Int(val) = option.value {
953 self.me.thresh = val as u32;
954 }
955 },
956 B_REFINE_OPTION => {
957 if let NAValue::Bool(val) = option.value {
958 self.refine_b = val;
959 }
960 },
961 I4_IN_B_OPTION => {
962 if let NAValue::Bool(val) = option.value {
963 self.i4_in_b = val;
964 }
965 },
966 B_OFFSET_OPTION => {
967 if let NAValue::Int(val) = option.value {
968 self.brc.b_offset = val as usize;
969 }
970 },
971 _ => {},
972 };
973 }
974 }
975 }
976 }
977 fn query_option_value(&self, name: &str) -> Option<NAValue> {
978 match name {
979 DEBUG_LOG_OPTION => Some(NAValue::String(self.debug_log.to_string())),
980 SLICE_SIZE_OPTION => Some(NAValue::Int(self.slice_bits as i64)),
981 FRAME_ORDER_OPTION => Some(NAValue::String(self.order.to_string())),
982 DEBLOCK_OPTION => Some(NAValue::Bool(self.deblock)),
983 QUANT_OPTION => Some(NAValue::Int(self.brc.get_force_quant().into())),
984 QUALITY_OPTION => Some(NAValue::Int(self.brc.get_force_quality().into())),
985 SET_OPTION => Some(NAValue::Int(if let Some(set) = self.force_set { set as i64 } else { -1 })),
986 SEARCH_MODE_OPTION => Some(NAValue::String(self.me.get_mode().to_string())),
987 SEARCH_THR_OPTION => Some(NAValue::Int(self.me.thresh.into())),
988 SEARCH_RANGE_OPTION => Some(NAValue::Int(self.me.range.into())),
989 B_REFINE_OPTION => Some(NAValue::Bool(self.refine_b)),
990 I4_IN_B_OPTION => Some(NAValue::Bool(self.i4_in_b)),
991 B_OFFSET_OPTION => Some(NAValue::Int(self.brc.b_offset as i64)),
992 _ => None,
993 }
994 }
995}
996
997pub fn get_encoder() -> Box<dyn NAEncoder + Send> {
998 Box::new(RV40Encoder::new())
999}
1000
1001#[cfg(test)]
1002mod test {
1003 use nihav_core::codecs::*;
1004 use nihav_core::demuxers::*;
1005 use nihav_core::muxers::*;
1006 use crate::*;
1007 use nihav_codec_support::test::enc_video::*;
1008 use nihav_commonfmt::*;
1009
1010 #[allow(unused_variables)]
1011 fn encode_test(out_name: &'static str, enc_options: &[NAOption], limit: Option<u64>, hash: &[u32; 4]) {
1012 let mut dmx_reg = RegisteredDemuxers::new();
1013 generic_register_all_demuxers(&mut dmx_reg);
1014 let mut dec_reg = RegisteredDecoders::new();
1015 generic_register_all_decoders(&mut dec_reg);
1016 let mut mux_reg = RegisteredMuxers::new();
1017 realmedia_register_all_muxers(&mut mux_reg);
1018 let mut enc_reg = RegisteredEncoders::new();
1019 realmedia_register_all_encoders(&mut enc_reg);
1020
1021 // sample from private collection
1022 let dec_config = DecoderTestParams {
1023 demuxer: "yuv4mpeg",
1024 in_name: "assets/day3b.y4m",
1025 stream_type: StreamType::Video,
1026 limit,
1027 dmx_reg, dec_reg,
1028 };
1029 let enc_config = EncoderTestParams {
1030 muxer: "realmedia",
1031 enc_name: "realvideo4",
1032 out_name,
1033 mux_reg, enc_reg,
1034 };
1035 let dst_vinfo = NAVideoInfo {
1036 width: 0,
1037 height: 0,
1038 format: YUV420_FORMAT,
1039 flipped: false,
1040 bits: 12,
1041 };
1042 let enc_params = EncodeParameters {
1043 format: NACodecTypeInfo::Video(dst_vinfo),
1044 quality: 0,
1045 bitrate: 300000,
1046 tb_num: 0,
1047 tb_den: 0,
1048 flags: 0,
1049 };
1050 //test_encoding_to_file(&dec_config, &enc_config, enc_params, enc_options);
1051 test_encoding_md5(&dec_config, &enc_config, enc_params, enc_options,
1052 hash);
1053 }
1054 #[test]
1055 fn test_rv40_encoder_simple() {
1056 let enc_options = &[
1057 NAOption { name: super::FRAME_ORDER_OPTION, value: NAValue::String("I".to_owned()) },
1058 NAOption { name: super::DEBLOCK_OPTION, value: NAValue::Bool(false) },
1059 NAOption { name: super::QUANT_OPTION, value: NAValue::Int(17) },
1060 NAOption { name: super::SEARCH_MODE_OPTION, value: NAValue::String("diamond".to_owned()) },
1061 ];
1062 encode_test("rv40simple.rmvb", enc_options, Some(10), &[0x03b0d743, 0x36c20dbb, 0x18fa1c9e, 0x4b2b7324]);
1063 }
1064 #[test]
1065 fn test_rv40_encoder_ipb() {
1066 let enc_options = &[
1067 NAOption { name: super::FRAME_ORDER_OPTION, value: NAValue::String("IBPB".to_owned()) },
1068 NAOption { name: super::DEBLOCK_OPTION, value: NAValue::Bool(true) },
1069 NAOption { name: super::QUANT_OPTION, value: NAValue::Int(17) },
1070 NAOption { name: super::SEARCH_MODE_OPTION, value: NAValue::String("hexagon".to_owned()) },
1071 ];
1072 encode_test("rv40ipb.rmvb", enc_options, Some(8), &[0xc382ab0b, 0xbcfbb02a, 0xf12a064f, 0xe6a5c2c3]);
1073 }
1074 #[test]
1075 fn test_rv40_encoder_advanced() {
1076 let enc_options = &[
1077 NAOption { name: super::FRAME_ORDER_OPTION, value: NAValue::String("dynamic".to_owned()) },
1078 NAOption { name: super::DEBLOCK_OPTION, value: NAValue::Bool(true) },
1079 NAOption { name: super::SEARCH_MODE_OPTION, value: NAValue::String("umh".to_owned()) },
1080 ];
1081 encode_test("rv40adv.rmvb", enc_options, Some(8), &[0xc4395f49, 0x0536d5f0, 0x32406834, 0xb7b634be]);
1082 }
1083}