1use alloc::boxed::Box;
2use core::num::NonZeroU16;
3
4#[derive(Debug, Clone)]
8pub enum QuantizationTableType {
9 Default,
11
12 Flat,
14
15 CustomMsSsim,
17
18 CustomPsnrHvs,
20
21 ImageMagick,
25
26 KleinSilversteinCarney,
28
29 DentalXRays,
31
32 VisualDetectionModel,
34
35 ImprovedDetectionModel,
37
38 Custom(Box<[u16; 64]>),
40}
41
42impl QuantizationTableType {
43 fn index(&self) -> usize {
44 use QuantizationTableType::*;
45
46 match self {
47 Default => 0,
48 Flat => 1,
49 CustomMsSsim => 2,
50 CustomPsnrHvs => 3,
51 ImageMagick => 4,
52 KleinSilversteinCarney => 5,
53 DentalXRays => 6,
54 VisualDetectionModel => 7,
55 ImprovedDetectionModel => 8,
56 Custom(_) => panic!("Custom types not supported"),
57 }
58 }
59}
60
61static DEFAULT_LUMA_TABLES: [[u16; 64]; 9] = [
63 [
64 16, 11, 10, 16, 24, 40, 51, 61, 12, 12, 14, 19, 26, 58, 60, 55, 14, 13, 16, 24, 40, 57, 69,
66 56, 14, 17, 22, 29, 51, 87, 80, 62, 18, 22, 37, 56, 68, 109, 103, 77, 24, 35, 55, 64, 81,
67 104, 113, 92, 49, 64, 78, 87, 103, 121, 120, 101, 72, 92, 95, 98, 112, 100, 103, 99,
68 ],
69 [
70 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16,
72 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16,
73 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16,
74 ],
75 [
76 12, 17, 20, 21, 30, 34, 56, 63, 18, 20, 20, 26, 28, 51, 61, 55, 19, 20, 21, 26, 33, 58, 69,
78 55, 26, 26, 26, 30, 46, 87, 86, 66, 31, 33, 36, 40, 46, 96, 100, 73, 40, 35, 46, 62, 81,
79 100, 111, 91, 46, 66, 76, 86, 102, 121, 120, 101, 68, 90, 90, 96, 113, 102, 105, 103,
80 ],
81 [
82 9, 10, 12, 14, 27, 32, 51, 62, 11, 12, 14, 19, 27, 44, 59, 73, 12, 14, 18, 25, 42, 59, 79,
84 78, 17, 18, 25, 42, 61, 92, 87, 92, 23, 28, 42, 75, 79, 112, 112, 99, 40, 42, 59, 84, 88,
85 124, 132, 111, 42, 64, 78, 95, 105, 126, 125, 99, 70, 75, 100, 102, 116, 100, 107, 98,
86 ],
87 [
88 16, 16, 16, 18, 25, 37, 56, 85, 16, 17, 20, 27, 34, 40, 53, 75, 16, 20, 24, 31, 43, 62, 91,
91 135, 18, 27, 31, 40, 53, 74, 106, 156, 25, 34, 43, 53, 69, 94, 131, 189, 37, 40, 62, 74,
92 94, 124, 169, 238, 56, 53, 91, 106, 131, 169, 226, 311, 85, 75, 135, 156, 189, 238, 311,
93 418,
94 ],
95 [
96 10, 12, 14, 19, 26, 38, 57, 86, 12, 18, 21, 28, 35, 41, 54, 76, 14, 21, 25, 32, 44, 63, 92,
98 136, 19, 28, 32, 41, 54, 75, 107, 157, 26, 35, 44, 54, 70, 95, 132, 190, 38, 41, 63, 75,
99 95, 125, 170, 239, 57, 54, 92, 107, 132, 170, 227, 312, 86, 76, 136, 157, 190, 239, 312,
100 419,
101 ],
102 [
103 7, 8, 10, 14, 23, 44, 95, 241, 8, 8, 11, 15, 25, 47, 102, 255, 10, 11, 13, 19, 31, 58, 127,
105 255, 14, 15, 19, 27, 44, 83, 181, 255, 23, 25, 31, 44, 72, 136, 255, 255, 44, 47, 58, 83,
106 136, 255, 255, 255, 95, 102, 127, 181, 255, 255, 255, 255, 241, 255, 255, 255, 255, 255,
107 255, 255,
108 ],
109 [
110 15, 11, 11, 12, 15, 19, 25, 32, 11, 13, 10, 10, 12, 15, 19, 24, 11, 10, 14, 14, 16, 18, 22,
112 27, 12, 10, 14, 18, 21, 24, 28, 33, 15, 12, 16, 21, 26, 31, 36, 42, 19, 15, 18, 24, 31, 38,
113 45, 53, 25, 19, 22, 28, 36, 45, 55, 65, 32, 24, 27, 33, 42, 53, 65, 77,
114 ],
115 [
116 14, 10, 11, 14, 19, 25, 34, 45, 10, 11, 11, 12, 15, 20, 26, 33, 11, 11, 15, 18, 21, 25, 31,
118 38, 14, 12, 18, 24, 28, 33, 39, 47, 19, 15, 21, 28, 36, 43, 51, 59, 25, 20, 25, 33, 43, 54,
119 64, 74, 34, 26, 31, 39, 51, 64, 77, 91, 45, 33, 38, 47, 59, 74, 91, 108,
120 ],
121];
122
123static DEFAULT_CHROMA_TABLES: [[u16; 64]; 9] = [
125 [
126 17, 18, 24, 47, 99, 99, 99, 99, 18, 21, 26, 66, 99, 99, 99, 99, 24, 26, 56, 99, 99, 99, 99,
128 99, 47, 66, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99,
129 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99,
130 ],
131 [
132 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16,
134 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16,
135 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16,
136 ],
137 [
138 8, 12, 15, 15, 86, 96, 96, 98, 13, 13, 15, 26, 90, 96, 99, 98, 12, 15, 18, 96, 99, 99, 99,
140 99, 17, 16, 90, 96, 99, 99, 99, 99, 96, 96, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99,
141 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99,
142 ],
143 [
144 9, 10, 17, 19, 62, 89, 91, 97, 12, 13, 18, 29, 84, 91, 88, 98, 14, 19, 29, 93, 95, 95, 98,
146 97, 20, 26, 84, 88, 95, 95, 98, 94, 26, 86, 91, 93, 97, 99, 98, 99, 99, 100, 98, 99, 99,
147 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 97, 97, 99, 99, 99, 99, 97, 99,
148 ],
149 [
150 16, 16, 16, 18, 25, 37, 56, 85, 16, 17, 20, 27, 34, 40, 53, 75, 16, 20, 24, 31, 43, 62, 91,
153 135, 18, 27, 31, 40, 53, 74, 106, 156, 25, 34, 43, 53, 69, 94, 131, 189, 37, 40, 62, 74,
154 94, 124, 169, 238, 56, 53, 91, 106, 131, 169, 226, 311, 85, 75, 135, 156, 189, 238, 311,
155 418,
156 ],
157 [
158 10, 12, 14, 19, 26, 38, 57, 86, 12, 18, 21, 28, 35, 41, 54, 76, 14, 21, 25, 32, 44, 63, 92,
160 136, 19, 28, 32, 41, 54, 75, 107, 157, 26, 35, 44, 54, 70, 95, 132, 190, 38, 41, 63, 75,
161 95, 125, 170, 239, 57, 54, 92, 107, 132, 170, 227, 312, 86, 76, 136, 157, 190, 239, 312,
162 419,
163 ],
164 [
165 7, 8, 10, 14, 23, 44, 95, 241, 8, 8, 11, 15, 25, 47, 102, 255, 10, 11, 13, 19, 31, 58, 127,
167 255, 14, 15, 19, 27, 44, 83, 181, 255, 23, 25, 31, 44, 72, 136, 255, 255, 44, 47, 58, 83,
168 136, 255, 255, 255, 95, 102, 127, 181, 255, 255, 255, 255, 241, 255, 255, 255, 255, 255,
169 255, 255,
170 ],
171 [
172 15, 11, 11, 12, 15, 19, 25, 32, 11, 13, 10, 10, 12, 15, 19, 24, 11, 10, 14, 14, 16, 18, 22,
174 27, 12, 10, 14, 18, 21, 24, 28, 33, 15, 12, 16, 21, 26, 31, 36, 42, 19, 15, 18, 24, 31, 38,
175 45, 53, 25, 19, 22, 28, 36, 45, 55, 65, 32, 24, 27, 33, 42, 53, 65, 77,
176 ],
177 [
178 14, 10, 11, 14, 19, 25, 34, 45, 10, 11, 11, 12, 15, 20, 26, 33, 11, 11, 15, 18, 21, 25, 31,
180 38, 14, 12, 18, 24, 28, 33, 39, 47, 19, 15, 21, 28, 36, 43, 51, 59, 25, 20, 25, 33, 43, 54,
181 64, 74, 34, 26, 31, 39, 51, 64, 77, 91, 45, 33, 38, 47, 59, 74, 91, 108,
182 ],
183];
184
185const SHIFT: u32 = 2 * 8 - 1;
186
187fn compute_reciprocal(divisor: u32) -> (i32, i32) {
188 if divisor <= 1 {
189 return (1, 0);
190 }
191
192 let mut reciprocals = (1 << SHIFT) / divisor;
193 let fractional = (1 << SHIFT) % divisor;
194
195 let mut correction = divisor / 2;
197
198 if fractional != 0 {
199 if fractional <= correction {
200 correction += 1;
201 } else {
202 reciprocals += 1;
203 }
204 }
205
206 (reciprocals as i32, correction as i32)
207}
208
209pub struct QuantizationTable {
210 table: [NonZeroU16; 64],
211 reciprocals: [i32; 64],
212 corrections: [i32; 64],
213}
214
215impl QuantizationTable {
216 pub fn new_with_quality(
217 table: &QuantizationTableType,
218 quality: u8,
219 luma: bool,
220 ) -> QuantizationTable {
221 let table = match table {
222 QuantizationTableType::Custom(table) => Self::get_user_table(table),
223 table => {
224 let table = if luma {
225 &DEFAULT_LUMA_TABLES[table.index()]
226 } else {
227 &DEFAULT_CHROMA_TABLES[table.index()]
228 };
229 Self::get_with_quality(table, quality)
230 }
231 };
232
233 let mut reciprocals = [0i32; 64];
234 let mut corrections = [0i32; 64];
235
236 for i in 0..64 {
237 let (reciprocal, correction) = compute_reciprocal(table[i].get() as u32);
238
239 reciprocals[i] = reciprocal;
240 corrections[i] = correction;
241 }
242
243 QuantizationTable {
244 table,
245 reciprocals,
246 corrections,
247 }
248 }
249
250 fn get_user_table(table: &[u16; 64]) -> [NonZeroU16; 64] {
251 let mut q_table = [NonZeroU16::new(1).unwrap(); 64];
252 for (i, &v) in table.iter().enumerate() {
253 q_table[i] = match NonZeroU16::new(v.max(1).min(2 << 10) << 3) {
254 Some(v) => v,
255 None => panic!("Invalid quantization table value: {}", v),
256 };
257 }
258 q_table
259 }
260
261 fn get_with_quality(table: &[u16; 64], quality: u8) -> [NonZeroU16; 64] {
262 let quality = quality.max(1).min(100) as u32;
263
264 let scale = if quality < 50 {
265 5000 / quality
266 } else {
267 200 - quality * 2
268 };
269
270 let mut q_table = [NonZeroU16::new(1).unwrap(); 64];
271
272 for (i, &v) in table.iter().enumerate() {
273 let v = v as u32;
274
275 let v = (v * scale + 50) / 100;
276
277 let v = v.max(1).min(255) as u16;
278
279 q_table[i] = NonZeroU16::new(v << 3).unwrap();
281 }
282 q_table
283 }
284
285 #[inline]
286 pub fn get(&self, index: usize) -> u8 {
287 (self.table[index].get() >> 3) as u8
288 }
289
290 #[inline]
291 pub fn quantize(&self, in_value: i16, index: usize) -> i16 {
292 let value = in_value as i32;
293
294 let reciprocal = self.reciprocals[index];
295 let corrections = self.corrections[index];
296
297 let abs_value = value.abs();
298
299 let mut product = (abs_value + corrections) * reciprocal;
300 product >>= SHIFT;
301
302 if value != abs_value {
303 product *= -1;
304 }
305
306 product as i16
307 }
308}
309
310#[cfg(test)]
311mod tests {
312 use crate::quantization::{QuantizationTable, QuantizationTableType};
313
314 #[test]
315 fn test_new_100() {
316 let q = QuantizationTable::new_with_quality(&QuantizationTableType::Default, 100, true);
317
318 for &v in &q.table {
319 let v = v.get();
320 assert_eq!(v, 1 << 3);
321 }
322
323 let q = QuantizationTable::new_with_quality(&QuantizationTableType::Default, 100, false);
324
325 for &v in &q.table {
326 let v = v.get();
327 assert_eq!(v, 1 << 3);
328 }
329 }
330
331 #[test]
332 fn test_new_100_quantize() {
333 let q = QuantizationTable::new_with_quality(&QuantizationTableType::Default, 100, true);
334
335 for i in -255..255 {
336 assert_eq!(i, q.quantize(i << 3, 0));
337 }
338 }
339}