| 1 | use super::blocks::*; |
| 2 | use super::coder::*; |
| 3 | |
| 4 | static mut COUNTER: usize = 0; |
| 5 | |
| 6 | pub const SMALL_DIST: u32 = 256; |
| 7 | pub const MAX_DIST: u32 = std::u32::MAX; |
| 8 | |
| 9 | const INTER_TO_INTRA_RATIO: f32 = 0.85; |
| 10 | |
| 11 | pub struct RateDistMetric { |
| 12 | pub lambda: f32, |
| 13 | } |
| 14 | |
| 15 | impl RateDistMetric { |
| 16 | pub fn new() -> Self { |
| 17 | Self { |
| 18 | lambda: 1.0, |
| 19 | } |
| 20 | } |
| 21 | pub fn calc_metric(&self, dist: u32, nits: u32) -> u32 { |
| 22 | ((dist as f32) + self.lambda * (nits as f32) + 0.5) as u32 |
| 23 | } |
| 24 | pub fn adjust_br(&mut self, cur_size: usize, tgt_size: usize) { |
| 25 | let low_limit = tgt_size - tgt_size / 8; |
| 26 | let up_limit = tgt_size + tgt_size / 8; |
| 27 | if cur_size < low_limit { |
| 28 | self.lambda = (self.lambda - 0.1).max(0.0); |
| 29 | } else if cur_size > up_limit { |
| 30 | self.lambda = (self.lambda + 0.1).min(16.0); |
| 31 | } |
| 32 | } |
| 33 | |
| 34 | pub fn block_dist(&self, src: &[u8; 16], new: &[u8; 16], q: usize, ctype: usize, pctx: u8, probs: &[[[u8; 11]; 3]; 8]) -> (u32, bool) { |
| 35 | let mut diff = [0i16; 16]; |
| 36 | get_block_difference(&mut diff, src, new); |
| 37 | diff.fdct(); |
| 38 | diff.quant(q, ctype); |
| 39 | let has_nz = diff.has_nz(); |
| 40 | let nits = estimate_subblock_nits(&diff, ctype, pctx, probs); |
| 41 | diff.dequant(q, ctype); |
| 42 | diff.idct(); |
| 43 | let dist = get_difference_dist(src, new, &diff); |
| 44 | unsafe {COUNTER += 1;} |
| 45 | (self.calc_metric(dist, nits), has_nz) |
| 46 | } |
| 47 | } |
| 48 | |
| 49 | #[derive(Default)] |
| 50 | pub struct BitRateControl { |
| 51 | tb_num: u32, |
| 52 | tb_den: u32, |
| 53 | key_int: u32, |
| 54 | bitrate: u32, |
| 55 | force_q: Option<usize>, |
| 56 | bitpool: u32, |
| 57 | fpos: u32, |
| 58 | kpos: u32, |
| 59 | num_mb: u32, |
| 60 | } |
| 61 | |
| 62 | impl BitRateControl { |
| 63 | pub fn new() -> Self { Self::default() } |
| 64 | fn reset(&mut self) { |
| 65 | self.fpos = 0; |
| 66 | self.kpos = 0; |
| 67 | self.bitpool = self.bitrate; |
| 68 | } |
| 69 | pub fn set_params(&mut self, tb_num: u32, tb_den: u32, bitrate: u32, key_int: u8, num_mb: usize) { |
| 70 | self.tb_num = tb_num; |
| 71 | self.tb_den = tb_den; |
| 72 | self.bitrate = bitrate; |
| 73 | self.key_int = u32::from(key_int); |
| 74 | self.num_mb = num_mb as u32; |
| 75 | self.reset(); |
| 76 | } |
| 77 | pub fn has_bitrate(&self) -> bool { self.bitrate != 0 } |
| 78 | pub fn get_quant(&self) -> Option<usize> { self.force_q } |
| 79 | pub fn set_quant(&mut self, q: Option<usize>) { |
| 80 | if self.force_q != q { |
| 81 | self.force_q = q; |
| 82 | self.reset(); |
| 83 | } |
| 84 | } |
| 85 | pub fn set_key_interval(&mut self, key_int: u8) { |
| 86 | let key_int = u32::from(key_int); |
| 87 | if self.key_int != key_int { |
| 88 | self.key_int = key_int; |
| 89 | self.reset(); |
| 90 | } |
| 91 | } |
| 92 | pub fn get_target_size(&self, is_intra: bool) -> u32 { |
| 93 | if self.bitrate != 0 && self.force_q.is_none() { |
| 94 | let pool_frames = self.tb_den - self.fpos; |
| 95 | if self.key_int <= 1 { // all intra |
| 96 | if self.bitpool == 0 || pool_frames == 0 { |
| 97 | self.bitrate * self.tb_num / self.tb_den |
| 98 | } else { |
| 99 | self.bitpool / pool_frames |
| 100 | } |
| 101 | } else { |
| 102 | let full_gop_weight = 1.0 + ((self.key_int - 1) as f32) * INTER_TO_INTRA_RATIO; |
| 103 | let i_bits = if self.bitpool == 0 || pool_frames == 0 { |
| 104 | let gop_size = self.bitrate * self.tb_num * self.key_int / self.tb_den; |
| 105 | (gop_size as f32) / full_gop_weight |
| 106 | } else { |
| 107 | let full_gops = pool_frames / self.key_int; |
| 108 | let weight = (full_gops as f32) * full_gop_weight + ((pool_frames % self.key_int) as f32) * INTER_TO_INTRA_RATIO; |
| 109 | (self.bitpool as f32) / weight |
| 110 | }; |
| 111 | if is_intra { |
| 112 | (i_bits + 0.5) as u32 |
| 113 | } else { |
| 114 | (i_bits * INTER_TO_INTRA_RATIO + 0.5) as u32 |
| 115 | } |
| 116 | } |
| 117 | } else { |
| 118 | 0 |
| 119 | } |
| 120 | } |
| 121 | fn pred_nits_per_mb(is_intra: bool, q: usize) -> f32 { |
| 122 | let fq = q as f32; |
| 123 | match (is_intra, q) { |
| 124 | (true, 0..=6) => 3434.0 + fq * fq * 7.5 - fq * 195.0, |
| 125 | (true, _) => 2500.0 - (fq - 6.0).ln() * 500.0, |
| 126 | (false, 0..=10) => 1595.0 + fq * fq * 3.4 - fq * 125.0, |
| 127 | (false, _) => 800.0 - (fq - 8.0).ln() * 155.0, |
| 128 | } |
| 129 | } |
| 130 | #[allow(dead_code)] |
| 131 | // todo use for refining maybe |
| 132 | pub fn predict_size(&self, is_intra: bool, q: usize) -> u32 { |
| 133 | let min_size = if is_intra { 200 * 8 } else { 50 * 8 }; |
| 134 | let nits_per_mb = Self::pred_nits_per_mb(is_intra, q); |
| 135 | ((nits_per_mb * (self.num_mb as f32) / 8.0) as u32).max(min_size) |
| 136 | } |
| 137 | pub fn get_frame_quant(&self, is_intra: bool) -> usize { |
| 138 | if let Some(q) = self.force_q { |
| 139 | q as usize |
| 140 | } else { |
| 141 | let expected_size = self.get_target_size(is_intra); |
| 142 | let nits_per_mb = ((expected_size * 8) as f32) / (self.num_mb as f32); |
| 143 | if is_intra { |
| 144 | if nits_per_mb > 2500.0 { // simple search |
| 145 | if nits_per_mb > Self::pred_nits_per_mb(is_intra, 3) { |
| 146 | if nits_per_mb > Self::pred_nits_per_mb(is_intra, 1) { |
| 147 | 0 |
| 148 | } else if nits_per_mb > Self::pred_nits_per_mb(is_intra, 2) { |
| 149 | 1 |
| 150 | } else { |
| 151 | 2 |
| 152 | } |
| 153 | } else { |
| 154 | if nits_per_mb > Self::pred_nits_per_mb(is_intra, 4) { |
| 155 | 3 |
| 156 | } else if nits_per_mb > Self::pred_nits_per_mb(is_intra, 5) { |
| 157 | 4 |
| 158 | } else { |
| 159 | 5 |
| 160 | } |
| 161 | } |
| 162 | } else { |
| 163 | ((((2500.0 - nits_per_mb) / 500.0).exp() + 6.0) as usize).min(127) |
| 164 | } |
| 165 | } else { |
| 166 | if nits_per_mb > 680.0 { // simple search |
| 167 | let (start, end) = if nits_per_mb > Self::pred_nits_per_mb(is_intra, 5) { |
| 168 | if nits_per_mb > Self::pred_nits_per_mb(is_intra, 3) { |
| 169 | (0, 3) |
| 170 | } else { |
| 171 | (3, 5) |
| 172 | } |
| 173 | } else if nits_per_mb > Self::pred_nits_per_mb(is_intra, 7) { |
| 174 | (5, 7) |
| 175 | } else { |
| 176 | (7, 10) |
| 177 | }; |
| 178 | let mut q = end; |
| 179 | for qq in start..end { |
| 180 | if nits_per_mb > Self::pred_nits_per_mb(is_intra, qq) { |
| 181 | q = qq; |
| 182 | break; |
| 183 | } |
| 184 | } |
| 185 | q |
| 186 | } else { |
| 187 | ((((800.0 - nits_per_mb) / 155.0).exp() + 6.0) as usize).max(10).min(127) |
| 188 | } |
| 189 | } |
| 190 | } |
| 191 | } |
| 192 | pub fn update(&mut self, size: usize) { |
| 193 | self.kpos += 1; |
| 194 | if self.kpos == self.key_int { |
| 195 | self.kpos = 0; |
| 196 | } |
| 197 | if self.bitrate == 0 || self.force_q.is_some() { |
| 198 | return; |
| 199 | } |
| 200 | self.fpos += self.tb_num; |
| 201 | while self.fpos >= self.tb_den { |
| 202 | self.fpos -= self.tb_den; |
| 203 | self.bitpool += self.bitrate; |
| 204 | } |
| 205 | self.bitpool = self.bitpool.saturating_sub((size * 8) as u32); |
| 206 | } |
| 207 | } |