projects / nihav.git / blame
| Commit | Line | Data |
|---|---|---|
| 4965a5e5 KS |
1 | use nihav_core::frame::FrameType; |
| 2 | ||
| 3 | pub struct RateDistMetric { | |
| 4 | pub lambda: f32, | |
| 5 | pub good_enough: u32, | |
| 6 | pub p_split_thr: u32, | |
| 7 | } | |
| 8 | ||
| 9 | impl RateDistMetric { | |
| 10 | pub fn new() -> Self { | |
| 11 | Self { | |
| 12 | lambda: 1.0, | |
| 13 | good_enough: 256, | |
| 14 | p_split_thr: 8192, | |
| 15 | } | |
| 16 | } | |
| 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 | |
| 19 | } | |
| 20 | } | |
| 21 | ||
| 22 | #[derive(Clone,Copy)] | |
| 23 | struct BitrateCounter { | |
| 24 | factors: [f32; 32], | |
| 25 | last_q: usize, | |
| 26 | proj_size: usize, | |
| 27 | intra: bool, | |
| 28 | } | |
| 29 | ||
| 30 | impl BitrateCounter { | |
| 31 | fn new(intra: bool) -> Self { | |
| 32 | let mut obj = Self { | |
| 33 | factors: [0.0; 32], | |
| 34 | last_q: 0, | |
| 35 | proj_size: 0, | |
| 36 | intra | |
| 37 | }; | |
| 38 | obj.reset(); | |
| 39 | obj | |
| 40 | } | |
| 41 | fn reset(&mut self) { | |
| 42 | if self.intra { | |
| 43 | self.last_q = 8; | |
| 44 | for (q, dst) in self.factors.iter_mut().enumerate() { | |
| 45 | let q = q as f32; | |
| 46 | *dst = (-0.1 * q + 2.95) / 100.0; | |
| 47 | } | |
| 48 | } else { | |
| 49 | self.last_q = 10; | |
| 50 | for (q, dst) in self.factors.iter_mut().enumerate() { | |
| 51 | let q = q as f32; | |
| 52 | *dst = 100.0 / (8.2 * q * q + 51.0 * q + 3411.0); | |
| 53 | } | |
| 54 | } | |
| 55 | } | |
| 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) | |
| 60 | ]; | |
| 61 | let (ge_thr, ps_thr) = THRESHOLDS[q / 8]; | |
| 62 | metric.good_enough = ge_thr; | |
| 63 | metric.p_split_thr = ps_thr; | |
| 64 | metric.lambda = 1.0; | |
| 65 | } | |
| 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]; | |
| 69 | if inv_fac > 1.0 { | |
| 70 | inv_fac -= 0.5; | |
| 71 | } | |
| 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]; | |
| 75 | if inv_fac < 200.0 { | |
| 76 | inv_fac += 0.5; | |
| 77 | } | |
| 78 | self.factors[self.last_q] = 1.0 / inv_fac; | |
| 79 | } | |
| 80 | } | |
| 81 | fn get_est_size(&self, complexity: u32, q: usize) -> usize { | |
| 82 | ((complexity as f32) * self.factors[q]).ceil() as usize | |
| 83 | } | |
| 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); | |
| 87 | if target < tgt_31 { | |
| 88 | self.last_q = 31; | |
| 89 | self.proj_size = tgt_31; | |
| 90 | } else if target > tgt_0 { | |
| 91 | self.last_q = 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; | |
| 99 | self.last_q = q; | |
| 100 | } | |
| 101 | } | |
| 102 | } | |
| 103 | self.last_q | |
| 104 | } | |
| 105 | fn get_last_quant(&self) -> usize { self.last_q } | |
| 106 | } | |
| 107 | ||
| 108 | const TIMEBASE: u32 = 1000; | |
| 109 | ||
| 110 | pub struct BitRateControl { | |
| 111 | force_quant: Option<usize>, | |
| 112 | force_quality: Option<u8>, | |
| 113 | br_counter: [BitrateCounter; 2], | |
| 114 | ||
| 115 | bitrate: u32, | |
| 116 | tpos: u32, | |
| 117 | bitpool: usize, | |
| 118 | ||
| 119 | duration: u32, | |
| 120 | dcount: u32, | |
| 121 | ||
| 122 | pub b_offset: usize, | |
| 123 | } | |
| 124 | ||
| 125 | impl BitRateControl { | |
| 126 | pub fn new() -> Self { | |
| 127 | Self { | |
| 128 | force_quant: None, | |
| 129 | force_quality: None, | |
| 130 | br_counter: [BitrateCounter::new(true), BitrateCounter::new(false)], | |
| 131 | ||
| 132 | bitrate: 0, | |
| 133 | tpos: 0, | |
| 134 | bitpool: 0, | |
| 135 | ||
| 136 | duration: 0, | |
| 137 | dcount: 0, | |
| 138 | ||
| 139 | b_offset: 4, | |
| 140 | } | |
| 141 | } | |
| 142 | pub fn rate_ctl_in_use(&self) -> bool { | |
| 143 | self.force_quant.is_none() && self.force_quality.is_none() && self.bitrate != 0 | |
| 144 | } | |
| 145 | pub fn set_bitrate(&mut self, bitrate: u32) { | |
| 146 | self.bitrate = bitrate; | |
| 147 | for br in self.br_counter.iter_mut() { | |
| 148 | br.reset(); | |
| 149 | } | |
| 150 | ||
| 151 | self.bitpool = (self.bitrate as usize) * 2; | |
| 152 | self.tpos = 0; | |
| 153 | } | |
| 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 { | |
| 157 | q as i8 | |
| 158 | } else { | |
| 159 | -1 | |
| 160 | } | |
| 161 | } | |
| 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 { | |
| 165 | q as i8 | |
| 166 | } else { | |
| 167 | -1 | |
| 168 | } | |
| 169 | } | |
| 170 | pub fn get_quant(&mut self, ftype: FrameType, complexity: u32) -> usize { | |
| 171 | if let Some(q) = self.force_quant { | |
| 172 | q | |
| 173 | } else if self.force_quality.is_some() { | |
| 174 | 4 | |
| 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) | |
| 178 | } else { | |
| 179 | (self.br_counter[1].get_last_quant() + self.b_offset).min(31) | |
| 180 | } | |
| 181 | } | |
| 182 | pub fn get_last_quant(&self, ftype: FrameType) -> usize { | |
| 183 | match ftype { | |
| 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), | |
| 187 | } | |
| 188 | } | |
| 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; | |
| 192 | } else { | |
| 193 | match ftype { | |
| 194 | FrameType::I => { | |
| 195 | self.br_counter[0].init_metric(metric, 0); | |
| 196 | }, | |
| 197 | FrameType::P => { | |
| 198 | self.br_counter[1].init_metric(metric, 0); | |
| 199 | }, | |
| 200 | _ => { | |
| 201 | self.br_counter[1].init_metric(metric, self.b_offset); | |
| 202 | }, | |
| 203 | }; | |
| 204 | } | |
| 205 | } | |
| 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; | |
| 210 | self.dcount += 1; | |
| 211 | } | |
| 212 | self.tpos += ts_diff; | |
| 213 | while self.tpos >= TIMEBASE { | |
| 214 | self.tpos -= TIMEBASE; | |
| 215 | self.bitpool += self.bitrate as usize; | |
| 216 | } | |
| 217 | self.bitpool = self.bitpool.saturating_sub(fsize * 8).max(1024); | |
| 218 | } | |
| 219 | match ftype { | |
| 220 | FrameType::I => self.br_counter[0].update_stats(fsize), | |
| 221 | FrameType::P => self.br_counter[1].update_stats(fsize), | |
| 222 | _ => {}, | |
| 223 | }; | |
| 224 | } | |
| 225 | pub fn get_target_size(&self, ftype: FrameType) -> usize { | |
| 226 | if self.bitrate == 0 || self.bitpool == 0 { | |
| 227 | return 0; | |
| 228 | } | |
| 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) | |
| 234 | } else { | |
| 235 | bitpool_avail / 10 | |
| 236 | }; | |
| 237 | let tgt_bits = match ftype { | |
| 238 | FrameType::I => target_size * 3, | |
| 239 | FrameType::B => target_size * 3 / 4, | |
| 240 | _ => target_size, | |
| 241 | }; | |
| 242 | (tgt_bits + 7) / 8 | |
| 243 | } | |
| 244 | } |