]>
Commit | Line | Data |
---|---|---|
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) { | |
42 | let sqrng = (range * range) as f64; | |
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 | } | |
55 | fn find_node(&mut self, clr: &[f64; 3]) -> usize { | |
56 | for i in 0..SPECIAL_NODES { | |
57 | if &self.weights[i] == clr { | |
58 | return i; | |
59 | } | |
60 | } | |
61 | let mut bestdist = std::f64::MAX; | |
62 | let mut distidx = 0; | |
63 | let mut bestbias = std::f64::MAX; | |
64 | let mut biasidx = 0; | |
65 | for i in SPECIAL_NODES..256 { | |
66 | let dist = (self.weights[i][0] - clr[0]) * (self.weights[i][0] - clr[0]) | |
67 | + (self.weights[i][1] - clr[1]) * (self.weights[i][1] - clr[1]) | |
68 | + (self.weights[i][2] - clr[2]) * (self.weights[i][2] - clr[2]); | |
69 | if bestdist > dist { | |
70 | bestdist = dist; | |
71 | distidx = i; | |
72 | } | |
73 | let biasdiff = dist - self.bias[i]; | |
74 | if bestbias > biasdiff { | |
75 | bestbias = biasdiff; | |
76 | biasidx = i; | |
77 | } | |
78 | self.freq[i] -= self.freq[i] / 1024.0; | |
79 | self.bias[i] += self.freq[i]; | |
80 | } | |
81 | self.freq[distidx] += 1.0 / 1024.0; | |
82 | self.bias[distidx] -= 1.0; | |
83 | biasidx | |
84 | } | |
85 | pub fn learn(&mut self, src: &[Pixel]) { | |
86 | let mut bias_radius = (256 / 8) << 6; | |
87 | let alphadec = (30 + (self.factor - 1) / 3) as f64; | |
88 | let initial_alpha = (1 << 10) as f64; | |
89 | ||
90 | let npixels = src.len(); | |
91 | ||
92 | let mut radius = bias_radius >> 6; | |
93 | if radius == 1 { radius = 0 }; | |
94 | let samples = npixels / self.factor; | |
95 | let delta = samples / 100; | |
96 | let mut alpha = initial_alpha; | |
97 | ||
98 | let mut pos = 0; | |
99 | const PRIMES: [usize; 4] = [ 499, 491, 487, 503 ]; | |
100 | let mut step = PRIMES[3]; | |
101 | for prime in PRIMES.iter().rev() { | |
102 | if npixels % *prime != 0 { | |
103 | step = *prime; | |
104 | } | |
105 | } | |
106 | ||
107 | for i in 0..samples { | |
108 | let clr = [src[pos].r as f64, src[pos].g as f64, src[pos].b as f64]; | |
109 | let idx = self.find_node(&clr); | |
110 | if idx >= SPECIAL_NODES { | |
111 | let new_alpha = alphadec / initial_alpha; | |
112 | self.update_node(idx, &clr, new_alpha); | |
113 | if radius > 0 { | |
114 | self.update_neighbours(idx, &clr, new_alpha, radius); | |
115 | } | |
116 | } | |
117 | pos = (pos + step) % npixels; | |
118 | if (i + 1) % delta == 0 { | |
119 | alpha -= alpha / alphadec; | |
120 | bias_radius -= bias_radius / 30; | |
121 | radius = bias_radius >> 6; | |
122 | if radius == 1 { radius = 0 }; | |
123 | } | |
124 | } | |
125 | } | |
126 | pub fn make_pal(&self, pal: &mut [[u8; 3]; 256]) { | |
127 | for (pal, node) in pal.iter_mut().zip(self.weights.iter()) { | |
128 | pal[0] = (node[0] + 0.5).max(0.0).min(255.0) as u8; | |
129 | pal[1] = (node[1] + 0.5).max(0.0).min(255.0) as u8; | |
130 | pal[2] = (node[2] + 0.5).max(0.0).min(255.0) as u8; | |
131 | } | |
132 | } | |
133 | } |