#![warn(missing_docs)]
#![warn(clippy::all)]
pub mod algorithms;
mod color;
mod error;
pub mod wasm;
pub use color::{Color, ColorPalette};
pub use error::{DominantColorError, Result};
use image::DynamicImage;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum Algorithm {
#[default]
KMeans,
MedianCut,
Octree,
}
#[derive(Debug, Clone)]
pub struct Config {
pub max_colors: usize,
pub sample_size: Option<u32>,
pub kmeans_seed: u64,
pub kmeans_max_iterations: usize,
pub kmeans_convergence_threshold: f64,
}
impl Default for Config {
fn default() -> Self {
Self {
max_colors: 8,
sample_size: Some(256),
kmeans_seed: 42,
kmeans_max_iterations: 100,
kmeans_convergence_threshold: 1.0,
}
}
}
impl Config {
#[must_use]
pub fn max_colors(mut self, n: usize) -> Self {
assert!(n > 0, "max_colors must be at least 1");
self.max_colors = n;
self
}
#[must_use]
pub fn sample_size(mut self, size: Option<u32>) -> Self {
self.sample_size = size;
self
}
#[must_use]
pub fn kmeans_seed(mut self, seed: u64) -> Self {
self.kmeans_seed = seed;
self
}
#[must_use]
pub fn kmeans_max_iterations(mut self, iters: usize) -> Self {
self.kmeans_max_iterations = iters;
self
}
}
pub struct DominantColors {
image: DynamicImage,
config: Config,
}
impl DominantColors {
pub fn new(image: DynamicImage) -> Self {
Self {
image,
config: Config::default(),
}
}
#[must_use]
pub fn config(mut self, config: Config) -> Self {
self.config = config;
self
}
pub fn extract(self, algorithm: Algorithm) -> Result<ColorPalette> {
let pixels = self.sample_pixels();
if pixels.is_empty() {
return Err(DominantColorError::EmptyImage);
}
let mut palette = match algorithm {
Algorithm::KMeans => algorithms::kmeans::extract(&pixels, &self.config)?,
Algorithm::MedianCut => algorithms::median_cut::extract(&pixels, &self.config)?,
Algorithm::Octree => algorithms::octree::extract(&pixels, &self.config)?,
};
palette.sort_by(|a, b| b.percentage.partial_cmp(&a.percentage).unwrap());
Ok(palette)
}
fn sample_pixels(&self) -> Vec<[u8; 3]> {
let img = if let Some(size) = self.config.sample_size {
let (w, h) = (self.image.width(), self.image.height());
if w > size || h > size {
self.image.thumbnail(size, size).into_rgb8()
} else {
self.image.to_rgb8()
}
} else {
self.image.to_rgb8()
};
img.pixels().map(|p| [p.0[0], p.0[1], p.0[2]]).collect()
}
}
#[cfg(test)]
mod tests {
use super::*;
use image::{ImageBuffer, Rgb};
fn make_image(pixels: &[[u8; 3]]) -> DynamicImage {
let width = pixels.len() as u32;
let buf: ImageBuffer<Rgb<u8>, Vec<u8>> = ImageBuffer::from_fn(width, 1, |x, _| {
let p = pixels[x as usize];
Rgb([p[0], p[1], p[2]])
});
DynamicImage::ImageRgb8(buf)
}
#[test]
fn test_empty_image_error() {
let img = DynamicImage::ImageRgb8(ImageBuffer::new(0, 0));
for alg in [Algorithm::KMeans, Algorithm::MedianCut, Algorithm::Octree] {
let result = DominantColors::new(img.clone())
.config(Config::default().sample_size(None))
.extract(alg);
assert!(
matches!(result, Err(DominantColorError::EmptyImage)),
"{alg:?} 应返回 EmptyImage"
);
}
}
#[test]
fn test_single_color_image() {
let pixels = vec![[255u8, 0, 0]; 100];
let img = make_image(&pixels);
for alg in [Algorithm::KMeans, Algorithm::MedianCut, Algorithm::Octree] {
let palette = DominantColors::new(img.clone())
.config(Config::default().max_colors(3).sample_size(None))
.extract(alg)
.expect("纯色图片应成功");
assert!(!palette.is_empty(), "{alg:?} 调色板不应为空");
assert!(
palette.iter().map(|c| c.percentage).sum::<f32>() > 0.99,
"{alg:?} 占比之和应约等于 1.0"
);
}
}
#[test]
fn test_two_color_image_separation() {
let mut pixels = vec![[255u8, 0, 0]; 50];
pixels.extend(vec![[0u8, 0, 255]; 50]);
let img = make_image(&pixels);
for alg in [Algorithm::KMeans, Algorithm::MedianCut, Algorithm::Octree] {
let palette = DominantColors::new(img.clone())
.config(Config::default().max_colors(2).sample_size(None))
.extract(alg)
.expect("双色图片应成功");
assert_eq!(palette.len(), 2, "{alg:?} 应识别出 2 种颜色");
for color in &palette {
assert!(
(color.percentage - 0.5).abs() < 0.1,
"{alg:?}: 期望约 50%,实际 {:.1}%",
color.percentage * 100.0
);
}
}
}
#[test]
fn test_palette_sorted_descending() {
let mut pixels = vec![[255u8, 0, 0]; 60];
pixels.extend(vec![[0u8, 255, 0]; 30]);
pixels.extend(vec![[0u8, 0, 255]; 10]);
let img = make_image(&pixels);
for alg in [Algorithm::KMeans, Algorithm::MedianCut, Algorithm::Octree] {
let palette = DominantColors::new(img.clone())
.config(Config::default().max_colors(3).sample_size(None))
.extract(alg)
.expect("三色图片应成功");
let percentages: Vec<f32> = palette.iter().map(|c| c.percentage).collect();
let mut sorted = percentages.clone();
sorted.sort_by(|a, b| b.partial_cmp(a).unwrap());
assert_eq!(percentages, sorted, "{alg:?} 调色板应按占比降序排列");
}
}
#[test]
fn test_config_builder() {
let cfg = Config::default()
.max_colors(10)
.sample_size(Some(128))
.kmeans_seed(99)
.kmeans_max_iterations(50);
assert_eq!(cfg.max_colors, 10);
assert_eq!(cfg.sample_size, Some(128));
assert_eq!(cfg.kmeans_seed, 99);
assert_eq!(cfg.kmeans_max_iterations, 50);
}
#[test]
#[should_panic(expected = "max_colors must be at least 1")]
fn test_config_zero_colors_panics() {
let _ = Config::default().max_colors(0);
}
}