about summary refs log tree commit diff
diff options
context:
space:
mode:
-rw-r--r--.gitignore1
-rw-r--r--Cargo.lock237
-rw-r--r--Cargo.toml7
-rw-r--r--squash/Cargo.toml2
-rw-r--r--squash/src/main.rs4
-rw-r--r--src/lib.rs1
-rw-r--r--src/nih_kmeans.rs2
-rw-r--r--src/selection.rs52
8 files changed, 289 insertions, 17 deletions
diff --git a/.gitignore b/.gitignore
index ba7f17b..0b1ac8f 100644
--- a/.gitignore
+++ b/.gitignore
@@ -8,3 +8,4 @@
 flamegraph.svg
 perf.data
 perf.data.old
+.vscode
\ No newline at end of file
diff --git a/Cargo.lock b/Cargo.lock
index 7520455..a727669 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -15,6 +15,12 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "080e9890a082662b09c1ad45f567faeeb47f22b5fb23895fbe1e651e718e25ca"
 
 [[package]]
+name = "autocfg"
+version = "1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d468802bab17cbc0cc575e9b053f41e72aa36bfa6b7f55e3529ffa43161b97fa"
+
+[[package]]
 name = "bitflags"
 version = "1.3.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -58,10 +64,11 @@ checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd"
 
 [[package]]
 name = "colorsquash"
-version = "0.1.0"
+version = "0.2.0"
 dependencies = [
  "gifed",
- "rand",
+ "kmeans",
+ "rand 0.8.5",
  "rgb",
 ]
 
@@ -75,6 +82,37 @@ dependencies = [
 ]
 
 [[package]]
+name = "crossbeam-deque"
+version = "0.8.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "613f8cc01fe9cf1a3eb3d7f488fd2fa8388403e97039e2f73692932e291a770d"
+dependencies = [
+ "crossbeam-epoch",
+ "crossbeam-utils",
+]
+
+[[package]]
+name = "crossbeam-epoch"
+version = "0.9.18"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e"
+dependencies = [
+ "crossbeam-utils",
+]
+
+[[package]]
+name = "crossbeam-utils"
+version = "0.8.19"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "248e3bacc7dc6baa3b21e405ee045c3047101a49145e7e9eca583ab4c2ca5345"
+
+[[package]]
+name = "either"
+version = "1.9.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a26ae43d7bcc3b814de94796a5e736d4029efb0ee900c12e2d54c993ad1a1e07"
+
+[[package]]
 name = "fdeflate"
 version = "0.3.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -101,13 +139,24 @@ checksum = "e6d5a32815ae3f33302d95fdcb2ce17862f8c65363dcfd29360480ba1001fc9c"
 
 [[package]]
 name = "getrandom"
+version = "0.1.16"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8fc3cb4d91f53b50155bdcfd23f6a4c39ae1969c2ae85982b135750cccaf5fce"
+dependencies = [
+ "cfg-if",
+ "libc",
+ "wasi 0.9.0+wasi-snapshot-preview1",
+]
+
+[[package]]
+name = "getrandom"
 version = "0.2.12"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "190092ea657667030ac6a35e305e62fc4dd69fd98ac98631e5d3a2b1575a12b5"
 dependencies = [
  "cfg-if",
  "libc",
- "wasi",
+ "wasi 0.11.0+wasi-snapshot-preview1",
 ]
 
 [[package]]
@@ -119,12 +168,30 @@ dependencies = [
 ]
 
 [[package]]
+name = "kmeans"
+version = "0.2.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "76ccc6d18ad4bdf1b31a515991e73192cc1ef9e0ff06ea8ade4d95f80ee70352"
+dependencies = [
+ "num",
+ "packed_simd",
+ "rand 0.7.3",
+ "rayon",
+]
+
+[[package]]
 name = "libc"
 version = "0.2.152"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "13e3bf6590cbc649f4d1a3eefc9d5d6eb746f5200ffb04e5e142700b8faa56e7"
 
 [[package]]
+name = "libm"
+version = "0.2.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4ec2a862134d2a7d32d7983ddcdd1c4923530833c9f2ea1a44fc5fa473989058"
+
+[[package]]
 name = "log"
 version = "0.4.20"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -141,6 +208,93 @@ dependencies = [
 ]
 
 [[package]]
+name = "num"
+version = "0.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8b7a8e9be5e039e2ff869df49155f1c06bd01ade2117ec783e56ab0932b67a8f"
+dependencies = [
+ "num-bigint",
+ "num-complex",
+ "num-integer",
+ "num-iter",
+ "num-rational",
+ "num-traits",
+]
+
+[[package]]
+name = "num-bigint"
+version = "0.3.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5f6f7833f2cbf2360a6cfd58cd41a53aa7a90bd4c202f5b1c7dd2ed73c57b2c3"
+dependencies = [
+ "autocfg",
+ "num-integer",
+ "num-traits",
+]
+
+[[package]]
+name = "num-complex"
+version = "0.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "747d632c0c558b87dbabbe6a82f3b4ae03720d0646ac5b7b4dae89394be5f2c5"
+dependencies = [
+ "num-traits",
+]
+
+[[package]]
+name = "num-integer"
+version = "0.1.45"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "225d3389fb3509a24c93f5c29eb6bde2586b98d9f016636dff58d7c6f7569cd9"
+dependencies = [
+ "autocfg",
+ "num-traits",
+]
+
+[[package]]
+name = "num-iter"
+version = "0.1.43"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7d03e6c028c5dc5cac6e2dec0efda81fc887605bb3d884578bb6d6bf7514e252"
+dependencies = [
+ "autocfg",
+ "num-integer",
+ "num-traits",
+]
+
+[[package]]
+name = "num-rational"
+version = "0.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "12ac428b1cb17fce6f731001d307d351ec70a6d202fc2e60f7d4c5e42d8f4f07"
+dependencies = [
+ "autocfg",
+ "num-bigint",
+ "num-integer",
+ "num-traits",
+]
+
+[[package]]
+name = "num-traits"
+version = "0.2.17"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "39e3200413f237f41ab11ad6d161bc7239c84dcb631773ccd7de3dfe4b5c267c"
+dependencies = [
+ "autocfg",
+ "libm",
+]
+
+[[package]]
+name = "packed_simd"
+version = "0.3.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1f9f08af0c877571712e2e3e686ad79efad9657dbf0f7c3c8ba943ff6c38932d"
+dependencies = [
+ "cfg-if",
+ "num-traits",
+]
+
+[[package]]
 name = "png"
 version = "0.17.10"
 source = "git+https://github.com/image-rs/image-png.git?rev=f10238a1e886b228e7da5301e5c0f5011316f2d6#f10238a1e886b228e7da5301e5c0f5011316f2d6"
@@ -166,13 +320,36 @@ checksum = "dc33ff2d4973d518d823d61aa239014831e521c75da58e3df4840d3f47749d09"
 
 [[package]]
 name = "rand"
+version = "0.7.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6a6b1679d49b24bbfe0c803429aa1874472f50d9b363131f0e89fc356b544d03"
+dependencies = [
+ "getrandom 0.1.16",
+ "libc",
+ "rand_chacha 0.2.2",
+ "rand_core 0.5.1",
+ "rand_hc",
+]
+
+[[package]]
+name = "rand"
 version = "0.8.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404"
 dependencies = [
  "libc",
- "rand_chacha",
- "rand_core",
+ "rand_chacha 0.3.1",
+ "rand_core 0.6.4",
+]
+
+[[package]]
+name = "rand_chacha"
+version = "0.2.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f4c8ed856279c9737206bf725bf36935d8666ead7aa69b52be55af369d193402"
+dependencies = [
+ "ppv-lite86",
+ "rand_core 0.5.1",
 ]
 
 [[package]]
@@ -182,7 +359,16 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88"
 dependencies = [
  "ppv-lite86",
- "rand_core",
+ "rand_core 0.6.4",
+]
+
+[[package]]
+name = "rand_core"
+version = "0.5.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "90bde5296fc891b0cef12a6d03ddccc162ce7b2aff54160af9338f8d40df6d19"
+dependencies = [
+ "getrandom 0.1.16",
 ]
 
 [[package]]
@@ -191,7 +377,36 @@ version = "0.6.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c"
 dependencies = [
- "getrandom",
+ "getrandom 0.2.12",
+]
+
+[[package]]
+name = "rand_hc"
+version = "0.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ca3129af7b92a17112d59ad498c6f81eaf463253766b90396d39ea7a39d6613c"
+dependencies = [
+ "rand_core 0.5.1",
+]
+
+[[package]]
+name = "rayon"
+version = "1.8.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9c27db03db7734835b3f53954b534c91069375ce6ccaa2e065441e07d9b6cdb1"
+dependencies = [
+ "either",
+ "rayon-core",
+]
+
+[[package]]
+name = "rayon-core"
+version = "1.12.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5ce3fb6ad83f861aac485e76e1985cd109d9a3713802152be56c3b1f0e0658ed"
+dependencies = [
+ "crossbeam-deque",
+ "crossbeam-utils",
 ]
 
 [[package]]
@@ -211,7 +426,7 @@ checksum = "d66dc143e6b11c1eddc06d5c423cfc97062865baf299914ab64caa38182078fe"
 
 [[package]]
 name = "squash"
-version = "0.2.0"
+version = "0.3.0"
 dependencies = [
  "anyhow",
  "camino",
@@ -229,6 +444,12 @@ checksum = "55937e1799185b12863d447f42597ed69d9928686b8d88a1df17376a097d8369"
 
 [[package]]
 name = "wasi"
+version = "0.9.0+wasi-snapshot-preview1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cccddf32554fecc6acb585f82a32a72e28b48f8c4c1883ddfeeeaa96f7d8e519"
+
+[[package]]
+name = "wasi"
 version = "0.11.0+wasi-snapshot-preview1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423"
diff --git a/Cargo.toml b/Cargo.toml
index e0cdfc4..21a50b0 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -13,6 +13,13 @@ repository = "https://github.com/gennyble/colorsquash"
 rgb = "0.8.36"
 gifed = { path = "../gifed/gifed", optional = true }
 rand = { version = "0.8.5", optional = true }
+kmeans = { version = "0.2.1", optional = true }
+
+[features]
+default = ["simd-kmeans"]
+# use the kmeans crate instead of the internal kmeans implementation. the crate
+# is faster and uses SIMD but requries nightly Rust.
+simd-kmeans = ["kmeans"]
 
 [workspace]
 members = ["squash"]
diff --git a/squash/Cargo.toml b/squash/Cargo.toml
index 7d33b5b..971afce 100644
--- a/squash/Cargo.toml
+++ b/squash/Cargo.toml
@@ -11,7 +11,7 @@ repository = "https://github.com/gennyble/colorsquash/tree/main/squash"
 
 [dependencies]
 # the meat 'o the thing! the meaning behind it all
-colorsquash = { path = "..", version = "0.2.0", features = ["gifed", "kmeans"] }
+colorsquash = { path = "..", version = "0.2.0", features = ["gifed"] }
 
 # just useful tools for writing binaries
 anyhow = "1.0.75"
diff --git a/squash/src/main.rs b/squash/src/main.rs
index a7cc66b..24d557f 100644
--- a/squash/src/main.rs
+++ b/squash/src/main.rs
@@ -31,10 +31,12 @@ fn main() -> Result<(), anyhow::Error> {
 
 			builder = builder.selector(sorsel);
 		}
-		cli::Selector::Kmeans => builder = builder.selector(Kmeans),
+		cli::Selector::Kmeans => builder = builder.selector(Kmeans { max_iter: 10 }),
 	};
 
+	let start = std::time::Instant::now();
 	let mut squasher = builder.build(&image.data);
+	println!("{:.2}ms", start.elapsed().as_secs_f32());
 
 	let size = squasher.map_over(&mut image.data);
 	image.data.resize(size, 0);
diff --git a/src/lib.rs b/src/lib.rs
index 9ab6b5c..cbecf76 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -3,6 +3,7 @@ use std::collections::HashSet;
 use rgb::{ComponentBytes, FromSlice, RGB8};
 
 pub mod difference;
+#[cfg(not(feature = "simd-kmeans"))]
 mod nih_kmeans;
 pub mod selection;
 
diff --git a/src/nih_kmeans.rs b/src/nih_kmeans.rs
index a34d528..ee752bc 100644
--- a/src/nih_kmeans.rs
+++ b/src/nih_kmeans.rs
@@ -1,6 +1,6 @@
 use std::collections::HashMap;
 
-#[cfg(rand)]
+#[cfg(feature = "rand")]
 use rand::{prelude::*, seq::index::sample};
 use rgb::{RGB, RGB8};
 
diff --git a/src/selection.rs b/src/selection.rs
index dacd735..d93603b 100644
--- a/src/selection.rs
+++ b/src/selection.rs
@@ -1,8 +1,10 @@
 use std::collections::HashMap;
 
-#[cfg(feature = "kmeans")]
+#[cfg(not(feature = "simd-kmeans"))]
 use crate::nih_kmeans::KMeans;
-use rgb::{ComponentBytes, RGB8};
+#[cfg(feature = "simd-kmeans")]
+use kmeans::{KMeans, KMeansConfig};
+use rgb::RGB8;
 
 use crate::{
 	difference::{self, DiffFn},
@@ -103,16 +105,54 @@ impl Default for SortSelect {
 	}
 }
 
-#[cfg(feature = "kmeans")]
 #[derive(Debug, Default)]
-pub struct Kmeans;
+pub struct Kmeans {
+	pub max_iter: usize,
+}
 
-#[cfg(feature = "kmeans")]
+#[cfg(not(feature = "simd-kmeans"))]
 impl Selector for Kmeans {
 	fn select(&mut self, max_colors: usize, image: ImageData) -> Vec<RGB8> {
 		let ImageData(rgb) = image;
 
 		let kmean = KMeans::new(rgb.to_vec());
-		kmean.get_k_colors(max_colors, max_iter)
+		kmean.get_k_colors(max_colors, self.max_iter)
+	}
+}
+
+#[cfg(feature = "simd-kmeans")]
+impl Selector for Kmeans {
+	fn select(&mut self, max_colors: usize, image: ImageData) -> Vec<RGB8> {
+		use rgb::ComponentBytes;
+
+		let ImageData(rgb) = image;
+
+		let kmean = KMeans::new(
+			rgb.as_bytes()
+				.iter()
+				.map(|u| *u as f32)
+				.collect::<Vec<f32>>(),
+			rgb.as_bytes().len() / 3,
+			3,
+		);
+
+		let result = kmean.kmeans_lloyd(
+			max_colors,
+			self.max_iter,
+			KMeans::init_kmeanplusplus,
+			&KMeansConfig::default(),
+		);
+
+		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()
 	}
 }