]>
Commit | Line | Data |
---|---|---|
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 | } |