dominant_colors/
lib.rs

1//! # dominant-colors
2//!
3//! 使用三种经典算法从图片中提取主色调：
4//!
5//! - **K-Means 聚类** — 基于质心的迭代式颜色量化
6//! - **中位切分（Median Cut）** — 沿最长轴递归划分颜色空间
7//! - **八叉树量化（Octree）** — 对 RGB 立方体进行层次化空间细分
8//!
9//! ## 快速上手
10//!
11//! ```rust,no_run
12//! use dominant_colors::{DominantColors, Algorithm, Config};
13//!
14//! // 加载图片，用 K-Means 提取 5 种主色调
15//! let img = image::open("photo.jpg").unwrap();
16//! let palette = DominantColors::new(img)
17//!     .config(Config::default().max_colors(5))
18//!     .extract(Algorithm::KMeans)
19//!     .unwrap();
20//!
21//! for color in &palette {
22//!     println!("#{:02X}{:02X}{:02X}  ({:.1}%)", color.r, color.g, color.b, color.percentage * 100.0);
23//! }
24//! ```
25//!
26//! ## 算法对比
27//!
28//! | 算法         | 速度 | 质量 | 确定性          |
29//! |-------------|------|------|----------------|
30//! | K-Means     | 中等 | 高   | 否（可设种子）  |
31//! | Median Cut  | 快   | 良   | 是              |
32//! | Octree      | 快   | 良   | 是              |
33
34#![warn(missing_docs)]
35#![warn(clippy::all)]
36
37pub mod algorithms;
38mod color;
39mod error;
40
41// WASM 绑定层：仅在编译目标为 wasm32 时启用
42pub mod wasm;
43
44pub use color::{Color, ColorPalette};
45pub use error::{DominantColorError, Result};
46
47use image::DynamicImage;
48
49/// 主色调提取算法枚举。
50#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
51pub enum Algorithm {
52    /// K-Means 聚类：迭代优化颜色质心。
53    /// 结果质量高，但因随机初始化而非确定性（可通过固定种子复现）。
54    #[default]
55    KMeans,
56    /// 中位切分：沿颜色范围最宽的通道递归对半分割颜色空间。
57    /// 速度快，结果完全确定。
58    MedianCut,
59    /// 八叉树量化：对 RGB 立方体进行层次化细分。
60    /// 速度快、结果确定、内存占用低。
61    Octree,
62}
63
64/// 所有算法共用的配置项。
65#[derive(Debug, Clone)]
66pub struct Config {
67    /// 提取的最大主色数量（默认：8）。
68    pub max_colors: usize,
69    /// 处理前将图片缩放到此尺寸（最长边像素数，默认：256）。
70    /// 设为 `None` 则处理原始分辨率。
71    pub sample_size: Option<u32>,
72    /// K-Means 随机种子（其他算法忽略此字段，默认：42）。
73    pub kmeans_seed: u64,
74    /// K-Means 最大迭代次数（默认：100）。
75    pub kmeans_max_iterations: usize,
76    /// K-Means 收敛阈值，单位为 RGB 欧氏距离（默认：1.0）。
77    /// 所有质心移动量均低于此值时提前终止迭代。
78    pub kmeans_convergence_threshold: f64,
79}
80
81impl Default for Config {
82    fn default() -> Self {
83        Self {
84            max_colors: 8,
85            sample_size: Some(256),
86            kmeans_seed: 42,
87            kmeans_max_iterations: 100,
88            kmeans_convergence_threshold: 1.0,
89        }
90    }
91}
92
93impl Config {
94    /// 设置提取的最大颜色数量。
95    ///
96    /// # Panics
97    ///
98    /// `n` 为零时 panic。
99    #[must_use]
100    pub fn max_colors(mut self, n: usize) -> Self {
101        assert!(n > 0, "max_colors must be at least 1");
102        self.max_colors = n;
103        self
104    }
105
106    /// 设置缩放尺寸（最长边像素数）。传入 `None` 则使用原始尺寸。
107    #[must_use]
108    pub fn sample_size(mut self, size: Option<u32>) -> Self {
109        self.sample_size = size;
110        self
111    }
112
113    /// 设置 K-Means 随机种子。
114    #[must_use]
115    pub fn kmeans_seed(mut self, seed: u64) -> Self {
116        self.kmeans_seed = seed;
117        self
118    }
119
120    /// 设置 K-Means 最大迭代次数。
121    #[must_use]
122    pub fn kmeans_max_iterations(mut self, iters: usize) -> Self {
123        self.kmeans_max_iterations = iters;
124        self
125    }
126}
127
128/// 主色调提取的构建器（Builder）。
129///
130/// # 示例
131///
132/// ```rust,no_run
133/// use dominant_colors::{DominantColors, Algorithm, Config};
134///
135/// let img = image::open("photo.jpg").unwrap();
136/// let palette = DominantColors::new(img)
137///     .config(Config::default().max_colors(6))
138///     .extract(Algorithm::Octree)
139///     .unwrap();
140/// ```
141pub struct DominantColors {
142    image: DynamicImage,
143    config: Config,
144}
145
146impl DominantColors {
147    /// 从 [`DynamicImage`] 创建提取器。
148    pub fn new(image: DynamicImage) -> Self {
149        Self {
150            image,
151            config: Config::default(),
152        }
153    }
154
155    /// 覆盖默认配置。
156    #[must_use]
157    pub fn config(mut self, config: Config) -> Self {
158        self.config = config;
159        self
160    }
161
162    /// 运行指定算法，返回按占比降序排列的 [`ColorPalette`]。
163    ///
164    /// 最主要的颜色排在最前面。
165    ///
166    /// # 错误
167    ///
168    /// 图片无像素或违反算法约束时返回 [`DominantColorError`]。
169    pub fn extract(self, algorithm: Algorithm) -> Result<ColorPalette> {
170        let pixels = self.sample_pixels();
171
172        if pixels.is_empty() {
173            return Err(DominantColorError::EmptyImage);
174        }
175
176        let mut palette = match algorithm {
177            Algorithm::KMeans => algorithms::kmeans::extract(&pixels, &self.config)?,
178            Algorithm::MedianCut => algorithms::median_cut::extract(&pixels, &self.config)?,
179            Algorithm::Octree => algorithms::octree::extract(&pixels, &self.config)?,
180        };
181
182        // 按占比降序排列，最主要的颜色排在最前
183        palette.sort_by(|a, b| b.percentage.partial_cmp(&a.percentage).unwrap());
184        Ok(palette)
185    }
186
187    /// 将图片等比缩放到 `config.sample_size` 后收集所有 RGB 像素。
188    fn sample_pixels(&self) -> Vec<[u8; 3]> {
189        let img = if let Some(size) = self.config.sample_size {
190            let (w, h) = (self.image.width(), self.image.height());
191            if w > size || h > size {
192                // 等比缩放，最长边不超过 size
193                self.image.thumbnail(size, size).into_rgb8()
194            } else {
195                self.image.to_rgb8()
196            }
197        } else {
198            self.image.to_rgb8()
199        };
200
201        img.pixels().map(|p| [p.0[0], p.0[1], p.0[2]]).collect()
202    }
203}
204
205#[cfg(test)]
206mod tests {
207    use super::*;
208    use image::{ImageBuffer, Rgb};
209
210    /// 从 RGB 三元组列表构造合成测试图片（宽度 = 像素数，高度 = 1）。
211    fn make_image(pixels: &[[u8; 3]]) -> DynamicImage {
212        let width = pixels.len() as u32;
213        let buf: ImageBuffer<Rgb<u8>, Vec<u8>> = ImageBuffer::from_fn(width, 1, |x, _| {
214            let p = pixels[x as usize];
215            Rgb([p[0], p[1], p[2]])
216        });
217        DynamicImage::ImageRgb8(buf)
218    }
219
220    #[test]
221    fn test_empty_image_error() {
222        // 空图片（0×0）应对所有算法均返回 EmptyImage 错误
223        let img = DynamicImage::ImageRgb8(ImageBuffer::new(0, 0));
224        for alg in [Algorithm::KMeans, Algorithm::MedianCut, Algorithm::Octree] {
225            let result = DominantColors::new(img.clone())
226                .config(Config::default().sample_size(None))
227                .extract(alg);
228            assert!(
229                matches!(result, Err(DominantColorError::EmptyImage)),
230                "{alg:?} 应返回 EmptyImage"
231            );
232        }
233    }
234
235    #[test]
236    fn test_single_color_image() {
237        // 纯红色图片：调色板不为空，且各颜色占比之和 ≈ 1.0
238        let pixels = vec![[255u8, 0, 0]; 100];
239        let img = make_image(&pixels);
240        for alg in [Algorithm::KMeans, Algorithm::MedianCut, Algorithm::Octree] {
241            let palette = DominantColors::new(img.clone())
242                .config(Config::default().max_colors(3).sample_size(None))
243                .extract(alg)
244                .expect("纯色图片应成功");
245            assert!(!palette.is_empty(), "{alg:?} 调色板不应为空");
246            assert!(
247                palette.iter().map(|c| c.percentage).sum::<f32>() > 0.99,
248                "{alg:?} 占比之和应约等于 1.0"
249            );
250        }
251    }
252
253    #[test]
254    fn test_two_color_image_separation() {
255        // 50 个红色像素 + 50 个蓝色像素，两种颜色各占约 50%
256        let mut pixels = vec![[255u8, 0, 0]; 50];
257        pixels.extend(vec![[0u8, 0, 255]; 50]);
258        let img = make_image(&pixels);
259
260        for alg in [Algorithm::KMeans, Algorithm::MedianCut, Algorithm::Octree] {
261            let palette = DominantColors::new(img.clone())
262                .config(Config::default().max_colors(2).sample_size(None))
263                .extract(alg)
264                .expect("双色图片应成功");
265
266            assert_eq!(palette.len(), 2, "{alg:?} 应识别出 2 种颜色");
267            for color in &palette {
268                assert!(
269                    (color.percentage - 0.5).abs() < 0.1,
270                    "{alg:?}: 期望约 50%，实际 {:.1}%",
271                    color.percentage * 100.0
272                );
273            }
274        }
275    }
276
277    #[test]
278    fn test_palette_sorted_descending() {
279        // 红 60% / 绿 30% / 蓝 10%，验证调色板按占比降序排列
280        let mut pixels = vec![[255u8, 0, 0]; 60];
281        pixels.extend(vec![[0u8, 255, 0]; 30]);
282        pixels.extend(vec![[0u8, 0, 255]; 10]);
283        let img = make_image(&pixels);
284
285        for alg in [Algorithm::KMeans, Algorithm::MedianCut, Algorithm::Octree] {
286            let palette = DominantColors::new(img.clone())
287                .config(Config::default().max_colors(3).sample_size(None))
288                .extract(alg)
289                .expect("三色图片应成功");
290
291            let percentages: Vec<f32> = palette.iter().map(|c| c.percentage).collect();
292            let mut sorted = percentages.clone();
293            sorted.sort_by(|a, b| b.partial_cmp(a).unwrap());
294            assert_eq!(percentages, sorted, "{alg:?} 调色板应按占比降序排列");
295        }
296    }
297
298    #[test]
299    fn test_config_builder() {
300        // 验证 Builder 链式调用可正确设置各字段
301        let cfg = Config::default()
302            .max_colors(10)
303            .sample_size(Some(128))
304            .kmeans_seed(99)
305            .kmeans_max_iterations(50);
306        assert_eq!(cfg.max_colors, 10);
307        assert_eq!(cfg.sample_size, Some(128));
308        assert_eq!(cfg.kmeans_seed, 99);
309        assert_eq!(cfg.kmeans_max_iterations, 50);
310    }
311
312    #[test]
313    #[should_panic(expected = "max_colors must be at least 1")]
314    fn test_config_zero_colors_panics() {
315        // max_colors 设为 0 时应 panic
316        let _ = Config::default().max_colors(0);
317    }
318}
dominant_colors/lib.rs

dominant_colors/
lib.rs
