#![allow(clippy::cast_lossless)]
#![allow(clippy::cast_precision_loss)]
#![allow(clippy::cast_possible_truncation)]
#![allow(clippy::cast_sign_loss)]
#![allow(clippy::cast_possible_wrap)]
#![allow(clippy::too_many_arguments)]
#![allow(clippy::struct_excessive_bools)]
#![forbid(unsafe_code)]
use crate::frame::FrameType;
use super::buffer::BufferModel;
use super::types::{FrameStats, GopStats, RcConfig, RcOutput};
#[derive(Clone, Debug)]
pub struct VbrController {
target_bitrate: u64,
max_bitrate: u64,
min_bitrate: u64,
current_qp: f32,
min_qp: u8,
max_qp: u8,
i_qp_offset: i8,
b_qp_offset: i8,
framerate: f64,
gop_length: u32,
frame_count: u64,
total_bits: u64,
current_gop: GopStats,
gop_history: Vec<GopStats>,
max_gop_history: usize,
pass: u8,
first_pass_data: Option<FirstPassData>,
quality_stability: f32,
bit_reservoir: i64,
max_reservoir: i64,
vbv_buffer: Option<BufferModel>,
enable_vbv: bool,
lookahead_size: usize,
lookahead_frames: Vec<LookaheadFrameData>,
rate_model: RatePredictionModel,
adaptive_gop: bool,
scene_change_threshold: f32,
enable_aq: bool,
#[allow(dead_code)]
aq_strength: f32,
pid_state: PidControllerState,
rdo_params: RdoParameters,
frame_type_state: FrameTypeDecisionState,
}
#[derive(Clone, Debug, Default)]
pub struct FirstPassData {
pub frame_complexity: Vec<f32>,
pub spatial_complexity: Vec<f32>,
pub temporal_complexity: Vec<f32>,
pub gop_complexity: Vec<f32>,
pub total_complexity: f32,
pub frame_count: u64,
pub suggested_bits: Vec<u64>,
pub scene_changes: Vec<u64>,
pub gop_boundaries: Vec<u64>,
pub recommended_qp: Vec<f32>,
}
impl FirstPassData {
pub fn add_frame(&mut self, spatial: f32, temporal: f32, combined: f32) {
self.frame_complexity.push(combined);
self.spatial_complexity.push(spatial);
self.temporal_complexity.push(temporal);
self.total_complexity += combined;
self.frame_count += 1;
}
pub fn mark_scene_change(&mut self, frame_num: u64) {
self.scene_changes.push(frame_num);
}
pub fn add_gop_boundary(&mut self, frame_num: u64) {
self.gop_boundaries.push(frame_num);
}
pub fn finalize_gop(&mut self) {
let gop_start = self
.gop_complexity
.last()
.map(|_| self.gop_boundaries.last().copied().unwrap_or(0))
.unwrap_or(0) as usize;
let gop_sum: f32 = self
.frame_complexity
.get(gop_start..)
.map(|slice| slice.iter().sum())
.unwrap_or(0.0);
self.gop_complexity.push(gop_sum);
}
pub fn calculate_bit_allocation(&mut self, total_bits: u64) {
if self.total_complexity <= 0.0 || self.frame_complexity.is_empty() {
return;
}
let bits_per_complexity = total_bits as f64 / self.total_complexity as f64;
self.suggested_bits = self
.frame_complexity
.iter()
.map(|c| ((*c as f64) * bits_per_complexity) as u64)
.collect();
let avg_complexity = self.total_complexity / self.frame_count as f32;
for complexity in &self.frame_complexity {
let complexity_ratio = complexity / avg_complexity;
let qp_adjustment = (complexity_ratio - 1.0) * 4.0;
let base_qp = 28.0;
let recommended = (base_qp + qp_adjustment).clamp(18.0, 51.0);
self.recommended_qp.push(recommended);
}
}
#[must_use]
pub fn get_suggested_bits(&self, frame_num: u64) -> Option<u64> {
self.suggested_bits.get(frame_num as usize).copied()
}
#[must_use]
pub fn get_recommended_qp(&self, frame_num: u64) -> Option<f32> {
self.recommended_qp.get(frame_num as usize).copied()
}
#[must_use]
pub fn is_scene_change(&self, frame_num: u64) -> bool {
self.scene_changes.contains(&frame_num)
}
}
#[derive(Clone, Debug, Default)]
struct LookaheadFrameData {
#[allow(dead_code)]
frame_index: u64,
#[allow(dead_code)]
spatial_complexity: f32,
#[allow(dead_code)]
temporal_complexity: f32,
combined_complexity: f32,
is_scene_change: bool,
#[allow(dead_code)]
predicted_type: FrameType,
#[allow(dead_code)]
predicted_bits: u64,
}
#[derive(Clone, Debug)]
struct RatePredictionModel {
linear_a: f64,
linear_b: f64,
quad_a: f64,
quad_b: f64,
quad_c: f64,
power_a: f64,
power_b: f64,
active_model: u8,
history_complexity: Vec<f32>,
history_bits: Vec<u64>,
max_history: usize,
update_count: u32,
}
impl Default for RatePredictionModel {
fn default() -> Self {
Self {
linear_a: 100_000.0,
linear_b: 50_000.0,
quad_a: 10_000.0,
quad_b: 50_000.0,
quad_c: 20_000.0,
power_a: 100_000.0,
power_b: 1.2,
active_model: 0,
history_complexity: Vec::new(),
history_bits: Vec::new(),
max_history: 100,
update_count: 0,
}
}
}
impl RatePredictionModel {
fn predict(&self, complexity: f32) -> u64 {
if complexity <= 0.0 {
return 50_000;
}
let prediction = match self.active_model {
0 => self.predict_linear(complexity),
1 => self.predict_quadratic(complexity),
2 => self.predict_power(complexity),
_ => self.predict_linear(complexity),
};
prediction.max(1000.0) as u64
}
fn predict_linear(&self, complexity: f32) -> f64 {
self.linear_a * complexity as f64 + self.linear_b
}
fn predict_quadratic(&self, complexity: f32) -> f64 {
let c = complexity as f64;
self.quad_a * c * c + self.quad_b * c + self.quad_c
}
fn predict_power(&self, complexity: f32) -> f64 {
self.power_a * (complexity as f64).powf(self.power_b)
}
fn update(&mut self, complexity: f32, bits: u64) {
self.history_complexity.push(complexity);
self.history_bits.push(bits);
if self.history_complexity.len() > self.max_history {
self.history_complexity.remove(0);
self.history_bits.remove(0);
}
self.update_count += 1;
if self.update_count >= 10 {
self.fit_models();
self.update_count = 0;
}
}
fn fit_models(&mut self) {
if self.history_complexity.len() < 5 {
return;
}
self.fit_linear_model();
self.fit_quadratic_model();
self.fit_power_model();
self.select_best_model();
}
fn fit_linear_model(&mut self) {
let n = self.history_complexity.len();
let mut sum_x = 0.0;
let mut sum_y = 0.0;
let mut sum_xy = 0.0;
let mut sum_xx = 0.0;
for i in 0..n {
let x = self.history_complexity[i] as f64;
let y = self.history_bits[i] as f64;
sum_x += x;
sum_y += y;
sum_xy += x * y;
sum_xx += x * x;
}
let n_f = n as f64;
let denominator = n_f * sum_xx - sum_x * sum_x;
if denominator.abs() > 1e-6 {
self.linear_a = (n_f * sum_xy - sum_x * sum_y) / denominator;
self.linear_b = (sum_y - self.linear_a * sum_x) / n_f;
}
}
fn fit_quadratic_model(&mut self) {
let n = self.history_complexity.len();
if n < 3 {
return;
}
let mut sum_x = 0.0;
let mut sum_y = 0.0;
let mut sum_x2 = 0.0;
let mut sum_xy = 0.0;
for i in 0..n {
let x = self.history_complexity[i] as f64;
let y = self.history_bits[i] as f64;
sum_x += x;
sum_y += y;
sum_x2 += x * x;
sum_xy += x * y;
}
let n_f = n as f64;
let avg_x = sum_x / n_f;
let avg_y = sum_y / n_f;
self.quad_b = self.linear_a;
self.quad_a = (sum_xy - n_f * avg_x * avg_y) / (sum_x2 - n_f * avg_x * avg_x) * 0.1;
self.quad_c = avg_y - self.quad_b * avg_x - self.quad_a * avg_x * avg_x;
}
fn fit_power_model(&mut self) {
let n = self.history_complexity.len();
let mut sum_log_x = 0.0;
let mut sum_log_y = 0.0;
let mut sum_log_xy = 0.0;
let mut sum_log_xx = 0.0;
let mut count = 0;
for i in 0..n {
let x = self.history_complexity[i] as f64;
let y = self.history_bits[i] as f64;
if x > 0.0 && y > 0.0 {
let log_x = x.ln();
let log_y = y.ln();
sum_log_x += log_x;
sum_log_y += log_y;
sum_log_xy += log_x * log_y;
sum_log_xx += log_x * log_x;
count += 1;
}
}
if count < 3 {
return;
}
let n_f = count as f64;
let denominator = n_f * sum_log_xx - sum_log_x * sum_log_x;
if denominator.abs() > 1e-6 {
self.power_b = (n_f * sum_log_xy - sum_log_x * sum_log_y) / denominator;
let log_a = (sum_log_y - self.power_b * sum_log_x) / n_f;
self.power_a = log_a.exp();
}
}
fn select_best_model(&mut self) {
if self.history_complexity.len() < 5 {
return;
}
let mut error_linear = 0.0;
let mut error_quadratic = 0.0;
let mut error_power = 0.0;
for i in 0..self.history_complexity.len() {
let complexity = self.history_complexity[i];
let actual_bits = self.history_bits[i] as f64;
let pred_linear = self.predict_linear(complexity);
let pred_quadratic = self.predict_quadratic(complexity);
let pred_power = self.predict_power(complexity);
error_linear += (actual_bits - pred_linear).abs();
error_quadratic += (actual_bits - pred_quadratic).abs();
error_power += (actual_bits - pred_power).abs();
}
if error_linear <= error_quadratic && error_linear <= error_power {
self.active_model = 0;
} else if error_quadratic <= error_power {
self.active_model = 1;
} else {
self.active_model = 2;
}
}
}
#[derive(Clone, Debug, Default)]
struct PidControllerState {
kp: f32,
ki: f32,
kd: f32,
prev_error: f32,
integral: f32,
max_integral: f32,
}
impl PidControllerState {
fn new() -> Self {
Self {
kp: 0.5,
ki: 0.1,
kd: 0.05,
prev_error: 0.0,
integral: 0.0,
max_integral: 10.0,
}
}
fn calculate(&mut self, error: f32) -> f32 {
let p_term = self.kp * error;
self.integral += error;
self.integral = self.integral.clamp(-self.max_integral, self.max_integral);
let i_term = self.ki * self.integral;
let d_term = self.kd * (error - self.prev_error);
self.prev_error = error;
p_term + i_term + d_term
}
fn reset(&mut self) {
self.prev_error = 0.0;
self.integral = 0.0;
}
}
#[derive(Clone, Debug)]
struct RdoParameters {
base_lambda: f64,
i_lambda_mult: f64,
b_lambda_mult: f64,
#[allow(dead_code)]
psy_rd: bool,
#[allow(dead_code)]
psy_strength: f64,
}
impl Default for RdoParameters {
fn default() -> Self {
Self {
base_lambda: 1.0,
i_lambda_mult: 0.6,
b_lambda_mult: 1.4,
psy_rd: true,
psy_strength: 1.0,
}
}
}
impl RdoParameters {
fn calculate_lambda(&self, qp: f32, frame_type: FrameType) -> f64 {
let base = 0.85 * 2.0_f64.powf((f64::from(qp) - 12.0) / 3.0);
let multiplier = match frame_type {
FrameType::Key => self.i_lambda_mult,
FrameType::BiDir => self.b_lambda_mult,
_ => 1.0,
};
base * multiplier * self.base_lambda
}
fn calculate_lambda_me(&self, lambda: f64) -> f64 {
lambda.sqrt()
}
}
#[derive(Clone, Debug, Default)]
struct FrameTypeDecisionState {
frames_since_keyframe: u32,
consecutive_b_frames: u32,
max_b_frames: u32,
force_keyframe: bool,
}
impl FrameTypeDecisionState {
fn new(max_b_frames: u32) -> Self {
Self {
frames_since_keyframe: 0,
consecutive_b_frames: 0,
max_b_frames,
force_keyframe: false,
}
}
fn decide_type(
&mut self,
gop_length: u32,
is_scene_change: bool,
adaptive_gop: bool,
) -> FrameType {
if self.force_keyframe
|| self.frames_since_keyframe == 0
|| (!adaptive_gop && self.frames_since_keyframe >= gop_length)
|| is_scene_change
{
self.frames_since_keyframe = 1;
self.consecutive_b_frames = 0;
self.force_keyframe = false;
return FrameType::Key;
}
self.frames_since_keyframe += 1;
if self.max_b_frames > 0 && self.consecutive_b_frames < self.max_b_frames {
let mini_gop_pos = self.frames_since_keyframe % (self.max_b_frames + 1);
if mini_gop_pos > 0 && mini_gop_pos <= self.max_b_frames {
self.consecutive_b_frames += 1;
return FrameType::BiDir;
}
}
self.consecutive_b_frames = 0;
FrameType::Inter
}
fn force_next_keyframe(&mut self) {
self.force_keyframe = true;
}
}
impl VbrController {
#[must_use]
pub fn new(config: &RcConfig) -> Self {
let max_bitrate = config.max_bitrate.unwrap_or(config.target_bitrate * 2);
let min_bitrate = config.min_bitrate.unwrap_or(config.target_bitrate / 4);
let framerate = config.framerate();
let vbv_buffer = if config.buffer_size > 0 {
Some(BufferModel::new(
config.buffer_size,
config.target_bitrate,
framerate,
config.initial_buffer_fullness as f64,
))
} else {
None
};
Self {
target_bitrate: config.target_bitrate,
max_bitrate,
min_bitrate,
current_qp: config.initial_qp as f32,
min_qp: config.min_qp,
max_qp: config.max_qp,
i_qp_offset: config.i_qp_offset,
b_qp_offset: config.b_qp_offset,
framerate,
gop_length: config.gop_length,
frame_count: 0,
total_bits: 0,
current_gop: GopStats::new(0, 0),
gop_history: Vec::new(),
max_gop_history: 10,
pass: 0,
first_pass_data: None,
quality_stability: 0.5,
bit_reservoir: 0,
max_reservoir: (config.target_bitrate as i64) * 2,
vbv_buffer,
enable_vbv: config.buffer_size > 0,
lookahead_size: config.lookahead_depth.min(250),
lookahead_frames: Vec::new(),
rate_model: RatePredictionModel::default(),
adaptive_gop: true,
scene_change_threshold: config.scene_cut_threshold,
enable_aq: config.enable_aq,
aq_strength: config.aq_strength,
pid_state: PidControllerState::new(),
rdo_params: RdoParameters::default(),
frame_type_state: FrameTypeDecisionState::new(3),
}
}
pub fn set_pass(&mut self, pass: u8) {
self.pass = pass.min(2);
if pass == 1 {
self.first_pass_data = Some(FirstPassData::default());
}
}
pub fn set_quality_stability(&mut self, stability: f32) {
self.quality_stability = stability.clamp(0.0, 1.0);
}
pub fn set_vbv_enabled(&mut self, enabled: bool) {
self.enable_vbv = enabled;
}
pub fn set_lookahead_size(&mut self, size: usize) {
self.lookahead_size = size.clamp(10, 250);
self.lookahead_frames.reserve(self.lookahead_size);
}
pub fn set_adaptive_gop(&mut self, enabled: bool) {
self.adaptive_gop = enabled;
}
pub fn set_scene_change_threshold(&mut self, threshold: f32) {
self.scene_change_threshold = threshold.clamp(0.0, 1.0);
}
pub fn set_first_pass_data(&mut self, data: FirstPassData) {
self.first_pass_data = Some(data);
}
pub fn push_lookahead_frame(
&mut self,
spatial: f32,
temporal: f32,
combined: f32,
is_scene_change: bool,
) {
let frame_data = LookaheadFrameData {
frame_index: self.frame_count + self.lookahead_frames.len() as u64,
spatial_complexity: spatial,
temporal_complexity: temporal,
combined_complexity: combined,
is_scene_change,
predicted_type: FrameType::Inter,
predicted_bits: self.rate_model.predict(combined),
};
self.lookahead_frames.push(frame_data);
if self.lookahead_frames.len() > self.lookahead_size {
self.lookahead_frames.remove(0);
}
}
#[must_use]
pub fn get_rc(&mut self, frame_type: FrameType, complexity: f32) -> RcOutput {
let is_scene_change = self.detect_scene_change_from_lookahead();
let actual_frame_type =
self.frame_type_state
.decide_type(self.gop_length, is_scene_change, self.adaptive_gop);
let final_frame_type = if frame_type == FrameType::Key {
FrameType::Key
} else {
actual_frame_type
};
let target_bits = self.calculate_target_bits(final_frame_type, complexity);
let qp_adjustment = self.calculate_closed_loop_qp_adjustment();
let adjusted_qp = self.current_qp + qp_adjustment;
let offset = match final_frame_type {
FrameType::Key => self.i_qp_offset,
FrameType::BiDir => self.b_qp_offset,
FrameType::Inter | FrameType::Switch => 0,
};
let complexity_adjustment = self.calculate_complexity_adjustment(complexity);
let vbv_adjustment = if self.enable_vbv {
self.calculate_vbv_qp_adjustment()
} else {
0.0
};
let final_qp = (adjusted_qp + offset as f32 + complexity_adjustment + vbv_adjustment)
.clamp(self.min_qp as f32, self.max_qp as f32);
let qp = final_qp.round() as u8;
let (min_bits, max_bits) = self.calculate_bit_limits(final_frame_type, target_bits);
let lambda = self.rdo_params.calculate_lambda(final_qp, final_frame_type);
let lambda_me = self.rdo_params.calculate_lambda_me(lambda);
let mut output = RcOutput {
qp,
qp_f: final_qp,
target_bits,
min_bits,
max_bits,
lambda,
lambda_me,
force_keyframe: final_frame_type == FrameType::Key && is_scene_change,
..Default::default()
};
if self.enable_aq {
output.qp_offsets = Some(self.calculate_aq_offsets(final_qp));
}
output
}
fn detect_scene_change_from_lookahead(&self) -> bool {
if self.lookahead_frames.is_empty() {
return false;
}
self.lookahead_frames
.first()
.map(|f| f.is_scene_change)
.unwrap_or(false)
}
fn calculate_target_bits(&self, frame_type: FrameType, complexity: f32) -> u64 {
let base_target = self.bits_per_frame_at_bitrate(self.target_bitrate);
if self.pass == 2 {
if let Some(ref data) = self.first_pass_data {
if let Some(suggested) = data.get_suggested_bits(self.frame_count) {
return suggested;
}
}
}
let model_prediction = self.rate_model.predict(complexity);
let type_multiplier = match frame_type {
FrameType::Key => 3.0,
FrameType::Inter => 1.0,
FrameType::BiDir => 0.5,
FrameType::Switch => 1.5,
};
let lookahead_mult = self.calculate_lookahead_multiplier();
let complexity_multiplier: f64 = if complexity > 0.0 {
(f64::from(complexity) / f64::from(self.average_complexity())).clamp(0.5, 2.0)
} else {
1.0
};
let reservoir_adjustment = self.calculate_reservoir_adjustment();
let target =
(base_target as f64 * type_multiplier * complexity_multiplier * lookahead_mult)
.max(model_prediction as f64)
+ reservoir_adjustment;
let adjusted_target = target.max(base_target as f64 / 4.0);
adjusted_target as u64
}
fn calculate_lookahead_multiplier(&self) -> f64 {
if self.lookahead_frames.is_empty() {
return 1.0;
}
let avg_future_complexity: f32 = self
.lookahead_frames
.iter()
.map(|f| f.combined_complexity)
.sum::<f32>()
/ self.lookahead_frames.len() as f32;
let current_complexity = self
.lookahead_frames
.first()
.map(|f| f.combined_complexity)
.unwrap_or(1.0);
if avg_future_complexity > current_complexity * 1.3 {
0.9
} else if avg_future_complexity < current_complexity * 0.7 {
1.1
} else {
1.0
}
}
fn calculate_closed_loop_qp_adjustment(&mut self) -> f32 {
if self.frame_count == 0 {
return 0.0;
}
let elapsed_time = self.frame_count as f64 / self.framerate;
if elapsed_time <= 0.0 {
return 0.0;
}
let actual_bitrate = self.total_bits as f64 / elapsed_time;
let error = (actual_bitrate - self.target_bitrate as f64) / self.target_bitrate as f64;
let pid_output = self.pid_state.calculate(error as f32);
pid_output * (1.0 - self.quality_stability)
}
fn calculate_complexity_adjustment(&self, complexity: f32) -> f32 {
let avg = self.average_complexity();
if avg <= 0.0 {
return 0.0;
}
let ratio = complexity / avg;
let adjustment = (ratio - 1.0) * 2.0 * self.quality_stability;
adjustment.clamp(-2.0, 2.0)
}
fn calculate_vbv_qp_adjustment(&self) -> f32 {
if let Some(ref buffer) = self.vbv_buffer {
let fullness = buffer.fullness();
let target_fullness = 0.5;
let error = fullness - target_fullness;
if error > 0.3 {
(error - 0.3) * 20.0
} else if error < -0.3 {
(error + 0.3) * 20.0
} else {
error * 5.0
}
} else {
0.0
}
}
fn calculate_bit_limits(&self, frame_type: FrameType, target_bits: u64) -> (u64, u64) {
let base_min = self.bits_per_frame_at_bitrate(self.min_bitrate) / 4;
let base_max = self.bits_per_frame_at_bitrate(self.max_bitrate) * 4;
let mut min_bits = base_min;
let mut max_bits = base_max;
if self.enable_vbv {
if let Some(ref buffer) = self.vbv_buffer {
let available = buffer.max_frame_bits();
max_bits = max_bits.min(available);
let min_to_prevent_underflow = target_bits / 2;
min_bits = min_bits.max(min_to_prevent_underflow);
}
}
match frame_type {
FrameType::Key => {
max_bits = max_bits.max(target_bits * 5);
}
FrameType::BiDir => {
max_bits = max_bits.min(target_bits * 2);
}
_ => {}
}
(min_bits, max_bits)
}
fn calculate_aq_offsets(&self, _base_qp: f32) -> Vec<f32> {
Vec::new()
}
fn calculate_reservoir_adjustment(&self) -> f64 {
let target_per_frame = self.bits_per_frame_at_bitrate(self.target_bitrate);
let reservoir_factor = self.bit_reservoir as f64 / self.max_reservoir as f64;
reservoir_factor * (target_per_frame as f64 * 0.1)
}
fn bits_per_frame_at_bitrate(&self, bitrate: u64) -> u64 {
if self.framerate <= 0.0 {
return 0;
}
(bitrate as f64 / self.framerate) as u64
}
fn average_complexity(&self) -> f32 {
if self.gop_history.is_empty() {
return 1.0;
}
let total: f32 = self.gop_history.iter().map(|g| g.average_complexity).sum();
(total / self.gop_history.len() as f32).max(0.01)
}
pub fn update(&mut self, stats: &FrameStats) {
self.frame_count += 1;
self.total_bits += stats.bits;
let target = self.bits_per_frame_at_bitrate(self.target_bitrate);
self.bit_reservoir += target as i64 - stats.bits as i64;
self.bit_reservoir = self
.bit_reservoir
.clamp(-self.max_reservoir, self.max_reservoir);
if self.enable_vbv {
if let Some(ref mut buffer) = self.vbv_buffer {
buffer.fill_for_frame();
buffer.remove_frame_bits(stats.bits);
}
}
self.rate_model.update(stats.complexity, stats.bits);
self.current_gop.add_frame(stats.clone());
if stats.frame_type == FrameType::Key && self.current_gop.frame_count > 1 {
self.finalize_gop();
}
if self.pass == 1 {
if let Some(ref mut data) = self.first_pass_data {
data.add_frame(
stats.spatial_complexity,
stats.temporal_complexity,
stats.complexity,
);
if stats.scene_cut {
data.mark_scene_change(stats.frame_num);
}
if stats.frame_type == FrameType::Key && data.frame_count > 1 {
data.finalize_gop();
data.add_gop_boundary(stats.frame_num);
}
}
}
self.adjust_base_qp(stats);
if !self.lookahead_frames.is_empty() {
self.lookahead_frames.remove(0);
}
}
fn finalize_gop(&mut self) {
self.gop_history.push(self.current_gop.clone());
if self.gop_history.len() > self.max_gop_history {
self.gop_history.remove(0);
}
let next_gop_index = self.current_gop.gop_index + 1;
self.current_gop = GopStats::new(next_gop_index, self.frame_count);
}
fn adjust_base_qp(&mut self, stats: &FrameStats) {
if stats.target_bits == 0 {
return;
}
let accuracy = stats.bits as f32 / stats.target_bits as f32;
let adjustment = if accuracy > 1.3 {
0.15
} else if accuracy > 1.1 {
0.05
} else if accuracy < 0.7 {
-0.15
} else if accuracy < 0.9 {
-0.05
} else {
0.0
};
self.current_qp =
(self.current_qp + adjustment).clamp(self.min_qp as f32, self.max_qp as f32);
}
#[must_use]
pub fn finalize_first_pass(&mut self) -> Option<FirstPassData> {
if self.pass != 1 {
return None;
}
if let Some(ref mut data) = self.first_pass_data {
data.finalize_gop();
let duration = self.frame_count as f64 / self.framerate;
let total_bits = (self.target_bitrate as f64 * duration) as u64;
data.calculate_bit_allocation(total_bits);
}
self.first_pass_data.take()
}
#[must_use]
pub fn target_bitrate(&self) -> u64 {
self.target_bitrate
}
#[must_use]
pub fn max_bitrate(&self) -> u64 {
self.max_bitrate
}
#[must_use]
pub fn min_bitrate(&self) -> u64 {
self.min_bitrate
}
#[must_use]
pub fn current_bitrate(&self) -> f64 {
if self.frame_count == 0 || self.framerate <= 0.0 {
return 0.0;
}
let elapsed = self.frame_count as f64 / self.framerate;
self.total_bits as f64 / elapsed
}
#[must_use]
pub fn current_qp(&self) -> f32 {
self.current_qp
}
#[must_use]
pub fn frame_count(&self) -> u64 {
self.frame_count
}
#[must_use]
pub fn bit_reservoir(&self) -> i64 {
self.bit_reservoir
}
#[must_use]
pub fn vbv_fullness(&self) -> f32 {
self.vbv_buffer
.as_ref()
.map(|b| b.fullness())
.unwrap_or(0.5)
}
pub fn force_keyframe(&mut self) {
self.frame_type_state.force_next_keyframe();
}
pub fn reset(&mut self) {
self.frame_count = 0;
self.total_bits = 0;
self.bit_reservoir = 0;
self.current_gop = GopStats::new(0, 0);
self.gop_history.clear();
self.lookahead_frames.clear();
self.pid_state.reset();
self.frame_type_state = FrameTypeDecisionState::new(self.frame_type_state.max_b_frames);
if let Some(ref mut buffer) = self.vbv_buffer {
buffer.reset();
}
}
}
impl Default for VbrController {
fn default() -> Self {
Self {
target_bitrate: 5_000_000,
max_bitrate: 10_000_000,
min_bitrate: 1_000_000,
current_qp: 28.0,
min_qp: 1,
max_qp: 63,
i_qp_offset: -2,
b_qp_offset: 2,
framerate: 30.0,
gop_length: 250,
frame_count: 0,
total_bits: 0,
current_gop: GopStats::new(0, 0),
gop_history: Vec::new(),
max_gop_history: 10,
pass: 0,
first_pass_data: None,
quality_stability: 0.5,
bit_reservoir: 0,
max_reservoir: 10_000_000,
vbv_buffer: None,
enable_vbv: false,
lookahead_size: 40,
lookahead_frames: Vec::new(),
rate_model: RatePredictionModel::default(),
adaptive_gop: true,
scene_change_threshold: 0.4,
enable_aq: true,
aq_strength: 1.0,
pid_state: PidControllerState::new(),
rdo_params: RdoParameters::default(),
frame_type_state: FrameTypeDecisionState::new(3),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
fn create_test_controller() -> VbrController {
let config = RcConfig::vbr(5_000_000, 10_000_000);
VbrController::new(&config)
}
#[test]
fn test_vbr_creation() {
let controller = create_test_controller();
assert_eq!(controller.target_bitrate(), 5_000_000);
assert_eq!(controller.max_bitrate(), 10_000_000);
}
#[test]
fn test_get_rc() {
let mut controller = create_test_controller();
let output = controller.get_rc(FrameType::Key, 1.0);
assert!(!output.drop_frame);
assert!(output.target_bits > 0);
assert!(output.qp > 0);
assert!(output.lambda > 0.0);
}
#[test]
fn test_closed_loop_feedback() {
let mut controller = create_test_controller();
for i in 0..30 {
let frame_type = if i % 10 == 0 {
FrameType::Key
} else {
FrameType::Inter
};
let output = controller.get_rc(frame_type, 1.0);
let mut stats = FrameStats::new(i, frame_type);
stats.bits = output.target_bits;
stats.target_bits = output.target_bits;
stats.qp = output.qp;
stats.qp_f = output.qp_f;
stats.complexity = 1.0;
controller.update(&stats);
}
assert_eq!(controller.frame_count(), 30);
assert!(controller.current_bitrate() > 0.0);
}
#[test]
fn test_lookahead_buffer() {
let mut controller = create_test_controller();
controller.set_lookahead_size(20);
for _ in 0..25 {
controller.push_lookahead_frame(1.0, 1.0, 1.0, false);
}
assert!(controller.lookahead_frames.len() <= 20);
}
#[test]
fn test_vbv_compliance() {
let mut config = RcConfig::vbr(5_000_000, 10_000_000);
config.buffer_size = 5_000_000;
let mut controller = VbrController::new(&config);
controller.set_vbv_enabled(true);
let output = controller.get_rc(FrameType::Key, 1.0);
assert!(output.max_bits > 0);
let mut stats = FrameStats::new(0, FrameType::Key);
stats.bits = output.target_bits;
stats.target_bits = output.target_bits;
controller.update(&stats);
let fullness = controller.vbv_fullness();
assert!(fullness >= 0.0 && fullness <= 1.0);
}
#[test]
fn test_rate_prediction_model() {
let mut model = RatePredictionModel::default();
for i in 1..20 {
let complexity = i as f32 * 0.1;
let bits = 50_000 + i * 5_000;
model.update(complexity, bits);
}
let prediction = model.predict(1.5);
assert!(prediction > 0);
}
#[test]
fn test_pid_controller() {
let mut pid = PidControllerState::new();
let mut error = 1.0;
for _ in 0..10 {
let output = pid.calculate(error);
error -= output * 0.1; }
assert!(error.abs() < 1.0);
}
#[test]
fn test_adaptive_gop() {
let mut controller = create_test_controller();
controller.set_adaptive_gop(true);
controller.push_lookahead_frame(1.0, 1.0, 1.0, false);
controller.push_lookahead_frame(5.0, 5.0, 5.0, true);
let output = controller.get_rc(FrameType::Inter, 1.0);
assert!(output.qp > 0);
}
#[test]
fn test_two_pass_encoding() {
let mut controller = create_test_controller();
controller.set_pass(1);
for i in 0..30 {
let frame_type = if i % 10 == 0 {
FrameType::Key
} else {
FrameType::Inter
};
let _output = controller.get_rc(frame_type, 1.0 + (i as f32 % 3.0) * 0.2);
let mut stats = FrameStats::new(i, frame_type);
stats.bits = 100_000;
stats.target_bits = 100_000;
stats.spatial_complexity = 1.0;
stats.temporal_complexity = 1.0;
stats.complexity = 1.0;
controller.update(&stats);
}
let first_pass_data = controller.finalize_first_pass().expect("should succeed");
assert_eq!(first_pass_data.frame_count, 30);
assert!(!first_pass_data.suggested_bits.is_empty());
let mut controller2 = create_test_controller();
controller2.set_pass(2);
controller2.set_first_pass_data(first_pass_data);
let output = controller2.get_rc(FrameType::Key, 1.0);
assert!(output.target_bits > 0);
}
#[test]
fn test_frame_type_decision() {
let mut state = FrameTypeDecisionState::new(3);
let ft = state.decide_type(250, false, false);
assert_eq!(ft, FrameType::Key);
for _ in 0..3 {
let ft = state.decide_type(250, false, false);
assert!(matches!(ft, FrameType::BiDir | FrameType::Inter));
}
let ft = state.decide_type(250, true, true);
assert_eq!(ft, FrameType::Key);
}
#[test]
fn test_complexity_adjustment() {
let mut controller = create_test_controller();
let mut gop = GopStats::new(0, 0);
gop.average_complexity = 1.0;
controller.gop_history.push(gop);
let low = controller.calculate_complexity_adjustment(0.5);
let high = controller.calculate_complexity_adjustment(2.0);
assert!(high > low);
}
#[test]
fn test_bit_reservoir() {
let mut controller = create_test_controller();
for i in 0..10 {
let output = controller.get_rc(FrameType::Inter, 1.0);
let mut stats = FrameStats::new(i, FrameType::Inter);
stats.bits = output.target_bits / 2;
stats.target_bits = output.target_bits;
controller.update(&stats);
}
assert!(controller.bit_reservoir() > 0);
}
#[test]
fn test_reset() {
let mut controller = create_test_controller();
for i in 0..10 {
let mut stats = FrameStats::new(i, FrameType::Inter);
stats.bits = 100_000;
controller.update(&stats);
}
controller.reset();
assert_eq!(controller.frame_count(), 0);
assert_eq!(controller.bit_reservoir(), 0);
}
}