implement kmeans palette optimization

examples/demo.rs

@@ -16,6 +16,9 @@ fn main() {
     println!("Generating palette");
     let palette = exoquant::create_palette(&hist, 256);
 
+    println!("Optimize palette (k-means)");
+    let palette = exoquant::optimize_palette(palette, &hist, 4);
+
     let mut state = lodepng::State::new();
     for color in &palette {
         unsafe {

src/colormap.rs

@@ -141,8 +141,12 @@ impl KDNode {
 impl ColorMap {
     pub fn new(colors: &[Color]) -> ColorMap {
         let float_colors: Vec<_> = colors.iter().map(|c| c.into()).collect();
-        let kdtree = KDNode::new((0..colors.len()).collect(), &float_colors);
-        let neighbor_distance = float_colors.iter()
+        Self::from_float_colors(float_colors)
+    }
+
+    pub fn from_float_colors(colors: Vec<FloatColor>) -> ColorMap {
+        let kdtree = KDNode::new((0..colors.len()).collect(), &colors);
+        let neighbor_distance = colors.iter()
             .enumerate()
             .map(|(i, c)| {
                 if let Some(nearest) = kdtree.find_nearest(*c, ::std::f64::MAX, i) {
@@ -155,7 +159,7 @@ impl ColorMap {
         ColorMap {
             kdtree: kdtree,
             neighbor_distance: neighbor_distance,
-            colors: float_colors,
+            colors: colors,
         }
     }
 

src/histogram.rs

@@ -1,15 +1,28 @@
 use std::collections::HashMap;
 
 use super::Color;
+use ::quantizer::HistColor;
 
 pub struct Histogram {
-    pub data: HashMap<Color, usize>,
+    data: HashMap<Color, usize>,
 }
 
 impl Histogram {
     pub fn new() -> Histogram {
         Histogram { data: HashMap::new() }
     }
+
+    pub fn to_hist_colors(&self) -> Vec<HistColor> {
+        self.data
+            .iter()
+            .map(|(color, count)| {
+                HistColor {
+                    color: color.into(),
+                    count: *count,
+                }
+            })
+            .collect()
+    }
 }
 
 impl Extend<Color> for Histogram {

src/kmeans.rs

@@ -0,0 +1,41 @@
+use ::color::FloatColor;
+use ::color::Color;
+use ::quantizer::HistColor;
+use ::colormap::ColorMap;
+use ::histogram::Histogram;
+
+struct KMeansCluster {
+    sum: FloatColor,
+    count: usize,
+}
+
+pub fn kmeans_step(colors: Vec<FloatColor>, histogram: &[HistColor]) -> Vec<FloatColor> {
+    let map = ColorMap::from_float_colors(colors.iter().cloned().collect());
+    let mut clusters: Vec<_> = (0..colors.len())
+        .map(|_| {
+            KMeansCluster {
+                sum: FloatColor::default(),
+                count: 0,
+            }
+        })
+        .collect();
+    for entry in histogram {
+        let index = map.find_nearest_float(entry.color);
+        let mut cluster = &mut clusters[index];
+        cluster.sum += entry.color * entry.count as f64;
+        cluster.count += entry.count;
+    }
+    clusters.iter().map(|cluster| cluster.sum * (1.0 / cluster.count as f64)).collect()
+}
+
+pub fn optimize_palette(palette: Vec<Color>,
+                        histogram: &Histogram,
+                        num_iterations: usize)
+                        -> Vec<Color> {
+    let hist = histogram.to_hist_colors();
+    let mut colors = palette.iter().map(|c| c.into()).collect();
+    for _ in 0..num_iterations {
+        colors = kmeans_step(colors, &hist);
+    }
+    colors.iter().map(|&c| c.into()).collect()
+}

src/lib.rs

@@ -3,8 +3,10 @@ mod histogram;
 mod quantizer;
 pub mod colormap;
 mod remapper;
+mod kmeans;
 
 pub use color::Color;
 pub use histogram::Histogram;
 pub use quantizer::create_palette;
 pub use remapper::{Remapper, RemapperNoDither, RemapperOrdered};
+pub use kmeans::optimize_palette;

src/quantizer.rs

@@ -1,9 +1,9 @@
 use ::color::FloatColor;
 use ::color::Color;
 
-struct HistColor {
-    color: FloatColor,
-    count: usize,
+pub struct HistColor {
+    pub color: FloatColor,
+    pub count: usize,
 }
 
 struct QuantizerNode {
@@ -101,15 +101,7 @@ impl QuantizerNode {
 }
 
 pub fn create_palette(histogram: &::histogram::Histogram, num_colors: usize) -> Vec<Color> {
-    let hist: Vec<_> = histogram.data
-        .iter()
-        .map(|(color, count)| {
-            HistColor {
-                color: color.into(),
-                count: *count,
-            }
-        })
-        .collect();
+    let hist: Vec<_> = histogram.to_hist_colors();
 
     let mut nodes = vec![QuantizerNode::new(hist)];