[Custom Levels] Improve color palette generation (#3797)

A few improvements to color palette selection. These were made by
tracing some particularly bad colors through and seeing where it made
obviously bad decisions for splitting. I tested on crystal cave, and a
test GLBs from Kuitar that previously had issues with alpha.

- The previous approach to splitting was based on trying to keep a tree
of deduplicated colors balanced (same count in each leaf). This is not
really a good idea for generating color palettes. A better approach is
to try to minimize the volume of the child node, limiting how inaccurate
a color can be. Splitting is now chosen based on the average of the
_deduplicated channel values_, which in practice seems to work pretty
well for Kuitar's levels. Other approaches could work here too.
- The previous approach of alternating through dimensions to split on
was kept.
- The depth of the KD tree during the initial split was increased,
allowing it to use up to 8192 colors, instead of just 1024.
- In most cases, not all child nodes of the tree have colors in them,
meaning that a tree of depth 13 would have less than 8192 colors. If
this happens, child nodes are split until the color count reaches 8192.
The selection of which nodes to split is somewhat arbitrary, but is
breadth-first. The axis for splitting is the one with the largest range.
(which might be a better idea in general?)

On crystal-cave, the worst case color error was reduced from 221 to 9.

---------

Co-authored-by: water111 <[email protected]>

goalc/build_level/common/color_quantization.cpp

@@ -1,6 +1,7 @@
 #include "color_quantization.h"
 
 #include <algorithm>
+#include <array>
 #include <set>
 #include <unordered_map>
 
@@ -232,6 +233,57 @@ struct KdNode {
   std::vector<Color> colors;
 };
 
+size_t pick_split_point(const std::vector<Color>& colors, int dim) {
+  if (colors.empty()) {
+    return 0;
+  }
+  double mean = 0;
+  int mean_count = 0;
+  std::array<bool, 256> seen;
+  seen.fill(false);
+
+  for (const auto& color : colors) {
+    const u8 val = color[dim];
+    if (!seen[val]) {
+      mean += val;
+      mean_count++;
+      seen[val] = true;
+    }
+  }
+  mean /= mean_count;
+
+  for (size_t i = 0; i < colors.size(); i++) {
+    if (colors.at(i)[dim] >= mean) {
+      return i;
+    }
+  }
+
+  return colors.size() - 1;
+}
+
+int pick_split_dim_final_splits(const std::vector<Color>& colors) {
+  int mins[4] = {255, 255, 255, 255};
+  int maxs[4] = {0, 0, 0, 0};
+  for (const auto& color : colors) {
+    for (int i = 0; i < 4; i++) {
+      mins[i] = std::min(mins[i], (int)color[i]);
+      maxs[i] = std::max(maxs[i], (int)color[i]);
+    }
+  }
+
+  int best_dim = 0;
+  int best_diff = 0;
+  for (int i = 0; i < 4; i++) {
+    const int diff = maxs[i] - mins[i];
+    if (diff > best_diff) {
+      best_diff = diff;
+      best_dim = i;
+    }
+  }
+
+  return best_dim;
+}
+
 void split_kd(KdNode* in, u32 depth, int next_split_dim) {
   if (!depth) {
     return;
@@ -264,7 +316,8 @@ void split_kd(KdNode* in, u32 depth, int next_split_dim) {
   in->right = std::make_unique<KdNode>();
 
   size_t i = 0;
-  size_t mid = in->colors.size() / 2;
+  // size_t mid = in->colors.size() / 2;
+  size_t mid = pick_split_point(in->colors, next_split_dim);
   if (depth & 1) {
     while (mid > 1 && in->colors[mid][next_split_dim] == in->colors[mid - 1][next_split_dim]) {
       mid--;
@@ -284,6 +337,24 @@ void split_kd(KdNode* in, u32 depth, int next_split_dim) {
     in->right->colors.push_back(in->colors[i]);
   }
 
+  /*
+  if (debug) {
+    if (in->left->colors.empty()) {
+      printf(" LEFT empty\n");
+    } else {
+      printf(" LEFT, has %ld, %d to %d\n", in->left->colors.size(),
+             in->left->colors.front()[next_split_dim], in->left->colors.back()[next_split_dim]);
+    }
+
+    if (in->right->colors.empty()) {
+      printf(" RIGHT empty\n");
+    } else {
+      printf(" RIGHT, has %ld, %d to %d\n", in->right->colors.size(),
+             in->right->colors.front()[next_split_dim], in->right->colors.back()[next_split_dim]);
+    }
+  }
+  */
+
   split_kd(in->left.get(), depth - 1, (next_split_dim + 1) % 4);
   split_kd(in->right.get(), depth - 1, (next_split_dim + 1) % 4);
 }
@@ -322,17 +393,83 @@ std::vector<Color> deduplicated_colors(const std::vector<Color>& in) {
   return out;
 }
 
+u8 saturate_to_u8(u32 in) {
+  if (in >= UINT8_MAX) {
+    return UINT8_MAX;
+  } else {
+    return in;
+  }
+}
+
+s32 color_difference(const Color& c1, const Color& c2) {
+  s32 ret = 0;
+  for (int i = 0; i < 4; i++) {
+    const int diff = (int)c1[i] - (int)c2[i];
+    ret += diff * diff;
+  }
+  return ret;
+}
+
+int total_color_count(const KdNode* node) {
+  int ret = 0;
+  if (node->left && node->right) {
+    ret += total_color_count(node->left.get());
+    ret += total_color_count(node->right.get());
+  } else {
+    if (!node->colors.empty()) {
+      ret += 1;
+    }
+  }
+  return ret;
+}
+
+void get_splittable(KdNode* node, std::vector<KdNode*>* out) {
+  if (node->left && node->right) {
+    get_splittable(node->left.get(), out);
+    get_splittable(node->right.get(), out);
+  } else {
+    if (node->colors.size() > 1 && node->colors.front() != node->colors.back()) {
+      out->push_back(node);
+    }
+  }
+}
+
 QuantizedColors quantize_colors_kd_tree(const std::vector<math::Vector<u8, 4>>& in,
                                         u32 target_depth) {
   Timer timer;
   // Build root node:
   KdNode root;
   root.colors = deduplicated_colors(in);
-  // root.colors = in;
+  const int num_unique_colors = root.colors.size();
 
   // Split tree:
   split_kd(&root, target_depth, 0);
 
+  // keep splitting!
+  int color_count = total_color_count(&root);
+  printf("color count %d / %d\n", color_count, (1 << target_depth));
+  while (color_count < (1 << target_depth)) {
+    printf("extra split iteration - have %d / %d\n", color_count, (1 << target_depth));
+    std::vector<KdNode*> extra_splits;
+    get_splittable(&root, &extra_splits);
+    printf(" found %ld splittable nodes\n", extra_splits.size());
+    if (extra_splits.empty()) {
+      break;
+    }
+
+    int i = 0;
+    while (color_count < (1 << target_depth)) {
+      if (i >= extra_splits.size()) {
+        break;
+      }
+      split_kd(extra_splits[i], 1, pick_split_dim_final_splits(extra_splits[i]->colors));
+      i++;
+      color_count++;
+    }
+
+    color_count = total_color_count(&root);
+  }
+
   // Get final colors:
   std::unordered_map<u32, u32> color_value_to_color_idx;
   QuantizedColors result;
@@ -343,22 +480,46 @@ QuantizedColors quantize_colors_kd_tree(const std::vector<math::Vector<u8, 4>>&
 
     const u32 slot = result.final_colors.size();
     u32 totals[4] = {0, 0, 0, 0};
-    u32 n = node->colors.size();
+    const u32 n = node->colors.size();
     for (auto& color : node->colors) {
       color_value_to_color_idx[color_as_u32(color)] = slot;
       for (int i = 0; i < 4; i++) {
         totals[i] += color[i];
       }
     }
-    result.final_colors.emplace_back(totals[0] / n, totals[1] / n, totals[2] / n,
-                                     totals[3] / (2 * n));
+    result.final_colors.emplace_back(saturate_to_u8(totals[0] / n), saturate_to_u8(totals[1] / n),
+                                     saturate_to_u8(totals[2] / n),
+                                     saturate_to_u8(totals[3] / (2 * n)));
   });
 
   for (auto& color : in) {
     result.vtx_to_color.push_back(color_value_to_color_idx.at(color_as_u32(color)));
   }
 
-  lg::warn("Quantize colors: {} input colors -> {} output in {:.3f} ms\n", in.size(),
-           result.final_colors.size(), timer.getMs());
+  lg::warn("Quantize colors: {} input colors ({} unique) -> {} output in {:.3f} ms\n", in.size(),
+           num_unique_colors, result.final_colors.size(), timer.getMs());
+
+  // debug:
+  if (!result.vtx_to_color.empty()) {
+    Color worst_in;
+    Color worst_out;
+    s32 worst_diff = -1;
+
+    for (size_t i = 0; i < result.vtx_to_color.size(); i++) {
+      Color input = in.at(i);
+      input.w() /= 2;
+      const Color output = result.final_colors.at(result.vtx_to_color.at(i));
+      const s32 diff = color_difference(input, output);
+      if (diff > worst_diff) {
+        worst_diff = diff;
+        worst_in = input;
+        worst_out = output;
+      }
+    }
+
+    lg::error("Worst diff {} between {} {}", std::sqrt((float)worst_diff),
+              worst_in.to_string_hex_byte(), worst_out.to_string_hex_byte());
+  }
+
   return result;
 }

goalc/build_level/common/gltf_mesh_extract.cpp

@@ -15,6 +15,7 @@
 #include <common/util/image_resize.h>
 
 using namespace gltf_util;
+constexpr int kColorTreeDepth = 13;
 namespace gltf_mesh_extract {
 
 void dedup_tfrag_vertices(TfragOutput& data) {
@@ -201,7 +202,7 @@ void extract(const Input& in,
            out.tfrag_vertices.size());
 
   Timer quantize_timer;
-  auto quantized = quantize_colors_kd_tree(all_vtx_colors, 10);
+  auto quantized = quantize_colors_kd_tree(all_vtx_colors, kColorTreeDepth);
   for (size_t i = 0; i < out.tfrag_vertices.size(); i++) {
     out.tfrag_vertices[i].color_index = quantized.vtx_to_color[i];
   }
@@ -365,7 +366,7 @@ void extract(const Input& in,
            out.vertices.size());
 
   Timer quantize_timer;
-  auto quantized = quantize_colors_kd_tree(all_vtx_colors, 10);
+  auto quantized = quantize_colors_kd_tree(all_vtx_colors, kColorTreeDepth);
   for (size_t i = 0; i < out.vertices.size(); i++) {
     out.color_indices.push_back(quantized.vtx_to_color[i]);
   }