about summary refs log tree commit diff
path: root/src
diff options
context:
space:
mode:
Diffstat (limited to 'src')
-rw-r--r--src/lib.rs35
1 files changed, 35 insertions, 0 deletions
diff --git a/src/lib.rs b/src/lib.rs
index e0383a3..38f8e1a 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,3 +1,5 @@
+#[cfg(kmeans)]
+use kmeans::{KMeans, KMeansConfig};
 use rgb::{ComponentBytes, FromSlice, RGB8};
 use std::collections::HashMap;
 
@@ -95,6 +97,7 @@ impl<T: Count> Squasher<T> {
 	}
 
 	/// Create a new palette from the colours in the given image.
+	#[cfg(not(kmeans))]
 	pub fn recolor<'a, Img>(&mut self, image: Img)
 	where
 		Img: Into<ImageData<'a>>,
@@ -104,6 +107,37 @@ impl<T: Count> Squasher<T> {
 		self.palette = selected;
 	}
 
+	#[cfg(kmeans)]
+	pub fn recolor<'a, Img>(&mut self, image: Img)
+	where
+		Img: Into<ImageData<'a>>,
+	{
+		let ImageData(rgb) = image.into();
+
+		let kmean = KMeans::new(
+			rgb.as_bytes()
+				.iter()
+				.map(|u| *u as f32)
+				.collect::<Vec<f32>>(),
+			rgb.as_bytes().len() / 3,
+			3,
+		);
+		let k = self.max_colours_min1.as_usize() + 1;
+		let result =
+			kmean.kmeans_lloyd(k, 100, KMeans::init_kmeanplusplus, &KMeansConfig::default());
+		self.palette = result
+			.centroids
+			.chunks_exact(3)
+			.map(|rgb| {
+				RGB8::new(
+					rgb[0].round() as u8,
+					rgb[1].round() as u8,
+					rgb[2].round() as u8,
+				)
+			})
+			.collect();
+	}
+
 	/// Create a Squasher from parts. Noteably, this leave your palette empty
 	fn from_parts(max_colours_min1: T, difference_fn: Box<DiffFn>, tolerance: f32) -> Self {
 		Self {
@@ -208,6 +242,7 @@ impl<T: Count> Squasher<T> {
 
 	/// Pick the colors in the palette from a Vec of colors sorted by number
 	/// of times they occur, high to low.
+	#[cfg(not(kmeans))]
 	fn select_colors(&self, sorted: Vec<RGB8>) -> Vec<RGB8> {
 		let tolerance = (self.tolerance_percent / 100.0) * 765.0;
 		let max_colours = self.max_colours_min1.as_usize() + 1;