deLyrium/src/palette.rs
2022-01-08 10:28:42 -05:00

95 lines
2.2 KiB
Rust

use image::Rgb;
use exoquant::ditherer::FloydSteinberg;
use exoquant::Remapper;
use exoquant::optimizer::WeightedKMeans;
use exoquant::generate_palette;
use exoquant::SimpleColorSpace;
use image::DynamicImage;
use image::GenericImageView;
use exoquant::Color;
const MAX_SIZE_PIXELS: u32 = 160_000;
const RESIZE_TARGET_PIXELS: f32 = 90_000.;
const NUM_COLORS: usize = 1;
pub struct Palette {
pub raw_colors: [Rgb<u8>; NUM_COLORS],
color_frequencies: [usize; NUM_COLORS],
}
impl Palette {
pub fn generate(img: &DynamicImage) -> Self {
let _thumb;
// Scale the image down if it's too big
let thumb = if img.width() * img.height() > MAX_SIZE_PIXELS {
let w = img.width() as f32;
let h = img.height() as f32;
let ratio = w / h;
let new_width = (ratio * RESIZE_TARGET_PIXELS).sqrt();
let new_height = RESIZE_TARGET_PIXELS / new_width;
_thumb = img.thumbnail(new_width as u32, new_height as u32);
&_thumb
} else {
img
};
// Convert to exoquant image
let width = thumb.width();
let exo_img: Vec<Color> = thumb.to_rgb8()
.pixels()
.map(|p| Color {r: p[0], g: p[1], b: p[2], a: 255})
.collect();
// Generate histogram
let histogram = exo_img.iter().cloned().collect();
// Generate raw palette
let colorspace = SimpleColorSpace::default();
let palette = generate_palette(
&histogram,
&colorspace,
&WeightedKMeans,
NUM_COLORS,
);
// Quantize the image & count the frequency of each color
let ditherer = FloydSteinberg::new();
let remapper = Remapper::new(&palette, &colorspace, &ditherer);
let remapped = remapper.remap_iter(
Box::new(exo_img.into_iter()),
width as usize,
);
let color_counts = remapped.fold([0; NUM_COLORS], |mut counts, pixel| {
counts[pixel as usize] += 1;
counts
});
// Convert to a more comfortable image::Rgba color format
let palette: Vec<Rgb<u8>> = palette.into_iter()
.map(|c| Rgb([c.r, c.g, c.b]))
.collect();
let palette: [Rgb<u8>; NUM_COLORS] = palette.try_into().unwrap();
Self {
raw_colors: palette,
color_frequencies: color_counts,
}
}
pub fn dominant_color(&self) -> &Rgb<u8> {
let max_index = self.color_frequencies
.iter()
.enumerate()
.max_by_key(|(_i, count)| *count)
.unwrap()
.0;
&self.raw_colors[max_index]
}
}