1#![allow(clippy::cast_lossless)]
9#![allow(clippy::cast_precision_loss)]
10#![allow(clippy::cast_possible_truncation)]
11#![allow(clippy::cast_sign_loss)]
12#![allow(clippy::similar_names)]
13#![allow(clippy::missing_errors_doc)]
14#![allow(clippy::double_must_use)]
15#![allow(clippy::match_same_arms)]
16#![forbid(unsafe_code)]
17
18use crate::frame::FrameType;
19
20#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
22pub enum RateControlMode {
23 Cqp,
26 Cbr,
29 Vbr,
32 Abr,
35 #[default]
38 Crf,
39}
40
41impl RateControlMode {
42 #[must_use]
44 pub fn is_bitrate_based(&self) -> bool {
45 matches!(self, Self::Cbr | Self::Vbr | Self::Abr)
46 }
47
48 #[must_use]
50 pub fn is_quality_based(&self) -> bool {
51 matches!(self, Self::Cqp | Self::Crf)
52 }
53}
54
55#[derive(Clone, Debug)]
57pub struct RcConfig {
58 pub mode: RateControlMode,
60 pub target_bitrate: u64,
62 pub max_bitrate: Option<u64>,
64 pub min_bitrate: Option<u64>,
66 pub min_qp: u8,
68 pub max_qp: u8,
70 pub initial_qp: u8,
72 pub crf: f32,
74 pub buffer_size: u64,
76 pub initial_buffer_fullness: f32,
78 pub framerate_num: u32,
80 pub framerate_den: u32,
82 pub gop_length: u32,
84 pub enable_aq: bool,
86 pub aq_strength: f32,
88 pub lookahead_depth: usize,
90 pub b_qp_offset: i8,
92 pub i_qp_offset: i8,
94 pub scene_cut_detection: bool,
96 pub scene_cut_threshold: f32,
98}
99
100impl Default for RcConfig {
101 fn default() -> Self {
102 Self {
103 mode: RateControlMode::Crf,
104 target_bitrate: 5_000_000, max_bitrate: None,
106 min_bitrate: None,
107 min_qp: 1,
108 max_qp: 63,
109 initial_qp: 28,
110 crf: 23.0,
111 buffer_size: 10_000_000, initial_buffer_fullness: 0.75,
113 framerate_num: 30,
114 framerate_den: 1,
115 gop_length: 250,
116 enable_aq: true,
117 aq_strength: 1.0,
118 lookahead_depth: 40,
119 b_qp_offset: 2,
120 i_qp_offset: -2,
121 scene_cut_detection: true,
122 scene_cut_threshold: 0.4,
123 }
124 }
125}
126
127impl RcConfig {
128 #[must_use]
130 pub fn cqp(qp: u8) -> Self {
131 Self {
132 mode: RateControlMode::Cqp,
133 initial_qp: qp,
134 ..Default::default()
135 }
136 }
137
138 #[must_use]
140 pub fn cbr(bitrate: u64) -> Self {
141 Self {
142 mode: RateControlMode::Cbr,
143 target_bitrate: bitrate,
144 max_bitrate: Some(bitrate),
145 buffer_size: bitrate, ..Default::default()
147 }
148 }
149
150 #[must_use]
152 pub fn vbr(target: u64, max: u64) -> Self {
153 Self {
154 mode: RateControlMode::Vbr,
155 target_bitrate: target,
156 max_bitrate: Some(max),
157 buffer_size: max * 2, ..Default::default()
159 }
160 }
161
162 #[must_use]
164 pub fn crf(crf_value: f32) -> Self {
165 Self {
166 mode: RateControlMode::Crf,
167 crf: crf_value.clamp(0.0, 63.0),
168 ..Default::default()
169 }
170 }
171
172 #[must_use]
174 pub fn framerate(&self) -> f64 {
175 if self.framerate_den == 0 {
176 30.0
177 } else {
178 f64::from(self.framerate_num) / f64::from(self.framerate_den)
179 }
180 }
181
182 #[must_use]
184 pub fn target_bits_per_frame(&self) -> u64 {
185 let fps = self.framerate();
186 if fps <= 0.0 {
187 return 0;
188 }
189 (self.target_bitrate as f64 / fps) as u64
190 }
191
192 #[must_use]
194 pub fn validate(&self) -> Result<(), RcConfigError> {
195 if self.min_qp > self.max_qp {
196 return Err(RcConfigError::InvalidQpRange);
197 }
198 if self.initial_qp < self.min_qp || self.initial_qp > self.max_qp {
199 return Err(RcConfigError::InitialQpOutOfRange);
200 }
201 if self.crf < 0.0 || self.crf > 63.0 {
202 return Err(RcConfigError::InvalidCrf);
203 }
204 if self.framerate_den == 0 {
205 return Err(RcConfigError::InvalidFramerate);
206 }
207 if self.mode.is_bitrate_based() && self.target_bitrate == 0 {
208 return Err(RcConfigError::ZeroBitrate);
209 }
210 Ok(())
211 }
212}
213
214#[derive(Clone, Copy, Debug, PartialEq, Eq)]
216pub enum RcConfigError {
217 InvalidQpRange,
219 InitialQpOutOfRange,
221 InvalidCrf,
223 InvalidFramerate,
225 ZeroBitrate,
227}
228
229#[derive(Clone, Debug, Default)]
231pub struct FrameStats {
232 pub frame_num: u64,
234 pub frame_type: FrameType,
236 pub bits: u64,
238 pub qp: u8,
240 pub qp_f: f32,
242 pub spatial_complexity: f32,
244 pub temporal_complexity: f32,
246 pub complexity: f32,
248 pub psnr: Option<f32>,
250 pub ssim: Option<f32>,
252 pub scene_cut: bool,
254 pub target_bits: u64,
256 pub bit_accuracy: f32,
258 pub encode_time_us: u64,
260}
261
262impl FrameStats {
263 #[must_use]
265 pub fn new(frame_num: u64, frame_type: FrameType) -> Self {
266 Self {
267 frame_num,
268 frame_type,
269 ..Default::default()
270 }
271 }
272
273 #[must_use]
275 pub fn bits_per_pixel(&self, width: u32, height: u32) -> f32 {
276 let pixels = u64::from(width) * u64::from(height);
277 if pixels == 0 {
278 return 0.0;
279 }
280 self.bits as f32 / pixels as f32
281 }
282
283 #[must_use]
285 pub fn exceeded_target(&self, threshold: f32) -> bool {
286 self.target_bits > 0 && self.bit_accuracy > (1.0 + threshold)
287 }
288
289 #[must_use]
291 pub fn under_target(&self, threshold: f32) -> bool {
292 self.target_bits > 0 && self.bit_accuracy < (1.0 - threshold)
293 }
294}
295
296#[derive(Clone, Debug, Default)]
298pub struct GopStats {
299 pub gop_index: u64,
301 pub frame_count: u32,
303 pub i_frame_count: u32,
305 pub p_frame_count: u32,
307 pub b_frame_count: u32,
309 pub total_bits: u64,
311 pub target_bits: u64,
313 pub average_qp: f32,
315 pub average_complexity: f32,
317 pub total_complexity: f32,
319 pub first_frame: u64,
321 pub last_frame: u64,
323 frames: Vec<FrameStats>,
325}
326
327impl GopStats {
328 #[must_use]
330 pub fn new(gop_index: u64, first_frame: u64) -> Self {
331 Self {
332 gop_index,
333 first_frame,
334 last_frame: first_frame,
335 ..Default::default()
336 }
337 }
338
339 pub fn add_frame(&mut self, stats: FrameStats) {
341 self.last_frame = stats.frame_num;
342 self.total_bits += stats.bits;
343 self.total_complexity += stats.complexity;
344
345 match stats.frame_type {
346 FrameType::Key => self.i_frame_count += 1,
347 FrameType::Inter => self.p_frame_count += 1,
348 FrameType::BiDir => self.b_frame_count += 1,
349 FrameType::Switch => self.p_frame_count += 1,
350 }
351
352 self.frames.push(stats);
353 self.frame_count = self.frames.len() as u32;
354
355 if self.frame_count > 0 {
357 let fc = self.frame_count as f32;
358 self.average_qp = self.frames.iter().map(|f| f.qp_f).sum::<f32>() / fc;
359 self.average_complexity = self.total_complexity / fc;
360 }
361 }
362
363 #[must_use]
365 pub fn frames(&self) -> &[FrameStats] {
366 &self.frames
367 }
368
369 #[must_use]
371 pub fn bit_accuracy(&self) -> f32 {
372 if self.target_bits == 0 {
373 return 1.0;
374 }
375 self.total_bits as f32 / self.target_bits as f32
376 }
377
378 #[must_use]
380 pub fn bits_per_i_frame(&self) -> f64 {
381 if self.i_frame_count == 0 {
382 return 0.0;
383 }
384 let i_bits: u64 = self
385 .frames
386 .iter()
387 .filter(|f| f.frame_type == FrameType::Key)
388 .map(|f| f.bits)
389 .sum();
390 i_bits as f64 / f64::from(self.i_frame_count)
391 }
392
393 #[must_use]
395 pub fn bits_per_p_frame(&self) -> f64 {
396 if self.p_frame_count == 0 {
397 return 0.0;
398 }
399 let p_bits: u64 = self
400 .frames
401 .iter()
402 .filter(|f| f.frame_type == FrameType::Inter)
403 .map(|f| f.bits)
404 .sum();
405 p_bits as f64 / f64::from(self.p_frame_count)
406 }
407
408 #[must_use]
410 pub fn bits_per_b_frame(&self) -> f64 {
411 if self.b_frame_count == 0 {
412 return 0.0;
413 }
414 let b_bits: u64 = self
415 .frames
416 .iter()
417 .filter(|f| f.frame_type == FrameType::BiDir)
418 .map(|f| f.bits)
419 .sum();
420 b_bits as f64 / f64::from(self.b_frame_count)
421 }
422}
423
424#[derive(Clone, Debug, Default)]
426pub struct RcOutput {
427 pub qp: u8,
429 pub qp_f: f32,
431 pub target_bits: u64,
433 pub min_bits: u64,
435 pub max_bits: u64,
437 pub drop_frame: bool,
439 pub force_keyframe: bool,
441 pub qp_offsets: Option<Vec<f32>>,
443 pub lambda: f64,
445 pub lambda_me: f64,
447}
448
449impl RcOutput {
450 #[must_use]
452 pub fn with_qp(qp: u8) -> Self {
453 Self {
454 qp,
455 qp_f: qp as f32,
456 ..Default::default()
457 }
458 }
459
460 #[must_use]
462 pub fn drop() -> Self {
463 Self {
464 drop_frame: true,
465 ..Default::default()
466 }
467 }
468
469 #[must_use]
472 pub fn qp_to_lambda(qp: f32) -> f64 {
473 0.85 * 2.0_f64.powf((f64::from(qp) - 12.0) / 3.0)
474 }
475
476 #[must_use]
478 pub fn lambda_to_lambda_me(lambda: f64) -> f64 {
479 lambda.sqrt()
480 }
481
482 pub fn compute_lambda(&mut self) {
484 self.lambda = Self::qp_to_lambda(self.qp_f);
485 self.lambda_me = Self::lambda_to_lambda_me(self.lambda);
486 }
487}
488
489#[derive(Clone, Debug, Default)]
491pub struct RcState {
492 pub frames_encoded: u64,
494 pub total_bits: u64,
496 pub buffer_level: i64,
498 pub average_bitrate: f64,
500 pub average_qp: f32,
502 pub frames_dropped: u64,
504 pub scene_cuts_detected: u64,
506 pub current_gop: u64,
508}
509
510impl RcState {
511 pub fn update(&mut self, stats: &FrameStats, elapsed_time: f64) {
513 self.frames_encoded += 1;
514 self.total_bits += stats.bits;
515
516 if elapsed_time > 0.0 {
517 self.average_bitrate = self.total_bits as f64 / elapsed_time;
518 }
519
520 let n = self.frames_encoded as f32;
522 self.average_qp = self.average_qp * (n - 1.0) / n + stats.qp_f / n;
523
524 if stats.scene_cut {
525 self.scene_cuts_detected += 1;
526 }
527 }
528}
529
530#[cfg(test)]
531mod tests {
532 use super::*;
533
534 #[test]
535 fn test_rate_control_mode() {
536 assert!(RateControlMode::Cbr.is_bitrate_based());
537 assert!(RateControlMode::Vbr.is_bitrate_based());
538 assert!(RateControlMode::Abr.is_bitrate_based());
539 assert!(!RateControlMode::Cqp.is_bitrate_based());
540 assert!(!RateControlMode::Crf.is_bitrate_based());
541
542 assert!(RateControlMode::Cqp.is_quality_based());
543 assert!(RateControlMode::Crf.is_quality_based());
544 assert!(!RateControlMode::Cbr.is_quality_based());
545 }
546
547 #[test]
548 fn test_rc_config_creation() {
549 let config = RcConfig::cqp(28);
550 assert_eq!(config.mode, RateControlMode::Cqp);
551 assert_eq!(config.initial_qp, 28);
552
553 let config = RcConfig::cbr(5_000_000);
554 assert_eq!(config.mode, RateControlMode::Cbr);
555 assert_eq!(config.target_bitrate, 5_000_000);
556
557 let config = RcConfig::vbr(5_000_000, 10_000_000);
558 assert_eq!(config.mode, RateControlMode::Vbr);
559 assert_eq!(config.max_bitrate, Some(10_000_000));
560
561 let config = RcConfig::crf(23.0);
562 assert_eq!(config.mode, RateControlMode::Crf);
563 assert!((config.crf - 23.0).abs() < f32::EPSILON);
564 }
565
566 #[test]
567 fn test_rc_config_validation() {
568 let mut config = RcConfig::default();
569 assert!(config.validate().is_ok());
570
571 config.min_qp = 50;
572 config.max_qp = 30;
573 assert_eq!(config.validate(), Err(RcConfigError::InvalidQpRange));
574
575 config.min_qp = 10;
576 config.max_qp = 40;
577 config.initial_qp = 5;
578 assert_eq!(config.validate(), Err(RcConfigError::InitialQpOutOfRange));
579
580 config.initial_qp = 25;
581 config.crf = 100.0;
582 assert_eq!(config.validate(), Err(RcConfigError::InvalidCrf));
583
584 config.crf = 23.0;
585 config.framerate_den = 0;
586 assert_eq!(config.validate(), Err(RcConfigError::InvalidFramerate));
587
588 config.framerate_den = 1;
589 config.mode = RateControlMode::Cbr;
590 config.target_bitrate = 0;
591 assert_eq!(config.validate(), Err(RcConfigError::ZeroBitrate));
592 }
593
594 #[test]
595 fn test_target_bits_per_frame() {
596 let config = RcConfig {
597 target_bitrate: 3_000_000,
598 framerate_num: 30,
599 framerate_den: 1,
600 ..Default::default()
601 };
602 assert_eq!(config.target_bits_per_frame(), 100_000);
603
604 let config = RcConfig {
605 target_bitrate: 6_000_000,
606 framerate_num: 60,
607 framerate_den: 1,
608 ..Default::default()
609 };
610 assert_eq!(config.target_bits_per_frame(), 100_000);
611 }
612
613 #[test]
614 fn test_frame_stats() {
615 let mut stats = FrameStats::new(0, FrameType::Key);
616 stats.bits = 100_000;
617 stats.target_bits = 80_000;
618 stats.bit_accuracy = stats.bits as f32 / stats.target_bits as f32;
619
620 assert!(stats.exceeded_target(0.1));
621 assert!(!stats.under_target(0.1));
622
623 let bpp = stats.bits_per_pixel(1920, 1080);
624 assert!(bpp > 0.0);
625 }
626
627 #[test]
628 fn test_gop_stats() {
629 let mut gop = GopStats::new(0, 0);
630
631 let mut i_frame = FrameStats::new(0, FrameType::Key);
632 i_frame.bits = 200_000;
633 i_frame.qp_f = 24.0;
634 i_frame.complexity = 1.5;
635 gop.add_frame(i_frame);
636
637 let mut p_frame = FrameStats::new(1, FrameType::Inter);
638 p_frame.bits = 50_000;
639 p_frame.qp_f = 26.0;
640 p_frame.complexity = 1.0;
641 gop.add_frame(p_frame);
642
643 assert_eq!(gop.frame_count, 2);
644 assert_eq!(gop.i_frame_count, 1);
645 assert_eq!(gop.p_frame_count, 1);
646 assert_eq!(gop.total_bits, 250_000);
647 assert!((gop.average_qp - 25.0).abs() < f32::EPSILON);
648 }
649
650 #[test]
651 fn test_rc_output_lambda() {
652 let lambda = RcOutput::qp_to_lambda(28.0);
653 assert!(lambda > 0.0);
654
655 let lambda_me = RcOutput::lambda_to_lambda_me(lambda);
656 assert!((lambda_me - lambda.sqrt()).abs() < f64::EPSILON);
657
658 let mut output = RcOutput::with_qp(28);
659 output.compute_lambda();
660 assert!(output.lambda > 0.0);
661 assert!(output.lambda_me > 0.0);
662 }
663
664 #[test]
665 fn test_rc_state_update() {
666 let mut state = RcState::default();
667 let mut stats = FrameStats::new(0, FrameType::Key);
668 stats.bits = 100_000;
669 stats.qp_f = 28.0;
670
671 state.update(&stats, 1.0 / 30.0);
672
673 assert_eq!(state.frames_encoded, 1);
674 assert_eq!(state.total_bits, 100_000);
675 assert!((state.average_qp - 28.0).abs() < f32::EPSILON);
676 }
677}