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 abr;
118pub mod allocation;
119pub mod analysis;
120pub mod aq;
121pub mod buffer;
122pub mod cbr;
123pub mod complexity;
124pub mod cqp;
125pub mod crf;
126pub mod lookahead;
127pub mod quantizer;
128pub mod scene_adaptive;
129pub mod simple_rc;
130pub mod types;
131pub mod vbr;
132
133// Re-export main types
134pub use abr::{
135    AbrController, FirstPassStats as AbrFirstPassStats, FramePassStats as AbrFramePassStats,
136};
137pub use allocation::{AllocationResult, AllocationStrategy, BitrateAllocator, GopAllocationStatus};
138pub use analysis::{
139    AnalysisResult, ContentAnalyzer, ContentType, SceneChangeThreshold, TextureMetrics,
140};
141pub use aq::{AdaptiveQuantization, AqMode, AqResult, AqStrength};
142pub use buffer::{BufferModel, RateBuffer, VbvParams};
143pub use cbr::CbrController;
144pub use complexity::{ComplexityEstimator, ComplexityResult, MotionAnalyzer, MotionResult};
145pub use cqp::CqpController;
146pub use crf::{CrfController, QualityPreset};
147pub use lookahead::{
148    AdaptiveAllocation, ContentAdaptiveAllocator, ContentMetrics, Lookahead, LookaheadFrame,
149    MiniGopInfo, SceneChangeDetector, SceneContentType,
150};
151pub use quantizer::{BlockQpMap, QpResult, QpSelector, QpStrategy};
152pub use scene_adaptive::{
153    FrameBitTarget, FrameContentMetrics, Scene as AdaptiveScene, SceneAdaptiveAllocator,
154    SceneAdaptiveConfig, SceneContentType as SceneAdaptiveContentType,
155};
156pub use simple_rc::{
157    SimpleRateControlConfig, SimpleRateControlMode, SimpleRateControlStats, SimpleRateController,
158};
159pub use types::{
160    FrameStats, GopStats, RateControlMode, RcConfig, RcConfigError, RcOutput, RcState,
161};
162pub use vbr::{FirstPassData, VbrController};
163
164/// Create a rate controller based on the configuration mode.
165///
166/// This factory function creates the appropriate controller type based on
167/// the `RateControlMode` specified in the configuration.
168///
169/// # Example
170///
171/// ```ignore
172/// use oximedia_codec::rate_control::{create_controller, RcConfig, RateControlMode};
173///
174/// let config = RcConfig::cbr(5_000_000);
175/// let controller = create_controller(&config);
176/// ```
177#[must_use]
178pub fn create_controller(config: &RcConfig) -> Box<dyn RateController> {
179    match config.mode {
180        RateControlMode::Cqp => Box::new(CqpController::new(config)),
181        RateControlMode::Cbr => Box::new(CbrController::new(config)),
182        RateControlMode::Vbr | RateControlMode::Abr => Box::new(VbrController::new(config)),
183        RateControlMode::Crf => Box::new(CrfController::new(config)),
184    }
185}
186
187/// Common trait for all rate controllers.
188pub trait RateController: Send {
189    /// Get rate control output for a frame.
190    fn get_output(&mut self, frame_type: crate::frame::FrameType, complexity: f32) -> RcOutput;
191
192    /// Update controller with frame encoding results.
193    fn update_stats(&mut self, stats: &FrameStats);
194
195    /// Reset the controller state.
196    fn reset(&mut self);
197
198    /// Get the current QP value.
199    fn current_qp(&self) -> f32;
200
201    /// Get total frames processed.
202    fn frame_count(&self) -> u64;
203
204    /// Get total bits produced.
205    fn total_bits(&self) -> u64;
206}
207
208impl RateController for CqpController {
209    fn get_output(&mut self, frame_type: crate::frame::FrameType, _complexity: f32) -> RcOutput {
210        self.get_qp(frame_type)
211    }
212
213    fn update_stats(&mut self, stats: &FrameStats) {
214        self.update(stats);
215    }
216
217    fn reset(&mut self) {
218        CqpController::reset(self);
219    }
220
221    fn current_qp(&self) -> f32 {
222        self.base_qp() as f32
223    }
224
225    fn frame_count(&self) -> u64 {
226        CqpController::frame_count(self)
227    }
228
229    fn total_bits(&self) -> u64 {
230        CqpController::total_bits(self)
231    }
232}
233
234impl RateController for CbrController {
235    fn get_output(&mut self, frame_type: crate::frame::FrameType, _complexity: f32) -> RcOutput {
236        self.get_rc(frame_type)
237    }
238
239    fn update_stats(&mut self, stats: &FrameStats) {
240        self.update(stats);
241    }
242
243    fn reset(&mut self) {
244        CbrController::reset(self);
245    }
246
247    fn current_qp(&self) -> f32 {
248        CbrController::current_qp(self)
249    }
250
251    fn frame_count(&self) -> u64 {
252        CbrController::frame_count(self)
253    }
254
255    fn total_bits(&self) -> u64 {
256        // CBR doesn't track total bits directly, compute from frame count
257        0
258    }
259}
260
261impl RateController for VbrController {
262    fn get_output(&mut self, frame_type: crate::frame::FrameType, complexity: f32) -> RcOutput {
263        self.get_rc(frame_type, complexity)
264    }
265
266    fn update_stats(&mut self, stats: &FrameStats) {
267        self.update(stats);
268    }
269
270    fn reset(&mut self) {
271        VbrController::reset(self);
272    }
273
274    fn current_qp(&self) -> f32 {
275        VbrController::current_qp(self)
276    }
277
278    fn frame_count(&self) -> u64 {
279        VbrController::frame_count(self)
280    }
281
282    fn total_bits(&self) -> u64 {
283        0
284    }
285}
286
287impl RateController for CrfController {
288    fn get_output(&mut self, frame_type: crate::frame::FrameType, complexity: f32) -> RcOutput {
289        self.get_rc(frame_type, complexity)
290    }
291
292    fn update_stats(&mut self, stats: &FrameStats) {
293        self.update(stats);
294    }
295
296    fn reset(&mut self) {
297        CrfController::reset(self);
298    }
299
300    fn current_qp(&self) -> f32 {
301        CrfController::current_qp(self)
302    }
303
304    fn frame_count(&self) -> u64 {
305        CrfController::frame_count(self)
306    }
307
308    fn total_bits(&self) -> u64 {
309        CrfController::total_bits(self)
310    }
311}
312
313#[cfg(test)]
314mod tests {
315    use super::*;
316    use crate::frame::FrameType;
317
318    #[test]
319    fn test_create_controller_cqp() {
320        let config = RcConfig::cqp(28);
321        let mut controller = create_controller(&config);
322
323        let output = controller.get_output(FrameType::Key, 1.0);
324        assert_eq!(output.qp, 28);
325    }
326
327    #[test]
328    fn test_create_controller_cbr() {
329        let config = RcConfig::cbr(5_000_000);
330        let mut controller = create_controller(&config);
331
332        let output = controller.get_output(FrameType::Key, 1.0);
333        assert!(!output.drop_frame);
334        assert!(output.target_bits > 0);
335    }
336
337    #[test]
338    fn test_create_controller_vbr() {
339        let config = RcConfig::vbr(5_000_000, 10_000_000);
340        let mut controller = create_controller(&config);
341
342        let output = controller.get_output(FrameType::Key, 1.0);
343        assert!(output.target_bits > 0);
344    }
345
346    #[test]
347    fn test_create_controller_crf() {
348        let config = RcConfig::crf(23.0);
349        let mut controller = create_controller(&config);
350
351        let output = controller.get_output(FrameType::Key, 1.0);
352        assert!(output.qp > 0);
353    }
354
355    #[test]
356    fn test_controller_trait_update() {
357        let config = RcConfig::cqp(28);
358        let mut controller = create_controller(&config);
359
360        let mut stats = FrameStats::new(0, FrameType::Key);
361        stats.bits = 100_000;
362        stats.qp_f = 28.0;
363
364        controller.update_stats(&stats);
365        assert_eq!(controller.frame_count(), 1);
366    }
367
368    #[test]
369    fn test_controller_trait_reset() {
370        let config = RcConfig::cqp(28);
371        let mut controller = create_controller(&config);
372
373        let mut stats = FrameStats::new(0, FrameType::Key);
374        stats.bits = 100_000;
375        controller.update_stats(&stats);
376
377        controller.reset();
378        assert_eq!(controller.frame_count(), 0);
379    }
380
381    #[test]
382    fn test_all_modes_covered() {
383        for mode in [
384            RateControlMode::Cqp,
385            RateControlMode::Cbr,
386            RateControlMode::Vbr,
387            RateControlMode::Abr,
388            RateControlMode::Crf,
389        ] {
390            let config = RcConfig {
391                mode,
392                target_bitrate: 5_000_000,
393                ..Default::default()
394            };
395
396            let mut controller = create_controller(&config);
397            let output = controller.get_output(FrameType::Inter, 1.0);
398
399            // All controllers should produce valid output
400            assert!(output.qp > 0 || output.drop_frame);
401        }
402    }
403}