Better quantization and dithering

Cargo.toml

@@ -17,5 +17,6 @@ debug-assertions = false
 "capi" = []
 
 [dependencies]
+vpsearch = "2.0.1"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html

src/cluster.rs

@@ -1,44 +1,28 @@
-use std::cmp::{Ord,Ordering};
+use std::cmp::Ordering;
 
 use crate::histogram::{Histogram, HistogramEntry};
 
 pub struct Cluster {
     pub entries: Vec<HistogramEntry>,
-    pub mean: [u8; 4],
-    pub weight: usize,
-    pub priority: u64,
+    pub mean: [f64; 4],
+    pub weight: f64,
+    pub chan_diff: f64,
     widest_chan: u8,
 }
 
-impl Ord for Cluster {
-    fn cmp(&self, other: &Self) -> Ordering {
-        self.priority.cmp(&other.priority)
-    }
-}
-
-impl PartialOrd for Cluster {
-    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
-        Some(self.cmp(other))
-    }
-}
-
-impl Eq for Cluster {}
-
-impl PartialEq for Cluster {
-    fn eq(&self, other: &Self) -> bool {
-        self.priority == other.priority
-    }
-}
-
 impl Cluster {
     pub fn new(entries: Vec<HistogramEntry>) -> Self {
-        Self{
+        let mut cluster = Self{
             entries: entries,
-            mean: [0; 4],
-            priority: 0,
-            weight: 0,
+            mean: [0.0; 4],
+            weight: 0.0,
+            chan_diff: 0.0,
             widest_chan: 0,
-        }
+        };
+
+        cluster.calc_stats();
+
+        cluster
     }
 
     pub fn from_histogram(hist: &Histogram) -> Self {
@@ -51,33 +35,47 @@ impl Cluster {
         Self::new(entries)
     }
 
-    pub fn calc_widest_and_priority(&mut self) {
-        let mut diff_sum: [usize; 4] = [0; 4];
+    fn calc_stats(&mut self) {
+        self.mean = [0.0; 4];
+        self.weight = 0.0;
+
+        if self.entries.is_empty() {
+            self.chan_diff = 0.0;
+            return
+        }
 
         for e in self.entries.iter() {
-            let weight = e.weight as usize;
+            let weight = e.weight as f64;
+
+            self.mean[0] += e.color[0] as f64 * weight;
+            self.mean[1] += e.color[1] as f64 * weight;
+            self.mean[2] += e.color[2] as f64 * weight;
+            self.mean[3] += e.color[3] as f64 * weight;
 
-            diff_sum[0] += diff(e.color[0], self.mean[0]) as usize * weight;
-            diff_sum[1] += diff(e.color[1], self.mean[1]) as usize * weight;
-            diff_sum[2] += diff(e.color[2], self.mean[2]) as usize * weight;
-            diff_sum[3] += diff(e.color[3], self.mean[3]) as usize * weight;
+            self.weight += weight;
         }
 
-        let mut chan = 0;
-        let mut max_diff_sum = 0;
+        self.mean[0] /= self.weight;
+        self.mean[1] /= self.weight;
+        self.mean[2] /= self.weight;
+        self.mean[3] /= self.weight;
 
-        for ch in 0..=3usize {
-            let d = diff_sum[ch];
+        let mut diff_sum: [f64; 4] = [0f64; 4];
 
-            if d > max_diff_sum {
-                chan = ch;
-                max_diff_sum = d;
-            }
+        for e in self.entries.iter() {
+            let weight = e.weight as f64;
+
+            diff_sum[0] += (e.color[0] as f64 - self.mean[0]).abs() * weight;
+            diff_sum[1] += (e.color[1] as f64 - self.mean[1]).abs() * weight;
+            diff_sum[2] += (e.color[2] as f64 - self.mean[2]).abs() * weight;
+            diff_sum[3] += (e.color[3] as f64 - self.mean[3]).abs() * weight;
         }
 
-        let chan_diff = max_diff_sum as f64 / self.weight as f64;
+        let (chan, max_diff_sum) = diff_sum.iter().enumerate()
+            .max_by(|&(_, a), &(_, b)| a.partial_cmp(&b).unwrap_or(Ordering::Equal))
+            .unwrap();
 
-        self.priority = (chan_diff * (self.weight as f64).sqrt()) as u64;
+        self.chan_diff = max_diff_sum / self.weight;
         self.widest_chan = chan as u8;
     }
 
@@ -93,71 +91,32 @@ impl Cluster {
         let mut gt_weight: usize = 0;
 
         while i <= gt {
-            let val = self.entries[i].color[widest_chan];
+            let entry = &self.entries[i];
+            let val = entry.color[widest_chan] as f64;
 
             if val < widest_chan_mean {
+                lt_weight += entry.weight as usize;
                 if lt != i {
                     self.entries.swap(lt, i);
                 }
                 lt += 1;
                 i += 1;
-                lt_weight += self.entries[i].weight as usize;
             } else if val > widest_chan_mean {
+                gt_weight += entry.weight as usize;
                 self.entries.swap(gt, i);
                 gt -= 1;
-                gt_weight += self.entries[i].weight as usize;
             } else {
                 i += 1;
             }
         }
 
         let mut split_pos = i;
-        if lt_weight < gt_weight {
+        if lt_weight > gt_weight {
             split_pos = lt;
         }
 
         let (sp1, sp2) = self.entries.split_at(split_pos);
 
         (Self::new(sp1.to_vec()), Self::new(sp2.to_vec()))
     }
-
-    pub fn calc_mean_and_weight(&mut self) {
-        self.weight = 0;
-
-        if self.entries.is_empty() {
-            self.mean = [0; 4];
-            return
-        }
-
-        let mut rsum: usize = 0;
-        let mut gsum: usize = 0;
-        let mut bsum: usize = 0;
-        let mut asum: usize = 0;
-
-        for e in self.entries.iter() {
-            let weight = e.weight as usize;
-
-            rsum += e.color[0] as usize * weight;
-            gsum += e.color[1] as usize * weight;
-            bsum += e.color[2] as usize * weight;
-            asum += e.color[3] as usize * weight;
-
-            self.weight += weight;
-        }
-
-        self.mean = [
-            (rsum / self.weight) as u8,
-            (gsum / self.weight) as u8,
-            (bsum / self.weight) as u8,
-            (asum / self.weight) as u8,
-        ]
-    }
-}
-
-fn diff(a: u8, b: u8) -> u8 {
-    if a > b {
-        return a - b
-    }
-
-    b - a
 }