| 1 | #[derive(Default)] |
| 2 | pub struct RateControl { |
| 3 | bitrate: u32, |
| 4 | br_pool: u32, |
| 5 | f_pos: u32, |
| 6 | fracs: u32, |
| 7 | key_int: u32, |
| 8 | tb_num: u32, |
| 9 | tb_den: u32, |
| 10 | quality: Option<u8>, |
| 11 | } |
| 12 | |
| 13 | impl RateControl { |
| 14 | pub fn new() -> Self { Self{ key_int: 10, ..Default::default() } } |
| 15 | pub fn set_quality(&mut self, quality: u8) { |
| 16 | if quality > 0 { |
| 17 | self.quality = Some((quality - 1).min(15)); |
| 18 | } else { |
| 19 | self.quality = None; |
| 20 | } |
| 21 | } |
| 22 | pub fn set_bitrate(&mut self, br: u32, tb_num: u32, tb_den: u32) { |
| 23 | if tb_num > 0 && tb_den > 0 { |
| 24 | self.bitrate = br / 8; |
| 25 | self.tb_num = tb_num; |
| 26 | self.tb_den = tb_den; |
| 27 | } else { |
| 28 | self.bitrate = 0; |
| 29 | self.tb_num = 0; |
| 30 | self.tb_den = 0; |
| 31 | } |
| 32 | self.quality = Some(0); |
| 33 | self.reset(); |
| 34 | } |
| 35 | pub fn set_key_int(&mut self, key_int: u32) { |
| 36 | self.key_int = key_int; |
| 37 | self.reset(); |
| 38 | } |
| 39 | pub fn reset(&mut self) { |
| 40 | self.br_pool = self.bitrate; |
| 41 | self.f_pos = 0; |
| 42 | self.fracs = self.tb_den; |
| 43 | } |
| 44 | pub fn get_key_int(&self) -> u32 { self.key_int } |
| 45 | pub fn get_quant(&self, frameno: u32) -> (bool, Option<u8>) { |
| 46 | let is_intra = self.key_int == 0 || (frameno % self.key_int) == 0; |
| 47 | (is_intra, self.quality) |
| 48 | } |
| 49 | pub fn advance(&mut self, size: u32) { |
| 50 | if self.bitrate != 0 { |
| 51 | let expected = self.get_expected_size(); |
| 52 | let cur_quant = self.quality.unwrap_or(0); |
| 53 | if (size > expected + expected / 10) && (cur_quant < 7) { |
| 54 | self.quality = Some(cur_quant + 1); |
| 55 | } else if (size < expected - expected / 10) && (cur_quant > 0) { |
| 56 | self.quality = Some(cur_quant - 1); |
| 57 | } |
| 58 | |
| 59 | self.f_pos += self.tb_num; |
| 60 | while self.f_pos >= self.tb_den { |
| 61 | self.f_pos -= self.tb_den; |
| 62 | self.br_pool += self.bitrate; |
| 63 | self.fracs += self.tb_den; |
| 64 | } |
| 65 | self.fracs -= self.tb_num; |
| 66 | |
| 67 | self.br_pool = self.br_pool.saturating_sub(size).min(self.bitrate * 3); |
| 68 | } |
| 69 | } |
| 70 | pub fn get_expected_size(&self) -> u32 { |
| 71 | if self.bitrate != 0 { |
| 72 | self.br_pool * self.tb_num / self.fracs |
| 73 | } else { |
| 74 | 0 |
| 75 | } |
| 76 | } |
| 77 | } |