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]); }