nihav-duck/src/codecs/vp7enc/rdo.rs

   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
 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 }