From c5af524dac5839e6db6bc2a92576d7f3b43e42c4 Mon Sep 17 00:00:00 2001
From: Vftdan <vftdan@gmail.com>
Date: Thu, 3 Oct 2024 01:27:30 +0200
Subject: [PATCH] Implement k-means palette generation

---
 img2cpi.c | 113 +++++++++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 112 insertions(+), 1 deletion(-)

diff --git a/img2cpi.c b/img2cpi.c
index ee81cf3..be9c3fa 100644
--- a/img2cpi.c
+++ b/img2cpi.c
@@ -13,6 +13,7 @@
 #include <stdbool.h>
 
 #define MAX_COLOR_DIFFERENCE 768
+#define K_MEANS_ITERATIONS 4
 
 struct rgba { uint8_t r, g, b, a; };
 union color {
@@ -85,6 +86,19 @@ struct image_pal {
   const struct palette *palette;
 };
 
+struct k_means_state {
+  const struct image *items;
+  struct palette *clusters;
+  uint8_t *predicted_cluster;
+  struct k_means_centroid_intermediate {
+    struct {
+      float r, g, b;
+    } sums;
+    size_t count;
+  } *centroid_intermediate;
+  size_t item_count;
+};
+
 bool parse_cmdline(int argc, char **argv);
 void show_help(const char *progname, bool show_all, FILE *fp);
 struct image *image_load(const char *fp);
@@ -103,6 +117,11 @@ void convert_8x11(const struct image_pal *img, struct cc_char *characters);
 struct palette *custom_palette_resize(uint8_t size);
 struct palette *custom_palette_from(const struct palette *orig);
 
+struct k_means_state k_means_init(const struct image *image, struct palette *starting_palette);
+bool k_means_iteration(struct k_means_state *state);
+void k_means_end(struct k_means_state *state);
+struct palette *palette_k_means(const struct image *image, const struct palette *prototype);
+
 const char *known_file_extensions[] = {
   ".png", ".jpg", ".jpeg", ".jfif", ".jpg", ".gif",
   ".tga", ".bmp", ".hdr", ".pnm", 0
@@ -176,7 +195,7 @@ int main(int argc, char **argv) {
   switch (args.palette_type) {
     case PALETTE_DEFAULT: palette = &DEFAULT_PALETTE; break;
     case PALETTE_DEFAULT_GRAY: palette = &DEFAULT_GRAY_PALETTE; break;
-    case PALETTE_AUTO: assert(0 && "Not implemented"); break;
+    case PALETTE_AUTO: palette = palette_k_means(src_image, &DEFAULT_PALETTE); break;
     case PALETTE_LIST: assert(0 && "Not implemented"); break;
     case PALETTE_PATH: assert(0 && "Not implemented"); break;
     default: assert(0 && "Unreachable");
@@ -654,6 +673,98 @@ struct palette *custom_palette_from(const struct palette *orig) {
   return (struct palette*)&custom_palette_data;
 }
 
+struct k_means_state k_means_init(const struct image *image, struct palette *starting_palette) {
+  size_t item_count = image->w * image->h;
+  struct k_means_state state = {
+    .items = image,
+    .clusters = starting_palette,
+    .predicted_cluster = calloc(image->w, image->h),
+    .centroid_intermediate = calloc(item_count, sizeof(struct k_means_centroid_intermediate)),
+    .item_count = item_count,
+  };
+  return state;
+}
+
+bool k_means_iteration(struct k_means_state *state) {
+  if (!state->predicted_cluster || !state->centroid_intermediate) {
+    return false;
+  }
+  bool changed = false;
+
+  // Find closest cluster
+  for (int i = 0; i < state->item_count; i++) {
+    int closest_cluster = 0;
+    float closest_distance = 1e20;
+    for (int cluster = 0; cluster < state->clusters->count; cluster++) {
+      float dist = get_color_difference(state->clusters->colors[cluster], state->items->pixels[i]);
+      if (dist <= closest_distance) {
+        closest_distance = dist;
+        closest_cluster = cluster;
+      }
+    }
+    if (!changed) {
+      changed = state->predicted_cluster[i] != closest_cluster;
+    }
+    state->predicted_cluster[i] = closest_cluster;
+    state->centroid_intermediate[closest_cluster].count += 1;
+    state->centroid_intermediate[closest_cluster].sums.r += state->items->pixels[i].rgba.r;
+    state->centroid_intermediate[closest_cluster].sums.g += state->items->pixels[i].rgba.g;
+    state->centroid_intermediate[closest_cluster].sums.b += state->items->pixels[i].rgba.b;
+  }
+
+  // Update centroids
+  for (int i = 0; i < state->clusters->count; ++i) {
+    struct k_means_centroid_intermediate intermediate = state->centroid_intermediate[i];
+    if (intermediate.count) {
+      union color centroid = {
+        .rgba = {
+          .r = intermediate.sums.r / intermediate.count,
+          .g = intermediate.sums.g / intermediate.count,
+          .b = intermediate.sums.b / intermediate.count,
+          .a = 0xff,
+        }
+      };
+      if (!changed) {
+        changed = state->clusters->colors[i].v != centroid.v;
+      }
+      state->clusters->colors[i] = centroid;
+    } else {
+      // No pixels are closest to this color
+      // TODO: wiggle the centroid?
+    }
+    state->centroid_intermediate[i] = (struct k_means_centroid_intermediate) { .sums = {0, 0, 0}, .count = 0 };
+  }
+
+  return changed;
+}
+
+void k_means_end(struct k_means_state *state) {
+  if (state->predicted_cluster) {
+    free(state->predicted_cluster);
+  }
+  if (state->centroid_intermediate) {
+    free(state->centroid_intermediate);
+  }
+}
+
+struct palette *palette_k_means(const struct image *image, const struct palette *prototype) {
+  if (!prototype) {
+    prototype = &DEFAULT_PALETTE;
+  }
+  struct palette *palette = custom_palette_from(prototype);
+
+  struct k_means_state state = k_means_init(image, palette);
+  for (int i = 0; i < K_MEANS_ITERATIONS; i++) {
+    if (!k_means_iteration(&state)) {
+      fprintf(stderr, "early k-means stop at iteration %d\n", i);
+      break;
+    }
+  }
+  k_means_end(&state);
+
+  return palette;
+}
+
 const struct palette DEFAULT_PALETTE = PALETTE(
   { { 0xf0, 0xf0, 0xf0, 0xff } },
   { { 0xf2, 0xb2, 0x33, 0xff } },