pbnify/
quantize.rs

1// Copyright (C) 2019  Adam Gausmann
2//
3// This program is free software: you can redistribute it and/or modify
4// it under the terms of the GNU General Public License as published by
5// the Free Software Foundation, either version 3 of the License, or
6// (at your option) any later version.
7//
8// This program is distributed in the hope that it will be useful,
9// but WITHOUT ANY WARRANTY; without even the implied warranty of
10// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
11// GNU General Public License for more details.
12//
13// You should have received a copy of the GNU General Public License
14// along with this program.  If not, see <https://www.gnu.org/licenses/>.
15
16//! Extensions for quantization and dithering of images to reduce the number of unique colors.
17
18use image::{ImageBuffer, Rgba, RgbaImage};
19
20fn sort_widest(pixels: &mut [Rgba<u8>]) {
21    let r_range =
22        pixels.iter().map(|p| p[0]).max().unwrap() - pixels.iter().map(|p| p[0]).min().unwrap();
23    let g_range =
24        pixels.iter().map(|p| p[1]).max().unwrap() - pixels.iter().map(|p| p[1]).min().unwrap();
25    let b_range =
26        pixels.iter().map(|p| p[2]).max().unwrap() - pixels.iter().map(|p| p[2]).min().unwrap();
27    let a_range =
28        pixels.iter().map(|p| p[3]).max().unwrap() - pixels.iter().map(|p| p[3]).min().unwrap();
29
30    if r_range > g_range && r_range > b_range && r_range > a_range {
31        pixels.sort_by_key(|p| p[0]);
32    } else if g_range > b_range && g_range > a_range {
33        pixels.sort_by_key(|p| p[1]);
34    } else if b_range > a_range {
35        pixels.sort_by_key(|p| p[2]);
36    } else {
37        pixels.sort_by_key(|p| p[3]);
38    }
39}
40
41/// Performs the median-cut algorithm, producing an `n`-color palette based on the given set of
42/// pixels.
43///
44/// If `n` is not a power of two, it will be rounded up to the smallest power of two greater than
45/// `n`.
46///
47/// If `n` is greater than or equal to the number of given pixels, then the input will be returned
48/// as the result.
49pub fn build_palette<I>(n: usize, pixels: I) -> Vec<Rgba<u8>>
50where
51    I: IntoIterator<Item = Rgba<u8>>,
52{
53    let n = n.next_power_of_two();
54    let mut pixels: Vec<_> = pixels.into_iter().collect();
55    pixels.sort_by_key(|p| p.data);
56    pixels.dedup();
57
58    if n >= pixels.len() {
59        return pixels;
60    }
61
62    for i in 0..n.trailing_zeros() {
63        let i = 1 << i;
64        for j in 0..i {
65            let pixels_a = j * pixels.len() / i;
66            let pixels_b = (j + 1) * pixels.len() / i;
67            sort_widest(&mut pixels[pixels_a..pixels_b]);
68        }
69    }
70
71    let mut output = Vec::with_capacity(n);
72    for i in 0..n {
73        let pixels_a = i * pixels.len() / n;
74        let pixels_b = (i + 1) * pixels.len() / n;
75        let sub_pixels = &pixels[pixels_a..pixels_b];
76
77        let avg_r: usize =
78            sub_pixels.iter().map(|p| p[0] as usize).sum::<usize>() / sub_pixels.len();
79        let avg_g: usize =
80            sub_pixels.iter().map(|p| p[1] as usize).sum::<usize>() / sub_pixels.len();
81        let avg_b: usize =
82            sub_pixels.iter().map(|p| p[2] as usize).sum::<usize>() / sub_pixels.len();
83        let avg_a: usize =
84            sub_pixels.iter().map(|p| p[3] as usize).sum::<usize>() / sub_pixels.len();
85
86        output.push(Rgba {
87            data: [avg_r as u8, avg_g as u8, avg_b as u8, avg_a as u8],
88        });
89    }
90    output
91}
92
93/// Substitutes each pixel in the given image with its closest approximation in the given palette.
94/// This method also employs Floyd-Steinberg dithering to smooth out gradients and more accurately
95/// represent colors which do not have close approximations in the given palette.
96pub fn quantize(image: &mut RgbaImage, palette: &[Rgba<u8>]) {
97    let mut fimage: ImageBuffer<Rgba<f32>, Vec<f32>> = ImageBuffer::from_raw(
98        image.width(),
99        image.height(),
100        vec![0.0; 4 * image.width() as usize * image.height() as usize],
101    )
102    .unwrap();
103
104    for x in 0..image.width() {
105        for y in 0..image.height() {
106            let [r, g, b, a] = image.get_pixel(x, y).data;
107            fimage.put_pixel(
108                x,
109                y,
110                Rgba {
111                    data: [
112                        r as f32 / 255.0,
113                        g as f32 / 255.0,
114                        b as f32 / 255.0,
115                        a as f32 / 255.0,
116                    ],
117                },
118            );
119        }
120    }
121
122    let palette: Vec<Rgba<f32>> = palette
123        .iter()
124        .map(|p| {
125            let [r, g, b, a] = p.data;
126            Rgba {
127                data: [
128                    r as f32 / 255.0,
129                    g as f32 / 255.0,
130                    b as f32 / 255.0,
131                    a as f32 / 255.0,
132                ],
133            }
134        })
135        .collect();
136
137    for x in 0..fimage.width() {
138        for y in 0..image.height() {
139            let old_pixel = fimage.get_pixel(x, y);
140            let [or, og, ob, oa] = old_pixel.data;
141
142            let new_pixel = palette
143                .iter()
144                .min_by_key(|p| {
145                    let [nr, ng, nb, na] = p.data;
146                    (((nr - or).powi(2)
147                        + (ng - og).powi(2)
148                        + (nb - ob).powi(2)
149                        + (na - oa).powi(2))
150                        * 1000000.0) as usize
151                })
152                .unwrap();
153            let [nr, ng, nb, na] = new_pixel.data;
154            let er = or - nr;
155            let eg = og - ng;
156            let eb = ob - nb;
157            let ea = oa - na;
158
159            if x < fimage.width() - 1 {
160                let [hr, hg, hb, ha] = &mut fimage.get_pixel_mut(x + 1, y).data;
161                *hr += er * 0.4375;
162                *hg += eg * 0.4375;
163                *hb += eb * 0.4375;
164                *ha += ea * 0.4375;
165            }
166            if y < fimage.height() - 1 {
167                let [hr, hg, hb, ha] = &mut fimage.get_pixel_mut(x, y + 1).data;
168                *hr += er * 0.1875;
169                *hg += eg * 0.1875;
170                *hb += eb * 0.1875;
171                *ha += ea * 0.1875;
172
173                if x < fimage.width() - 1 {
174                    let [hr, hg, hb, ha] = &mut fimage.get_pixel_mut(x + 1, y + 1).data;
175                    *hr += er * 0.3125;
176                    *hg += eg * 0.3125;
177                    *hb += eb * 0.3125;
178                    *ha += ea * 0.3125;
179                }
180
181                if x > 0 {
182                    let [hr, hg, hb, ha] = &mut fimage.get_pixel_mut(x - 1, y + 1).data;
183                    *hr += er * 0.0625;
184                    *hg += eg * 0.0625;
185                    *hb += eb * 0.0625;
186                    *ha += ea * 0.0625;
187                }
188            }
189
190            fimage.put_pixel(x, y, *new_pixel);
191            image.put_pixel(
192                x,
193                y,
194                Rgba {
195                    data: [
196                        (nr * 255.0) as u8,
197                        (ng * 255.0) as u8,
198                        (nb * 255.0) as u8,
199                        (na * 255.0) as u8,
200                    ],
201                },
202            );
203        }
204    }
205}
pbnify/quantize.rs

pbnify/
quantize.rs
