color_quant/lib.rs
1/*
2NeuQuant Neural-Net Quantization Algorithm by Anthony Dekker, 1994.
3See "Kohonen neural networks for optimal colour quantization"
4in "Network: Computation in Neural Systems" Vol. 5 (1994) pp 351-367.
5for a discussion of the algorithm.
6See also http://members.ozemail.com.au/~dekker/NEUQUANT.HTML
7
8Incorporated bugfixes and alpha channel handling from pngnq
9http://pngnq.sourceforge.net
10
11Copyright (c) 2014 The Piston Developers
12
13Permission is hereby granted, free of charge, to any person obtaining a copy
14of this software and associated documentation files (the "Software"), to deal
15in the Software without restriction, including without limitation the rights
16to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
17copies of the Software, and to permit persons to whom the Software is
18furnished to do so, subject to the following conditions:
19
20The above copyright notice and this permission notice shall be included in
21all copies or substantial portions of the Software.
22
23THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
24IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
25FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
26AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
27LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
28OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
29THE SOFTWARE.
30
31NeuQuant Neural-Net Quantization Algorithm
32------------------------------------------
33
34Copyright (c) 1994 Anthony Dekker
35
36NEUQUANT Neural-Net quantization algorithm by Anthony Dekker, 1994.
37See "Kohonen neural networks for optimal colour quantization"
38in "Network: Computation in Neural Systems" Vol. 5 (1994) pp 351-367.
39for a discussion of the algorithm.
40See also http://members.ozemail.com.au/~dekker/NEUQUANT.HTML
41
42Any party obtaining a copy of these files from the author, directly or
43indirectly, is granted, free of charge, a full and unrestricted irrevocable,
44world-wide, paid up, royalty-free, nonexclusive right and license to deal
45in this software and documentation files (the "Software"), including without
46limitation the rights to use, copy, modify, merge, publish, distribute, sublicense,
47and/or sell copies of the Software, and to permit persons who receive
48copies from any such party to do so, with the only requirement being
49that this copyright notice remain intact.
50
51*/
52
53//! # Color quantization library
54//!
55//! This library provides a color quantizer based on the [NEUQUANT](http://members.ozemail.com.au/~dekker/NEUQUANT.HTML)
56//!
57//! Original literature: Dekker, A. H. (1994). Kohonen neural networks for
58//! optimal colour quantization. *Network: Computation in Neural Systems*, 5(3), 351-367.
59//! [doi: 10.1088/0954-898X_5_3_003](https://doi.org/10.1088/0954-898X_5_3_003)
60//!
61//! See also <https://scientificgems.wordpress.com/stuff/neuquant-fast-high-quality-image-quantization/>
62//!
63//! ## Usage
64//!
65//! ```
66//! let data = vec![0; 40];
67//! let nq = color_quant::NeuQuant::new(10, 256, &data);
68//! let indixes: Vec<u8> = data.chunks(4).map(|pix| nq.index_of(pix) as u8).collect();
69//! let color_map = nq.color_map_rgba();
70//! ```
71
72mod math;
73use crate::math::clamp;
74
75use std::cmp::{max, min};
76
77const CHANNELS: usize = 4;
78
79const RADIUS_DEC: i32 = 30; // factor of 1/30 each cycle
80
81const ALPHA_BIASSHIFT: i32 = 10; // alpha starts at 1
82const INIT_ALPHA: i32 = 1 << ALPHA_BIASSHIFT; // biased by 10 bits
83
84const GAMMA: f64 = 1024.0;
85const BETA: f64 = 1.0 / GAMMA;
86const BETAGAMMA: f64 = BETA * GAMMA;
87
88// four primes near 500 - assume no image has a length so large
89// that it is divisible by all four primes
90const PRIMES: [usize; 4] = [499, 491, 478, 503];
91
92#[derive(Clone, Copy)]
93struct Quad<T> {
94 r: T,
95 g: T,
96 b: T,
97 a: T,
98}
99
100type Neuron = Quad<f64>;
101type Color = Quad<i32>;
102
103pub struct NeuQuant {
104 network: Vec<Neuron>,
105 colormap: Vec<Color>,
106 netindex: Vec<usize>,
107 bias: Vec<f64>, // bias and freq arrays for learning
108 freq: Vec<f64>,
109 samplefac: i32,
110 netsize: usize,
111}
112
113impl NeuQuant {
114 /// Creates a new neuronal network and trains it with the supplied data.
115 ///
116 /// Pixels are assumed to be in RGBA format.
117 /// `colors` should be $>=64$. `samplefac` determines the faction of
118 /// the sample that will be used to train the network. Its value must be in the
119 /// range $[1, 30]$. A value of $1$ thus produces the best result but is also
120 /// slowest. $10$ is a good compromise between speed and quality.
121 pub fn new(samplefac: i32, colors: usize, pixels: &[u8]) -> Self {
122 let netsize = colors;
123 let mut this = NeuQuant {
124 network: Vec::with_capacity(netsize),
125 colormap: Vec::with_capacity(netsize),
126 netindex: vec![0; 256],
127 bias: Vec::with_capacity(netsize),
128 freq: Vec::with_capacity(netsize),
129 samplefac: samplefac,
130 netsize: colors,
131 };
132 this.init(pixels);
133 this
134 }
135
136 /// Initializes the neuronal network and trains it with the supplied data.
137 ///
138 /// This method gets called by `Self::new`.
139 pub fn init(&mut self, pixels: &[u8]) {
140 self.network.clear();
141 self.colormap.clear();
142 self.bias.clear();
143 self.freq.clear();
144 let freq = (self.netsize as f64).recip();
145 for i in 0..self.netsize {
146 let tmp = (i as f64) * 256.0 / (self.netsize as f64);
147 // Sets alpha values at 0 for dark pixels.
148 let a = if i < 16 { i as f64 * 16.0 } else { 255.0 };
149 self.network.push(Neuron {
150 r: tmp,
151 g: tmp,
152 b: tmp,
153 a: a,
154 });
155 self.colormap.push(Color {
156 r: 0,
157 g: 0,
158 b: 0,
159 a: 255,
160 });
161 self.freq.push(freq);
162 self.bias.push(0.0);
163 }
164 self.learn(pixels);
165 self.build_colormap();
166 self.build_netindex();
167 }
168
169 /// Maps the rgba-pixel in-place to the best-matching color in the color map.
170 #[inline(always)]
171 pub fn map_pixel(&self, pixel: &mut [u8]) {
172 assert!(pixel.len() == 4);
173 let (r, g, b, a) = (pixel[0], pixel[1], pixel[2], pixel[3]);
174 let i = self.search_netindex(b, g, r, a);
175 pixel[0] = self.colormap[i].r as u8;
176 pixel[1] = self.colormap[i].g as u8;
177 pixel[2] = self.colormap[i].b as u8;
178 pixel[3] = self.colormap[i].a as u8;
179 }
180
181 /// Finds the best-matching index in the color map.
182 ///
183 /// `pixel` is assumed to be in RGBA format.
184 #[inline(always)]
185 pub fn index_of(&self, pixel: &[u8]) -> usize {
186 assert!(pixel.len() == 4);
187 let (r, g, b, a) = (pixel[0], pixel[1], pixel[2], pixel[3]);
188 self.search_netindex(b, g, r, a)
189 }
190
191 /// Lookup pixel values for color at `idx` in the colormap.
192 pub fn lookup(&self, idx: usize) -> Option<[u8; 4]> {
193 self.colormap
194 .get(idx)
195 .map(|p| [p.r as u8, p.g as u8, p.b as u8, p.a as u8])
196 }
197
198 /// Returns the RGBA color map calculated from the sample.
199 pub fn color_map_rgba(&self) -> Vec<u8> {
200 let mut map = Vec::with_capacity(self.netsize * 4);
201 for entry in &self.colormap {
202 map.push(entry.r as u8);
203 map.push(entry.g as u8);
204 map.push(entry.b as u8);
205 map.push(entry.a as u8);
206 }
207 map
208 }
209
210 /// Returns the RGBA color map calculated from the sample.
211 pub fn color_map_rgb(&self) -> Vec<u8> {
212 let mut map = Vec::with_capacity(self.netsize * 3);
213 for entry in &self.colormap {
214 map.push(entry.r as u8);
215 map.push(entry.g as u8);
216 map.push(entry.b as u8);
217 }
218 map
219 }
220
221 /// Move neuron i towards biased (a,b,g,r) by factor alpha
222 fn salter_single(&mut self, alpha: f64, i: i32, quad: Quad<f64>) {
223 let n = &mut self.network[i as usize];
224 n.b -= alpha * (n.b - quad.b);
225 n.g -= alpha * (n.g - quad.g);
226 n.r -= alpha * (n.r - quad.r);
227 n.a -= alpha * (n.a - quad.a);
228 }
229
230 /// Move neuron adjacent neurons towards biased (a,b,g,r) by factor alpha
231 fn alter_neighbour(&mut self, alpha: f64, rad: i32, i: i32, quad: Quad<f64>) {
232 let lo = max(i - rad, 0);
233 let hi = min(i + rad, self.netsize as i32);
234 let mut j = i + 1;
235 let mut k = i - 1;
236 let mut q = 0;
237
238 while (j < hi) || (k > lo) {
239 let rad_sq = rad as f64 * rad as f64;
240 let alpha = (alpha * (rad_sq - q as f64 * q as f64)) / rad_sq;
241 q += 1;
242 if j < hi {
243 let p = &mut self.network[j as usize];
244 p.b -= alpha * (p.b - quad.b);
245 p.g -= alpha * (p.g - quad.g);
246 p.r -= alpha * (p.r - quad.r);
247 p.a -= alpha * (p.a - quad.a);
248 j += 1;
249 }
250 if k > lo {
251 let p = &mut self.network[k as usize];
252 p.b -= alpha * (p.b - quad.b);
253 p.g -= alpha * (p.g - quad.g);
254 p.r -= alpha * (p.r - quad.r);
255 p.a -= alpha * (p.a - quad.a);
256 k -= 1;
257 }
258 }
259 }
260
261 /// Search for biased BGR values
262 /// finds closest neuron (min dist) and updates freq
263 /// finds best neuron (min dist-bias) and returns position
264 /// for frequently chosen neurons, freq[i] is high and bias[i] is negative
265 /// bias[i] = gamma*((1/self.netsize)-freq[i])
266 fn contest(&mut self, b: f64, g: f64, r: f64, a: f64) -> i32 {
267 use std::f64;
268
269 let mut bestd = f64::MAX;
270 let mut bestbiasd: f64 = bestd;
271 let mut bestpos = -1;
272 let mut bestbiaspos: i32 = bestpos;
273
274 for i in 0..self.netsize {
275 let bestbiasd_biased = bestbiasd + self.bias[i];
276 let mut dist;
277 let n = &self.network[i];
278 dist = (n.b - b).abs();
279 dist += (n.r - r).abs();
280 if dist < bestd || dist < bestbiasd_biased {
281 dist += (n.g - g).abs();
282 dist += (n.a - a).abs();
283 if dist < bestd {
284 bestd = dist;
285 bestpos = i as i32;
286 }
287 let biasdist = dist - self.bias[i];
288 if biasdist < bestbiasd {
289 bestbiasd = biasdist;
290 bestbiaspos = i as i32;
291 }
292 }
293 self.freq[i] -= BETA * self.freq[i];
294 self.bias[i] += BETAGAMMA * self.freq[i];
295 }
296 self.freq[bestpos as usize] += BETA;
297 self.bias[bestpos as usize] -= BETAGAMMA;
298 return bestbiaspos;
299 }
300
301 /// Main learning loop
302 /// Note: the number of learning cycles is crucial and the parameters are not
303 /// optimized for net sizes < 26 or > 256. 1064 colors seems to work fine
304 fn learn(&mut self, pixels: &[u8]) {
305 let initrad: i32 = self.netsize as i32 / 8; // for 256 cols, radius starts at 32
306 let radiusbiasshift: i32 = 6;
307 let radiusbias: i32 = 1 << radiusbiasshift;
308 let init_bias_radius: i32 = initrad * radiusbias;
309 let mut bias_radius = init_bias_radius;
310 let alphadec = 30 + ((self.samplefac - 1) / 3);
311 let lengthcount = pixels.len() / CHANNELS;
312 let samplepixels = lengthcount / self.samplefac as usize;
313 // learning cycles
314 let n_cycles = match self.netsize >> 1 {
315 n if n <= 100 => 100,
316 n => n,
317 };
318 let delta = match samplepixels / n_cycles {
319 0 => 1,
320 n => n,
321 };
322 let mut alpha = INIT_ALPHA;
323
324 let mut rad = bias_radius >> radiusbiasshift;
325 if rad <= 1 {
326 rad = 0
327 };
328
329 let mut pos = 0;
330 let step = *PRIMES
331 .iter()
332 .find(|&&prime| lengthcount % prime != 0)
333 .unwrap_or(&PRIMES[3]);
334
335 let mut i = 0;
336 while i < samplepixels {
337 let (r, g, b, a) = {
338 let p = &pixels[CHANNELS * pos..][..CHANNELS];
339 (p[0] as f64, p[1] as f64, p[2] as f64, p[3] as f64)
340 };
341
342 let j = self.contest(b, g, r, a);
343
344 let alpha_ = (1.0 * alpha as f64) / INIT_ALPHA as f64;
345 self.salter_single(alpha_, j, Quad { b, g, r, a });
346 if rad > 0 {
347 self.alter_neighbour(alpha_, rad, j, Quad { b, g, r, a })
348 };
349
350 pos += step;
351 while pos >= lengthcount {
352 pos -= lengthcount
353 }
354
355 i += 1;
356 if i % delta == 0 {
357 alpha -= alpha / alphadec;
358 bias_radius -= bias_radius / RADIUS_DEC;
359 rad = bias_radius >> radiusbiasshift;
360 if rad <= 1 {
361 rad = 0
362 };
363 }
364 }
365 }
366
367 /// initializes the color map
368 fn build_colormap(&mut self) {
369 for i in 0usize..self.netsize {
370 self.colormap[i].b = clamp(self.network[i].b.round() as i32);
371 self.colormap[i].g = clamp(self.network[i].g.round() as i32);
372 self.colormap[i].r = clamp(self.network[i].r.round() as i32);
373 self.colormap[i].a = clamp(self.network[i].a.round() as i32);
374 }
375 }
376
377 /// Insertion sort of network and building of netindex[0..255]
378 fn build_netindex(&mut self) {
379 let mut previouscol = 0;
380 let mut startpos = 0;
381
382 for i in 0..self.netsize {
383 let mut p = self.colormap[i];
384 let mut q;
385 let mut smallpos = i;
386 let mut smallval = p.g as usize; // index on g
387 // find smallest in i..netsize-1
388 for j in (i + 1)..self.netsize {
389 q = self.colormap[j];
390 if (q.g as usize) < smallval {
391 // index on g
392 smallpos = j;
393 smallval = q.g as usize; // index on g
394 }
395 }
396 q = self.colormap[smallpos];
397 // swap p (i) and q (smallpos) entries
398 if i != smallpos {
399 ::std::mem::swap(&mut p, &mut q);
400 self.colormap[i] = p;
401 self.colormap[smallpos] = q;
402 }
403 // smallval entry is now in position i
404 if smallval != previouscol {
405 self.netindex[previouscol] = (startpos + i) >> 1;
406 for j in (previouscol + 1)..smallval {
407 self.netindex[j] = i
408 }
409 previouscol = smallval;
410 startpos = i;
411 }
412 }
413 let max_netpos = self.netsize - 1;
414 self.netindex[previouscol] = (startpos + max_netpos) >> 1;
415 for j in (previouscol + 1)..256 {
416 self.netindex[j] = max_netpos
417 } // really 256
418 }
419
420 /// Search for best matching color
421 fn search_netindex(&self, b: u8, g: u8, r: u8, a: u8) -> usize {
422 let mut bestd = 1 << 30; // ~ 1_000_000
423 let mut best = 0;
424 // start at netindex[g] and work outwards
425 let mut i = self.netindex[g as usize];
426 let mut j = if i > 0 { i - 1 } else { 0 };
427
428 while (i < self.netsize) || (j > 0) {
429 if i < self.netsize {
430 let p = self.colormap[i];
431 let mut e = p.g - g as i32;
432 let mut dist = e * e; // inx key
433 if dist >= bestd {
434 break;
435 } else {
436 e = p.b - b as i32;
437 dist += e * e;
438 if dist < bestd {
439 e = p.r - r as i32;
440 dist += e * e;
441 if dist < bestd {
442 e = p.a - a as i32;
443 dist += e * e;
444 if dist < bestd {
445 bestd = dist;
446 best = i;
447 }
448 }
449 }
450 i += 1;
451 }
452 }
453 if j > 0 {
454 let p = self.colormap[j];
455 let mut e = p.g - g as i32;
456 let mut dist = e * e; // inx key
457 if dist >= bestd {
458 break;
459 } else {
460 e = p.b - b as i32;
461 dist += e * e;
462 if dist < bestd {
463 e = p.r - r as i32;
464 dist += e * e;
465 if dist < bestd {
466 e = p.a - a as i32;
467 dist += e * e;
468 if dist < bestd {
469 bestd = dist;
470 best = j;
471 }
472 }
473 }
474 j -= 1;
475 }
476 }
477 }
478 best
479 }
480}