Add modified zeng palette sorting method (#602)

This PR adds the modified zeng ("mzeng") palette sorting method, in
addition to the existing luma and battiato methods. Speed is very
similar to the battiato method with slightly better results on average.

Resulting sizes from two different image sets (all indexed or able to be
indexed):
| | master | PR |
|-|-|-|
| Set 1 | 29,647,156 | 29,555,697 |
| Set 2 | 23,732,133 | 23,570,862 |

Additionally, I've added a new "first colour" heuristic for both the
mzeng and battiato methods: We use the most popular colour overall, but
only if it covers at least 15% of the image. This provided 13k savings
on Set 2 vs the edge colour heuristic (which is still used in the luma
sort).

benches/reductions.rs

@@ -300,6 +300,16 @@ fn reductions_palette_sort(b: &mut Bencher) {
     b.iter(|| palette::sorted_palette(&png.raw));
 }
 
+#[bench]
+fn reductions_palette_sort_mzeng(b: &mut Bencher) {
+    let input = test::black_box(PathBuf::from(
+        "tests/files/palette_8_should_be_palette_8.png",
+    ));
+    let png = PngData::new(&input, &Options::default()).unwrap();
+
+    b.iter(|| palette::sorted_palette_mzeng(&png.raw));
+}
+
 #[bench]
 fn reductions_palette_sort_battiato(b: &mut Bencher) {
     let input = test::black_box(PathBuf::from(

src/reduction/mod.rs

@@ -1,6 +1,6 @@
 use std::sync::Arc;
 
-use crate::{evaluate::Evaluator, png::PngImage, Deadline, Deflaters, Options};
+use crate::{evaluate::Evaluator, png::PngImage, ColorType, Deadline, Deflaters, Options};
 
 pub mod alpha;
 use crate::alpha::*;
@@ -132,12 +132,41 @@ pub(crate) fn perform_reductions(
         }
     }
 
-    // Attempt to sort the palette using an alternative method
+    // Attempt additional palette sorting techniques
-    if !cheap && opts.palette_reduction && !deadline.passed() {
+    if !cheap && opts.palette_reduction {
-        // Make sure we use the `indexed` var if it exists
+        // Collect a list of palettes so we can avoid evaluating the same one twice
-        if let Some(reduced) = sorted_palette_battiato(indexed.as_ref().unwrap_or(&png)) {
+        let mut palettes = Vec::new();
-            eval.try_image(Arc::new(reduced));
+        if let ColorType::Indexed { palette } = &baseline.ihdr.color_type {
-            evaluation_added = true;
+            palettes.push(palette.clone());
+        }
+        // Make sure we use the `indexed` var as input if it exists
+        // This one doesn't need to be kept in the palette list as the sorters will fail if there's no change
+        let input = indexed.as_ref().unwrap_or(&png);
+
+        // Attempt to sort the palette using the battiato method
+        if !deadline.passed() {
+            if let Some(reduced) = sorted_palette_battiato(input) {
+                if let ColorType::Indexed { palette } = &reduced.ihdr.color_type {
+                    if !palettes.contains(palette) {
+                        palettes.push(palette.clone());
+                        eval.try_image(Arc::new(reduced));
+                        evaluation_added = true;
+                    }
+                }
+            }
+        }
+
+        // Attempt to sort the palette using the mzeng method
+        if !deadline.passed() {
+            if let Some(reduced) = sorted_palette_mzeng(input) {
+                if let ColorType::Indexed { palette } = &reduced.ihdr.color_type {
+                    if !palettes.contains(palette) {
+                        palettes.push(palette.clone());
+                        eval.try_image(Arc::new(reduced));
+                        evaluation_added = true;
+                    }
+                }
+            }
         }
     }
 

src/reduction/palette.rs

@@ -107,14 +107,14 @@ pub fn sorted_palette(png: &PngImage) -> Option<PngImage> {
     enumerated.insert(0, first);
 
     // Extract the new palette and determine if anything changed
-    let (old_map, palette): (Vec<_>, Vec<RGBA8>) = enumerated.into_iter().unzip();
+    let (remapping, palette): (Vec<_>, Vec<RGBA8>) = enumerated.into_iter().unzip();
-    if old_map.iter().enumerate().all(|(a, b)| a == *b) {
+    if remapping.iter().enumerate().all(|(a, b)| a == *b) {
         return None;
     }
 
     // Construct the new mapping and convert the data
     let mut byte_map = [0; 256];
-    for (i, &v) in old_map.iter().enumerate() {
+    for (i, &v) in remapping.iter().enumerate() {
         byte_map[v] = i as u8;
     }
     let data = png.data.iter().map(|&b| byte_map[b as usize]).collect();
@@ -128,7 +128,29 @@ pub fn sorted_palette(png: &PngImage) -> Option<PngImage> {
     })
 }
 
-/// Sort the colors in the palette by minimizing entropy, returning the sorted image if successful
+/// Sort the colors in the palette using the mzeng technique, returning the sorted image if successful
+#[must_use]
+pub fn sorted_palette_mzeng(png: &PngImage) -> Option<PngImage> {
+    // Interlacing not currently supported
+    if png.ihdr.bit_depth != BitDepth::Eight || png.ihdr.interlaced != Interlacing::None {
+        return None;
+    }
+    let palette = match &png.ihdr.color_type {
+        // Images with only two colors will remain unchanged from previous luma sort
+        ColorType::Indexed { palette } if palette.len() > 2 => palette,
+        _ => return None,
+    };
+
+    let matrix = co_occurrence_matrix(palette.len(), png);
+    let edges = weighted_edges(&matrix);
+    let mut remapping = mzeng_reindex(palette.len(), edges, &matrix);
+
+    apply_most_popular_color(png, &mut remapping);
+
+    apply_palette_reorder(png, &remapping)
+}
+
+/// Sort the colors in the palette using the battiato technique, returning the sorted image if successful
 #[must_use]
 pub fn sorted_palette_battiato(png: &PngImage) -> Option<PngImage> {
     // Interlacing not currently supported
@@ -143,28 +165,28 @@ pub fn sorted_palette_battiato(png: &PngImage) -> Option<PngImage> {
 
     let matrix = co_occurrence_matrix(palette.len(), png);
     let edges = weighted_edges(&matrix);
-    let mut old_map = battiato_tsp(palette.len(), edges);
+    let mut remapping = battiato_reindex(palette.len(), edges);
 
-    // Put the most popular edge color first, which can help slightly if the filter bytes are 0
+    apply_most_popular_color(png, &mut remapping);
-    let keep_first = most_popular_edge_color(palette.len(), png);
+
-    let first_idx = old_map.iter().position(|&i| i == keep_first).unwrap();
+    apply_palette_reorder(png, &remapping)
-    // If the index is past halfway, reverse the order so as to minimize the change
+}
-    if first_idx >= old_map.len() / 2 {
+
-        old_map.reverse();
+// Apply the palette reordering to the image data
-        old_map.rotate_right(first_idx + 1);
+fn apply_palette_reorder(png: &PngImage, remapping: &[usize]) -> Option<PngImage> {
-    } else {
+    let ColorType::Indexed { palette } = &png.ihdr.color_type else {
-        old_map.rotate_left(first_idx);
+        return None;
-    }
+    };
 
     // Check if anything changed
-    if old_map.iter().enumerate().all(|(a, b)| a == *b) {
+    if remapping.iter().enumerate().all(|(a, b)| a == *b) {
         return None;
     }
 
     // Construct the palette and byte maps and convert the data
     let mut new_palette = Vec::new();
     let mut byte_map = [0; 256];
-    for (i, &v) in old_map.iter().enumerate() {
+    for (i, &v) in remapping.iter().enumerate() {
         new_palette.push(palette[v]);
         byte_map[v] = i as u8;
     }
@@ -200,6 +222,38 @@ fn most_popular_edge_color(num_colors: usize, png: &PngImage) -> usize {
         .0
 }
 
+// Find the most popular color in the image, along with its count
+fn most_popular_color(num_colors: usize, png: &PngImage) -> (usize, u32) {
+    let mut counts = [0u32; 256];
+    for &val in &png.data {
+        counts[val as usize] += 1;
+    }
+    counts
+        .iter()
+        .copied()
+        .take(num_colors)
+        .enumerate()
+        .max_by_key(|&(_, v)| v)
+        .unwrap_or_default()
+}
+
+// Put the most popular color first
+fn apply_most_popular_color(png: &PngImage, remapping: &mut [usize]) {
+    let most_popular = most_popular_color(remapping.len(), png);
+    // If the most popular color is less than 15% of the image, don't use it
+    if most_popular.1 < png.data.len() as u32 * 3 / 20 {
+        return;
+    }
+    let first_idx = remapping.iter().position(|&i| i == most_popular.0).unwrap();
+    // If the index is past halfway, reverse the order so as to minimize the change
+    if first_idx >= remapping.len() / 2 {
+        remapping.reverse();
+        remapping.rotate_right(first_idx + 1);
+    } else {
+        remapping.rotate_left(first_idx);
+    }
+}
+
 // Calculate co-occurences matrix
 fn co_occurrence_matrix(num_colors: usize, png: &PngImage) -> Vec<Vec<u32>> {
     let mut matrix = vec![vec![0u32; num_colors]; num_colors];
@@ -213,9 +267,15 @@ fn co_occurrence_matrix(num_colors: usize, png: &PngImage) -> Vec<Vec<u32>> {
             }
             if let Some(prev_val) = prev_val.replace(val) {
                 matrix[prev_val][val] += 1;
+                matrix[val][prev_val] += 1;
             }
             if let Some(prev) = &prev {
-                matrix[prev.data[i] as usize][val] += 1;
+                let prev_val = prev.data[i] as usize;
+                if prev_val > num_colors {
+                    continue;
+                }
+                matrix[prev_val][val] += 1;
+                matrix[val][prev_val] += 1;
             }
         }
         prev = Some(line)
@@ -226,19 +286,76 @@ fn co_occurrence_matrix(num_colors: usize, png: &PngImage) -> Vec<Vec<u32>> {
 // Calculate edge list sorted by weight
 fn weighted_edges(matrix: &[Vec<u32>]) -> Vec<(usize, usize)> {
     let mut edges = Vec::new();
-    for i in 0..matrix.len() {
+    for (i, m_row) in matrix.iter().enumerate() {
-        for j in 0..i {
+        for (j, val) in m_row.iter().enumerate().take(i) {
-            edges.push(((j, i), matrix[i][j] + matrix[j][i]));
+            edges.push(((j, i), val));
         }
     }
     edges.sort_by(|(_, w1), (_, w2)| w2.cmp(w1));
     edges.into_iter().map(|(e, _)| e).collect()
 }
 
+// Apply a greedy index assignment using the modified version of Zeng's techinque from
+// "A note on Zeng's technique for color reindexing of palette-based images" by Pinho et al
+// https://ieeexplore.ieee.org/document/1261987
+// Based on the C implementation in libwebp
+fn mzeng_reindex(num_colors: usize, edges: Vec<(usize, usize)>, matrix: &[Vec<u32>]) -> Vec<usize> {
+    // Initialize the mapping list with the two best indices.
+    let mut remapping = vec![edges[0].0, edges[0].1];
+
+    // Initialize the sums with the first two remappings and find the best one
+    let mut sums = Vec::new();
+    let mut best_sum_pos = 0;
+    let mut best_sum = (0, 0);
+    for (i, m_row) in matrix.iter().enumerate() {
+        if i == remapping[0] || i == remapping[1] {
+            continue;
+        }
+        let sum = (i, m_row[remapping[0]] + m_row[remapping[1]]);
+        if sum.1 > best_sum.1 {
+            best_sum_pos = sums.len();
+            best_sum = sum;
+        }
+        sums.push(sum);
+    }
+
+    while !sums.is_empty() {
+        let best_index = best_sum.0;
+        // Compute delta to know if we need to prepend or append the best index.
+        let mut delta: isize = 0;
+        let n = (num_colors - sums.len()) as isize;
+        for (i, &index) in remapping.iter().enumerate() {
+            delta += (n - 1 - 2 * i as isize) * matrix[best_index][index] as isize;
+        }
+        if delta > 0 {
+            remapping.insert(0, best_index);
+        } else {
+            remapping.push(best_index);
+        }
+        // Remove best_sum from sums.
+        sums.swap_remove(best_sum_pos);
+        if !sums.is_empty() {
+            // Update all the sums and find the best one.
+            best_sum_pos = 0;
+            best_sum = (0, 0);
+            for (i, sum) in sums.iter_mut().enumerate() {
+                sum.1 += matrix[best_index][sum.0];
+                if sum.1 > best_sum.1 {
+                    best_sum_pos = i;
+                    best_sum = *sum;
+                }
+            }
+        }
+    }
+
+    // Return the completed remapping
+    remapping
+}
+
 // Calculate an approximate solution of the Traveling Salesman Problem using the algorithm
 // from "An efficient Re-indexing algorithm for color-mapped images" by Battiato et al
 // https://ieeexplore.ieee.org/document/1344033
-fn battiato_tsp(num_colors: usize, edges: Vec<(usize, usize)>) -> Vec<usize> {
+fn battiato_reindex(num_colors: usize, edges: Vec<(usize, usize)>) -> Vec<usize> {
     let mut chains = Vec::new();
     // Keep track of the state of each vertex (.0) and it's chain number (.1)
     // 0 = an unvisited vertex (White)