From 5458d865d4b6258b579e72062030fd4accdafc62 Mon Sep 17 00:00:00 2001 From: water111 <48171810+water111@users.noreply.github.com> Date: Sun, 8 Dec 2024 12:31:48 -0500 Subject: [PATCH] [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 --- .../build_level/common/color_quantization.cpp | 175 +++++++++++++++++- .../build_level/common/gltf_mesh_extract.cpp | 5 +- 2 files changed, 171 insertions(+), 9 deletions(-) diff --git a/goalc/build_level/common/color_quantization.cpp b/goalc/build_level/common/color_quantization.cpp index ccedb3b88ad..e4750281cb4 100644 --- a/goalc/build_level/common/color_quantization.cpp +++ b/goalc/build_level/common/color_quantization.cpp @@ -1,6 +1,7 @@ #include "color_quantization.h" #include +#include #include #include @@ -232,6 +233,57 @@ struct KdNode { std::vector colors; }; +size_t pick_split_point(const std::vector& colors, int dim) { + if (colors.empty()) { + return 0; + } + double mean = 0; + int mean_count = 0; + std::array 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& 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(); 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 deduplicated_colors(const std::vector& 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* 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>& 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 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 color_value_to_color_idx; QuantizedColors result; @@ -343,22 +480,46 @@ QuantizedColors quantize_colors_kd_tree(const std::vector>& 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; } diff --git a/goalc/build_level/common/gltf_mesh_extract.cpp b/goalc/build_level/common/gltf_mesh_extract.cpp index 3d84f74fda2..69124f1ec87 100644 --- a/goalc/build_level/common/gltf_mesh_extract.cpp +++ b/goalc/build_level/common/gltf_mesh_extract.cpp @@ -15,6 +15,7 @@ #include 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]); }