Skip to main content

oximedia_codec/rate_control/
mod.rs

1// Clippy allows for common pedantic lints in rate control math
2#![allow(clippy::cast_lossless)]
3#![allow(clippy::cast_precision_loss)]
4#![allow(clippy::cast_possible_truncation)]
5#![allow(clippy::cast_sign_loss)]
6#![allow(clippy::similar_names)]
7#![allow(clippy::unused_self)]
8
9//! Rate control module for video encoders.
10//!
11//! This module provides comprehensive rate control algorithms for video encoding:
12//!
13//! - **CQP** (Constant QP): Fixed quantization parameter per frame type
14//! - **CBR** (Constant Bitrate): Maintains steady bitrate with buffer model
15//! - **VBR** (Variable Bitrate): Variable bitrate with optional two-pass
16//! - **CRF** (Constant Rate Factor): Quality-based rate control
17//!
18//! # Architecture
19//!
20//! The rate control system consists of several interconnected components:
21//!
22//! ```text
23//! ┌─────────────────────────────────────────────────────────────────┐
24//! │                        Rate Controller                          │
25//! │  ┌─────────────────────────────────────────────────────────────┐ │
26//! │  │                      Lookahead                               │ │
27//! │  │  ┌──────────────┐  ┌──────────────┐  ┌──────────────────┐   │ │
28//! │  │  │ Scene Change │  │ Complexity   │  │ Mini-GOP         │   │ │
29//! │  │  │ Detection    │  │ Estimation   │  │ Structure        │   │ │
30//! │  │  └──────────────┘  └──────────────┘  └──────────────────┘   │ │
31//! │  └─────────────────────────────────────────────────────────────┘ │
32//! │  ┌─────────────────────────────────────────────────────────────┐ │
33//! │  │                    Rate Controllers                          │ │
34//! │  │  ┌─────┐  ┌─────┐  ┌─────┐  ┌─────┐                         │ │
35//! │  │  │ CQP │  │ CBR │  │ VBR │  │ CRF │                         │ │
36//! │  │  └─────┘  └─────┘  └─────┘  └─────┘                         │ │
37//! │  └─────────────────────────────────────────────────────────────┘ │
38//! │  ┌─────────────────────────────────────────────────────────────┐ │
39//! │  │                    Support Systems                           │ │
40//! │  │  ┌──────────────┐  ┌──────────────┐                         │ │
41//! │  │  │ Rate Buffer  │  │ Adaptive     │                         │ │
42//! │  │  │ (HRD)        │  │ Quantization │                         │ │
43//! │  │  └──────────────┘  └──────────────┘                         │ │
44//! │  └─────────────────────────────────────────────────────────────┘ │
45//! └─────────────────────────────────────────────────────────────────┘
46//! ```
47//!
48//! # Usage
49//!
50//! ## Simple CQP Encoding
51//!
52//! ```ignore
53//! use oximedia_codec::rate_control::{CqpController, RcConfig};
54//!
55//! let config = RcConfig::cqp(28);
56//! let mut controller = CqpController::new(&config);
57//!
58//! // Get QP for each frame
59//! let output = controller.get_qp(FrameType::Key);
60//! encoder.set_qp(output.qp);
61//! ```
62//!
63//! ## CBR Streaming
64//!
65//! ```ignore
66//! use oximedia_codec::rate_control::{CbrController, RcConfig};
67//!
68//! let config = RcConfig::cbr(5_000_000); // 5 Mbps
69//! let mut controller = CbrController::new(&config);
70//!
71//! // For each frame
72//! let output = controller.get_rc(frame_type);
73//! if !output.drop_frame {
74//!     encoder.set_qp(output.qp);
75//!     // ... encode frame ...
76//!     controller.update(&stats);
77//! }
78//! ```
79//!
80//! ## Quality-Based CRF
81//!
82//! ```ignore
83//! use oximedia_codec::rate_control::{CrfController, RcConfig};
84//!
85//! let config = RcConfig::crf(23.0);
86//! let mut controller = CrfController::new(&config);
87//!
88//! // For each frame
89//! let output = controller.get_rc(frame_type, complexity);
90//! encoder.set_qp(output.qp);
91//! encoder.set_lambda(output.lambda);
92//! ```
93//!
94//! ## Two-Pass VBR
95//!
96//! ```ignore
97//! use oximedia_codec::rate_control::{VbrController, RcConfig};
98//!
99//! // First pass
100//! let config = RcConfig::vbr(5_000_000, 10_000_000);
101//! let mut pass1 = VbrController::new(&config);
102//! pass1.set_pass(1);
103//!
104//! // Analyze all frames...
105//! let first_pass_data = pass1.finalize_first_pass()?;
106//!
107//! // Second pass
108//! let mut pass2 = VbrController::new(&config);
109//! pass2.set_pass(2);
110//! pass2.set_first_pass_data(first_pass_data);
111//!
112//! // Encode with optimal bit allocation
113//! ```
114
115#![forbid(unsafe_code)]
116
117pub mod allocation;
118pub mod analysis;
119pub mod aq;
120pub mod buffer;
121pub mod cbr;
122pub mod complexity;
123pub mod cqp;
124pub mod crf;
125pub mod lookahead;
126pub mod quantizer;
127pub mod simple_rc;
128pub mod types;
129pub mod vbr;
130
131// Re-export main types
132pub use allocation::{AllocationResult, AllocationStrategy, BitrateAllocator, GopAllocationStatus};
133pub use analysis::{
134    AnalysisResult, ContentAnalyzer, ContentType, SceneChangeThreshold, TextureMetrics,
135};
136pub use aq::{AdaptiveQuantization, AqMode, AqResult, AqStrength};
137pub use buffer::{BufferModel, RateBuffer, VbvParams};
138pub use cbr::CbrController;
139pub use complexity::{ComplexityEstimator, ComplexityResult, MotionAnalyzer, MotionResult};
140pub use cqp::CqpController;
141pub use crf::{CrfController, QualityPreset};
142pub use lookahead::{
143    AdaptiveAllocation, ContentAdaptiveAllocator, ContentMetrics, Lookahead, LookaheadFrame,
144    MiniGopInfo, SceneChangeDetector, SceneContentType,
145};
146pub use quantizer::{BlockQpMap, QpResult, QpSelector, QpStrategy};
147pub use simple_rc::{
148    SimpleRateControlConfig, SimpleRateControlMode, SimpleRateControlStats, SimpleRateController,
149};
150pub use types::{
151    FrameStats, GopStats, RateControlMode, RcConfig, RcConfigError, RcOutput, RcState,
152};
153pub use vbr::{FirstPassData, VbrController};
154
155/// Create a rate controller based on the configuration mode.
156///
157/// This factory function creates the appropriate controller type based on
158/// the `RateControlMode` specified in the configuration.
159///
160/// # Example
161///
162/// ```ignore
163/// use oximedia_codec::rate_control::{create_controller, RcConfig, RateControlMode};
164///
165/// let config = RcConfig::cbr(5_000_000);
166/// let controller = create_controller(&config);
167/// ```
168#[must_use]
169pub fn create_controller(config: &RcConfig) -> Box<dyn RateController> {
170    match config.mode {
171        RateControlMode::Cqp => Box::new(CqpController::new(config)),
172        RateControlMode::Cbr => Box::new(CbrController::new(config)),
173        RateControlMode::Vbr | RateControlMode::Abr => Box::new(VbrController::new(config)),
174        RateControlMode::Crf => Box::new(CrfController::new(config)),
175    }
176}
177
178/// Common trait for all rate controllers.
179pub trait RateController: Send {
180    /// Get rate control output for a frame.
181    fn get_output(&mut self, frame_type: crate::frame::FrameType, complexity: f32) -> RcOutput;
182
183    /// Update controller with frame encoding results.
184    fn update_stats(&mut self, stats: &FrameStats);
185
186    /// Reset the controller state.
187    fn reset(&mut self);
188
189    /// Get the current QP value.
190    fn current_qp(&self) -> f32;
191
192    /// Get total frames processed.
193    fn frame_count(&self) -> u64;
194
195    /// Get total bits produced.
196    fn total_bits(&self) -> u64;
197}
198
199impl RateController for CqpController {
200    fn get_output(&mut self, frame_type: crate::frame::FrameType, _complexity: f32) -> RcOutput {
201        self.get_qp(frame_type)
202    }
203
204    fn update_stats(&mut self, stats: &FrameStats) {
205        self.update(stats);
206    }
207
208    fn reset(&mut self) {
209        CqpController::reset(self);
210    }
211
212    fn current_qp(&self) -> f32 {
213        self.base_qp() as f32
214    }
215
216    fn frame_count(&self) -> u64 {
217        CqpController::frame_count(self)
218    }
219
220    fn total_bits(&self) -> u64 {
221        CqpController::total_bits(self)
222    }
223}
224
225impl RateController for CbrController {
226    fn get_output(&mut self, frame_type: crate::frame::FrameType, _complexity: f32) -> RcOutput {
227        self.get_rc(frame_type)
228    }
229
230    fn update_stats(&mut self, stats: &FrameStats) {
231        self.update(stats);
232    }
233
234    fn reset(&mut self) {
235        CbrController::reset(self);
236    }
237
238    fn current_qp(&self) -> f32 {
239        CbrController::current_qp(self)
240    }
241
242    fn frame_count(&self) -> u64 {
243        CbrController::frame_count(self)
244    }
245
246    fn total_bits(&self) -> u64 {
247        // CBR doesn't track total bits directly, compute from frame count
248        0
249    }
250}
251
252impl RateController for VbrController {
253    fn get_output(&mut self, frame_type: crate::frame::FrameType, complexity: f32) -> RcOutput {
254        self.get_rc(frame_type, complexity)
255    }
256
257    fn update_stats(&mut self, stats: &FrameStats) {
258        self.update(stats);
259    }
260
261    fn reset(&mut self) {
262        VbrController::reset(self);
263    }
264
265    fn current_qp(&self) -> f32 {
266        VbrController::current_qp(self)
267    }
268
269    fn frame_count(&self) -> u64 {
270        VbrController::frame_count(self)
271    }
272
273    fn total_bits(&self) -> u64 {
274        0
275    }
276}
277
278impl RateController for CrfController {
279    fn get_output(&mut self, frame_type: crate::frame::FrameType, complexity: f32) -> RcOutput {
280        self.get_rc(frame_type, complexity)
281    }
282
283    fn update_stats(&mut self, stats: &FrameStats) {
284        self.update(stats);
285    }
286
287    fn reset(&mut self) {
288        CrfController::reset(self);
289    }
290
291    fn current_qp(&self) -> f32 {
292        CrfController::current_qp(self)
293    }
294
295    fn frame_count(&self) -> u64 {
296        CrfController::frame_count(self)
297    }
298
299    fn total_bits(&self) -> u64 {
300        CrfController::total_bits(self)
301    }
302}
303
304#[cfg(test)]
305mod tests {
306    use super::*;
307    use crate::frame::FrameType;
308
309    #[test]
310    fn test_create_controller_cqp() {
311        let config = RcConfig::cqp(28);
312        let mut controller = create_controller(&config);
313
314        let output = controller.get_output(FrameType::Key, 1.0);
315        assert_eq!(output.qp, 28);
316    }
317
318    #[test]
319    fn test_create_controller_cbr() {
320        let config = RcConfig::cbr(5_000_000);
321        let mut controller = create_controller(&config);
322
323        let output = controller.get_output(FrameType::Key, 1.0);
324        assert!(!output.drop_frame);
325        assert!(output.target_bits > 0);
326    }
327
328    #[test]
329    fn test_create_controller_vbr() {
330        let config = RcConfig::vbr(5_000_000, 10_000_000);
331        let mut controller = create_controller(&config);
332
333        let output = controller.get_output(FrameType::Key, 1.0);
334        assert!(output.target_bits > 0);
335    }
336
337    #[test]
338    fn test_create_controller_crf() {
339        let config = RcConfig::crf(23.0);
340        let mut controller = create_controller(&config);
341
342        let output = controller.get_output(FrameType::Key, 1.0);
343        assert!(output.qp > 0);
344    }
345
346    #[test]
347    fn test_controller_trait_update() {
348        let config = RcConfig::cqp(28);
349        let mut controller = create_controller(&config);
350
351        let mut stats = FrameStats::new(0, FrameType::Key);
352        stats.bits = 100_000;
353        stats.qp_f = 28.0;
354
355        controller.update_stats(&stats);
356        assert_eq!(controller.frame_count(), 1);
357    }
358
359    #[test]
360    fn test_controller_trait_reset() {
361        let config = RcConfig::cqp(28);
362        let mut controller = create_controller(&config);
363
364        let mut stats = FrameStats::new(0, FrameType::Key);
365        stats.bits = 100_000;
366        controller.update_stats(&stats);
367
368        controller.reset();
369        assert_eq!(controller.frame_count(), 0);
370    }
371
372    #[test]
373    fn test_all_modes_covered() {
374        for mode in [
375            RateControlMode::Cqp,
376            RateControlMode::Cbr,
377            RateControlMode::Vbr,
378            RateControlMode::Abr,
379            RateControlMode::Crf,
380        ] {
381            let config = RcConfig {
382                mode,
383                target_bitrate: 5_000_000,
384                ..Default::default()
385            };
386
387            let mut controller = create_controller(&config);
388            let output = controller.get_output(FrameType::Inter, 1.0);
389
390            // All controllers should produce valid output
391            assert!(output.qp > 0 || output.drop_frame);
392        }
393    }
394}