forked from hkc/cc-stuff
1
0
Fork 0

Implement k-means palette generation

This commit is contained in:
Vftdan 2024-10-03 01:27:30 +02:00
parent 849563285c
commit c5af524dac
1 changed files with 112 additions and 1 deletions

113
img2cpi.c
View File

@ -13,6 +13,7 @@
#include <stdbool.h> #include <stdbool.h>
#define MAX_COLOR_DIFFERENCE 768 #define MAX_COLOR_DIFFERENCE 768
#define K_MEANS_ITERATIONS 4
struct rgba { uint8_t r, g, b, a; }; struct rgba { uint8_t r, g, b, a; };
union color { union color {
@ -85,6 +86,19 @@ struct image_pal {
const struct palette *palette; 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); bool parse_cmdline(int argc, char **argv);
void show_help(const char *progname, bool show_all, FILE *fp); void show_help(const char *progname, bool show_all, FILE *fp);
struct image *image_load(const char *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_resize(uint8_t size);
struct palette *custom_palette_from(const struct palette *orig); 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[] = { const char *known_file_extensions[] = {
".png", ".jpg", ".jpeg", ".jfif", ".jpg", ".gif", ".png", ".jpg", ".jpeg", ".jfif", ".jpg", ".gif",
".tga", ".bmp", ".hdr", ".pnm", 0 ".tga", ".bmp", ".hdr", ".pnm", 0
@ -176,7 +195,7 @@ int main(int argc, char **argv) {
switch (args.palette_type) { switch (args.palette_type) {
case PALETTE_DEFAULT: palette = &DEFAULT_PALETTE; break; case PALETTE_DEFAULT: palette = &DEFAULT_PALETTE; break;
case PALETTE_DEFAULT_GRAY: palette = &DEFAULT_GRAY_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_LIST: assert(0 && "Not implemented"); break;
case PALETTE_PATH: assert(0 && "Not implemented"); break; case PALETTE_PATH: assert(0 && "Not implemented"); break;
default: assert(0 && "Unreachable"); default: assert(0 && "Unreachable");
@ -654,6 +673,98 @@ struct palette *custom_palette_from(const struct palette *orig) {
return (struct palette*)&custom_palette_data; 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( const struct palette DEFAULT_PALETTE = PALETTE(
{ { 0xf0, 0xf0, 0xf0, 0xff } }, { { 0xf0, 0xf0, 0xf0, 0xff } },
{ { 0xf2, 0xb2, 0x33, 0xff } }, { { 0xf2, 0xb2, 0x33, 0xff } },