]>
Commit | Line | Data |
---|---|---|
4b459d0b KS |
1 | use super::Pixel; |
2 | ||
3 | pub struct NeuQuantQuantiser { | |
4 | weights: [[f64; 3]; 256], | |
5 | freq: [f64; 256], | |
6 | bias: [f64; 256], | |
7 | factor: usize, | |
8 | } | |
9 | ||
10 | const SPECIAL_NODES: usize = 2; | |
11 | impl NeuQuantQuantiser { | |
12 | pub fn new(factor: usize) -> Self { | |
13 | let mut weights = [[0.0; 3]; 256]; | |
14 | if SPECIAL_NODES > 1 { | |
15 | weights[1] = [255.0; 3]; // for white | |
16 | } | |
17 | for i in SPECIAL_NODES..256 { | |
18 | let w = 255.0 * ((i - SPECIAL_NODES) as f64) / ((256 - SPECIAL_NODES) as f64); | |
19 | weights[i] = [w, w, w]; | |
20 | } | |
21 | Self { | |
22 | weights, | |
23 | freq: [1.0 / 256.0; 256], | |
24 | bias: [0.0; 256], | |
25 | factor, | |
26 | } | |
27 | } | |
28 | fn update_node(&mut self, idx: usize, clr: &[f64; 3], alpha: f64) { | |
29 | self.weights[idx][0] -= alpha * (self.weights[idx][0] - clr[0]); | |
30 | self.weights[idx][1] -= alpha * (self.weights[idx][1] - clr[1]); | |
31 | self.weights[idx][2] -= alpha * (self.weights[idx][2] - clr[2]); | |
32 | } | |
33 | fn update_neighbours(&mut self, idx: usize, clr: &[f64; 3], alpha: f64, radius: usize) { | |
34 | let low = idx.saturating_sub(radius).max(SPECIAL_NODES - 1); | |
35 | let high = (idx + radius).min(self.weights.len() - 1); | |
36 | ||
37 | let mut idx0 = idx + 1; | |
38 | let mut idx1 = idx - 1; | |
39 | let mut range = 0; | |
40 | let sqradius = (radius * radius) as f64; | |
41 | while (idx0 < high) || (idx1 > low) { | |
b36f412c | 42 | let sqrng = f64::from(range * range); |
4b459d0b KS |
43 | let a = alpha * (sqradius - sqrng) / sqradius; |
44 | range += 1; | |
45 | if idx0 < high { | |
46 | self.update_node(idx0, clr, a); | |
47 | idx0 += 1; | |
48 | } | |
49 | if idx1 > low { | |
50 | self.update_node(idx1, clr, a); | |
51 | idx1 -= 1; | |
52 | } | |
53 | } | |
54 | } | |
b7c882c1 | 55 | #[allow(clippy::float_cmp)] |
4b459d0b KS |
56 | fn find_node(&mut self, clr: &[f64; 3]) -> usize { |
57 | for i in 0..SPECIAL_NODES { | |
58 | if &self.weights[i] == clr { | |
59 | return i; | |
60 | } | |
61 | } | |
62 | let mut bestdist = std::f64::MAX; | |
63 | let mut distidx = 0; | |
64 | let mut bestbias = std::f64::MAX; | |
65 | let mut biasidx = 0; | |
66 | for i in SPECIAL_NODES..256 { | |
67 | let dist = (self.weights[i][0] - clr[0]) * (self.weights[i][0] - clr[0]) | |
68 | + (self.weights[i][1] - clr[1]) * (self.weights[i][1] - clr[1]) | |
69 | + (self.weights[i][2] - clr[2]) * (self.weights[i][2] - clr[2]); | |
70 | if bestdist > dist { | |
71 | bestdist = dist; | |
72 | distidx = i; | |
73 | } | |
74 | let biasdiff = dist - self.bias[i]; | |
75 | if bestbias > biasdiff { | |
76 | bestbias = biasdiff; | |
77 | biasidx = i; | |
78 | } | |
79 | self.freq[i] -= self.freq[i] / 1024.0; | |
80 | self.bias[i] += self.freq[i]; | |
81 | } | |
82 | self.freq[distidx] += 1.0 / 1024.0; | |
83 | self.bias[distidx] -= 1.0; | |
84 | biasidx | |
85 | } | |
86 | pub fn learn(&mut self, src: &[Pixel]) { | |
87 | let mut bias_radius = (256 / 8) << 6; | |
88 | let alphadec = (30 + (self.factor - 1) / 3) as f64; | |
b36f412c | 89 | let initial_alpha = f64::from(1 << 10); |
4b459d0b KS |
90 | |
91 | let npixels = src.len(); | |
92 | ||
93 | let mut radius = bias_radius >> 6; | |
94 | if radius == 1 { radius = 0 }; | |
95 | let samples = npixels / self.factor; | |
96 | let delta = samples / 100; | |
97 | let mut alpha = initial_alpha; | |
98 | ||
99 | let mut pos = 0; | |
100 | const PRIMES: [usize; 4] = [ 499, 491, 487, 503 ]; | |
101 | let mut step = PRIMES[3]; | |
102 | for prime in PRIMES.iter().rev() { | |
103 | if npixels % *prime != 0 { | |
104 | step = *prime; | |
105 | } | |
106 | } | |
107 | ||
108 | for i in 0..samples { | |
b36f412c | 109 | let clr = [f64::from(src[pos].r), f64::from(src[pos].g), f64::from(src[pos].b)]; |
4b459d0b KS |
110 | let idx = self.find_node(&clr); |
111 | if idx >= SPECIAL_NODES { | |
112 | let new_alpha = alphadec / initial_alpha; | |
113 | self.update_node(idx, &clr, new_alpha); | |
114 | if radius > 0 { | |
115 | self.update_neighbours(idx, &clr, new_alpha, radius); | |
116 | } | |
117 | } | |
118 | pos = (pos + step) % npixels; | |
119 | if (i + 1) % delta == 0 { | |
120 | alpha -= alpha / alphadec; | |
121 | bias_radius -= bias_radius / 30; | |
122 | radius = bias_radius >> 6; | |
123 | if radius == 1 { radius = 0 }; | |
124 | } | |
125 | } | |
126 | } | |
127 | pub fn make_pal(&self, pal: &mut [[u8; 3]; 256]) { | |
128 | for (pal, node) in pal.iter_mut().zip(self.weights.iter()) { | |
129 | pal[0] = (node[0] + 0.5).max(0.0).min(255.0) as u8; | |
130 | pal[1] = (node[1] + 0.5).max(0.0).min(255.0) as u8; | |
131 | pal[2] = (node[2] + 0.5).max(0.0).min(255.0) as u8; | |
132 | } | |
133 | } | |
134 | } |