pbnify 0.2.1

// Copyright (C) 2019  Adam Gausmann
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.

//! Extensions for quantization and dithering of images to reduce the number of unique colors.

use image::{ImageBuffer, Rgba, RgbaImage};

fn sort_widest(pixels: &mut [Rgba<u8>]) {
    let r_range =
        pixels.iter().map(|p| p[0]).max().unwrap() - pixels.iter().map(|p| p[0]).min().unwrap();
    let g_range =
        pixels.iter().map(|p| p[1]).max().unwrap() - pixels.iter().map(|p| p[1]).min().unwrap();
    let b_range =
        pixels.iter().map(|p| p[2]).max().unwrap() - pixels.iter().map(|p| p[2]).min().unwrap();
    let a_range =
        pixels.iter().map(|p| p[3]).max().unwrap() - pixels.iter().map(|p| p[3]).min().unwrap();

    if r_range > g_range && r_range > b_range && r_range > a_range {
        pixels.sort_by_key(|p| p[0]);
    } else if g_range > b_range && g_range > a_range {
        pixels.sort_by_key(|p| p[1]);
    } else if b_range > a_range {
        pixels.sort_by_key(|p| p[2]);
    } else {
        pixels.sort_by_key(|p| p[3]);
    }
}

/// Performs the median-cut algorithm, producing an `n`-color palette based on the given set of
/// pixels.
///
/// If `n` is not a power of two, it will be rounded up to the smallest power of two greater than
/// `n`.
///
/// If `n` is greater than or equal to the number of given pixels, then the input will be returned
/// as the result.
pub fn build_palette<I>(n: usize, pixels: I) -> Vec<Rgba<u8>>
where
    I: IntoIterator<Item = Rgba<u8>>,
{
    let n = n.next_power_of_two();
    let mut pixels: Vec<_> = pixels.into_iter().collect();
    pixels.sort_by_key(|p| p.data);
    pixels.dedup();

    if n >= pixels.len() {
        return pixels;
    }

    for i in 0..n.trailing_zeros() {
        let i = 1 << i;
        for j in 0..i {
            let pixels_a = j * pixels.len() / i;
            let pixels_b = (j + 1) * pixels.len() / i;
            sort_widest(&mut pixels[pixels_a..pixels_b]);
        }
    }

    let mut output = Vec::with_capacity(n);
    for i in 0..n {
        let pixels_a = i * pixels.len() / n;
        let pixels_b = (i + 1) * pixels.len() / n;
        let sub_pixels = &pixels[pixels_a..pixels_b];

        let avg_r: usize =
            sub_pixels.iter().map(|p| p[0] as usize).sum::<usize>() / sub_pixels.len();
        let avg_g: usize =
            sub_pixels.iter().map(|p| p[1] as usize).sum::<usize>() / sub_pixels.len();
        let avg_b: usize =
            sub_pixels.iter().map(|p| p[2] as usize).sum::<usize>() / sub_pixels.len();
        let avg_a: usize =
            sub_pixels.iter().map(|p| p[3] as usize).sum::<usize>() / sub_pixels.len();

        output.push(Rgba {
            data: [avg_r as u8, avg_g as u8, avg_b as u8, avg_a as u8],
        });
    }
    output
}

/// Substitutes each pixel in the given image with its closest approximation in the given palette.
/// This method also employs Floyd-Steinberg dithering to smooth out gradients and more accurately
/// represent colors which do not have close approximations in the given palette.
pub fn quantize(image: &mut RgbaImage, palette: &[Rgba<u8>]) {
    let mut fimage: ImageBuffer<Rgba<f32>, Vec<f32>> = ImageBuffer::from_raw(
        image.width(),
        image.height(),
        vec![0.0; 4 * image.width() as usize * image.height() as usize],
    )
    .unwrap();

    for x in 0..image.width() {
        for y in 0..image.height() {
            let [r, g, b, a] = image.get_pixel(x, y).data;
            fimage.put_pixel(
                x,
                y,
                Rgba {
                    data: [
                        r as f32 / 255.0,
                        g as f32 / 255.0,
                        b as f32 / 255.0,
                        a as f32 / 255.0,
                    ],
                },
            );
        }
    }

    let palette: Vec<Rgba<f32>> = palette
        .iter()
        .map(|p| {
            let [r, g, b, a] = p.data;
            Rgba {
                data: [
                    r as f32 / 255.0,
                    g as f32 / 255.0,
                    b as f32 / 255.0,
                    a as f32 / 255.0,
                ],
            }
        })
        .collect();

    for x in 0..fimage.width() {
        for y in 0..image.height() {
            let old_pixel = fimage.get_pixel(x, y);
            let [or, og, ob, oa] = old_pixel.data;

            let new_pixel = palette
                .iter()
                .min_by_key(|p| {
                    let [nr, ng, nb, na] = p.data;
                    (((nr - or).powi(2)
                        + (ng - og).powi(2)
                        + (nb - ob).powi(2)
                        + (na - oa).powi(2))
                        * 1000000.0) as usize
                })
                .unwrap();
            let [nr, ng, nb, na] = new_pixel.data;
            let er = or - nr;
            let eg = og - ng;
            let eb = ob - nb;
            let ea = oa - na;

            if x < fimage.width() - 1 {
                let [hr, hg, hb, ha] = &mut fimage.get_pixel_mut(x + 1, y).data;
                *hr += er * 0.4375;
                *hg += eg * 0.4375;
                *hb += eb * 0.4375;
                *ha += ea * 0.4375;
            }
            if y < fimage.height() - 1 {
                let [hr, hg, hb, ha] = &mut fimage.get_pixel_mut(x, y + 1).data;
                *hr += er * 0.1875;
                *hg += eg * 0.1875;
                *hb += eb * 0.1875;
                *ha += ea * 0.1875;

                if x < fimage.width() - 1 {
                    let [hr, hg, hb, ha] = &mut fimage.get_pixel_mut(x + 1, y + 1).data;
                    *hr += er * 0.3125;
                    *hg += eg * 0.3125;
                    *hb += eb * 0.3125;
                    *ha += ea * 0.3125;
                }

                if x > 0 {
                    let [hr, hg, hb, ha] = &mut fimage.get_pixel_mut(x - 1, y + 1).data;
                    *hr += er * 0.0625;
                    *hg += eg * 0.0625;
                    *hb += eb * 0.0625;
                    *ha += ea * 0.0625;
                }
            }

            fimage.put_pixel(x, y, *new_pixel);
            image.put_pixel(
                x,
                y,
                Rgba {
                    data: [
                        (nr * 255.0) as u8,
                        (ng * 255.0) as u8,
                        (nb * 255.0) as u8,
                        (na * 255.0) as u8,
                    ],
                },
            );
        }
    }
}