92 lines
2.2 KiB
Rust
92 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 {
|
|
|
|
// 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;
|
|
|
|
img.thumbnail(new_width as u32, new_height as u32)
|
|
} else {
|
|
img
|
|
};
|
|
|
|
// Convert to exoquant image
|
|
let width = thumb.width();
|
|
let exo_img: Vec<Color> = thumb.into_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;
|
|
|
|
return &self.raw_colors[max_index];
|
|
}
|
|
}
|