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