1 use nihav_core::frame::FrameType;
3 pub struct RateDistMetric {
10 pub fn new() -> Self {
17 pub fn get_metric(&self, bits: u32, dist: u32) -> u32 {
18 ((bits as f32) + (dist as f32) * 0.1 * self.lambda).ceil() as u32
23 struct BitrateCounter {
31 fn new(intra: bool) -> Self {
44 for (q, dst) in self.factors.iter_mut().enumerate() {
46 *dst = (-0.1 * q + 2.95) / 100.0;
50 for (q, dst) in self.factors.iter_mut().enumerate() {
52 *dst = 100.0 / (8.2 * q * q + 51.0 * q + 3411.0);
56 fn init_metric(&self, metric: &mut RateDistMetric, q_add: usize) {
57 let q = (self.last_q + q_add).min(31);
58 const THRESHOLDS: [(u32, u32); 4] = [
59 (256, 8192), (128, 8192), (64, 4196), (32, 2048)
61 let (ge_thr, ps_thr) = THRESHOLDS[q / 8];
62 metric.good_enough = ge_thr;
63 metric.p_split_thr = ps_thr;
66 fn update_stats(&mut self, fsize: usize) {
67 if fsize < self.proj_size - self.proj_size / 8 {
68 let mut inv_fac = 1.0 / self.factors[self.last_q];
72 self.factors[self.last_q] = 1.0 / inv_fac;
73 } else if fsize > self.proj_size + self.proj_size / 8 {
74 let mut inv_fac = 1.0 / self.factors[self.last_q];
78 self.factors[self.last_q] = 1.0 / inv_fac;
81 fn get_est_size(&self, complexity: u32, q: usize) -> usize {
82 ((complexity as f32) * self.factors[q]).ceil() as usize
84 fn get_quant(&mut self, target: usize, complexity: u32) -> usize {
85 let tgt_31 = self.get_est_size(complexity, 31);
86 let tgt_0 = self.get_est_size(complexity, 0);
89 self.proj_size = tgt_31;
90 } else if target > tgt_0 {
92 self.proj_size = tgt_0;
93 } else { //xxx: do binary search?
94 for q in (0..31).rev() {
95 let expected_size = self.get_est_size(complexity, q);
96 if target >= (expected_size - expected_size / 8) &&
97 target <= (expected_size + expected_size / 8) {
98 self.proj_size = expected_size;
105 fn get_last_quant(&self) -> usize { self.last_q }
108 const TIMEBASE: u32 = 1000;
110 pub struct BitRateControl {
111 force_quant: Option<usize>,
112 force_quality: Option<u8>,
113 br_counter: [BitrateCounter; 2],
125 impl BitRateControl {
126 pub fn new() -> Self {
130 br_counter: [BitrateCounter::new(true), BitrateCounter::new(false)],
142 pub fn rate_ctl_in_use(&self) -> bool {
143 self.force_quant.is_none() && self.force_quality.is_none() && self.bitrate != 0
145 pub fn set_bitrate(&mut self, bitrate: u32) {
146 self.bitrate = bitrate;
147 for br in self.br_counter.iter_mut() {
151 self.bitpool = (self.bitrate as usize) * 2;
154 pub fn set_force_quant(&mut self, force_q: Option<usize>) { self.force_quant = force_q; }
155 pub fn get_force_quant(&self) -> i8 {
156 if let Some(q) = self.force_quant {
162 pub fn set_force_quality(&mut self, force_q: Option<u8>) { self.force_quality = force_q; }
163 pub fn get_force_quality(&self) -> i8 {
164 if let Some(q) = self.force_quality {
170 pub fn get_quant(&mut self, ftype: FrameType, complexity: u32) -> usize {
171 if let Some(q) = self.force_quant {
173 } else if self.force_quality.is_some() {
175 } else if ftype != FrameType::B {
176 let tgt = self.get_target_size(ftype);
177 self.br_counter[if ftype == FrameType::I { 0 } else { 1 }].get_quant(tgt, complexity)
179 (self.br_counter[1].get_last_quant() + self.b_offset).min(31)
182 pub fn get_last_quant(&self, ftype: FrameType) -> usize {
184 FrameType::I => self.br_counter[0].get_last_quant(),
185 FrameType::P => self.br_counter[1].get_last_quant(),
186 _ => (self.br_counter[1].get_last_quant() + self.b_offset).min(31),
189 pub fn init_metric(&self, ftype: FrameType, metric: &mut RateDistMetric) {
190 if let Some(q) = self.force_quality {
191 metric.lambda = (q as f32) / 50.0;
195 self.br_counter[0].init_metric(metric, 0);
198 self.br_counter[1].init_metric(metric, 0);
201 self.br_counter[1].init_metric(metric, self.b_offset);
206 pub fn update_stats(&mut self, ftype: FrameType, fsize: usize, ts_diff: u32) {
207 if self.bitrate > 0 {
208 if ts_diff > 0 && self.duration < std::u32::MAX / 2 {
209 self.duration += ts_diff;
212 self.tpos += ts_diff;
213 while self.tpos >= TIMEBASE {
214 self.tpos -= TIMEBASE;
215 self.bitpool += self.bitrate as usize;
217 self.bitpool = self.bitpool.saturating_sub(fsize * 8).max(1024);
220 FrameType::I => self.br_counter[0].update_stats(fsize),
221 FrameType::P => self.br_counter[1].update_stats(fsize),
225 pub fn get_target_size(&self, ftype: FrameType) -> usize {
226 if self.bitrate == 0 || self.bitpool == 0 {
229 let bitpool_limit = (self.bitrate + self.bitrate / 8) as usize;
230 let bitpool_avail = self.bitpool.min(bitpool_limit);
231 let target_size = if self.dcount > 0 {
232 let avg_len = ((self.duration + self.dcount / 2) / self.dcount).max(1);
233 bitpool_avail * (avg_len as usize) / ((TIMEBASE - self.tpos) as usize)
237 let tgt_bits = match ftype {
238 FrameType::I => target_size * 3,
239 FrameType::B => target_size * 3 / 4,