git.sesse.net Git - alphaindex/blob - alphaindex.cpp

   1 #include <stdio.h>
   2 #include <math.h>
   3 #include <png.h>
   4 #include <SDL/SDL_image.h>
   5 #include <vector>
   6 #include <algorithm>
   7
   8 #define NUM_CLUSTERS 16
   9
  10 using namespace std;
  11
  12 struct Pixel {
  13         float r, g, b, a;
  14
  15         Pixel(float r, float g, float b, float a) : r(r), g(g), b(b), a(a) {}
  16         Pixel() {}
  17 };
  18
  19 struct CompareByRGBA {
  20         bool operator() (const Pixel &a, const Pixel &b) const {
  21                 if (a.a != b.a)
  22                         return (a.a < b.a);
  23                 if (a.r != b.r)
  24                         return (a.r < b.r);
  25                 if (a.g != b.g)
  26                         return (a.g < b.g);
  27                 return (a.b < b.b);
  28         }
  29 };
  30
  31 void write_png(const char *filename, unsigned char *screenbuf, int width, int height)
  32 {
  33         FILE *fp = fopen(filename, "wb");
  34         png_structp png_ptr = png_create_write_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);
  35         png_infop info_ptr = png_create_info_struct(png_ptr);
  36
  37         if (setjmp(png_jmpbuf(png_ptr))) {
  38                 fclose(fp);
  39                 fprintf(stderr, "Write to %s failed; exiting.\n", filename);
  40                 exit(1);
  41         }
  42
  43         png_set_IHDR(png_ptr, info_ptr, width, height, 8, PNG_COLOR_TYPE_RGB_ALPHA, PNG_INTERLACE_NONE, PNG_COMPRESSION_TYPE_DEFAULT, PNG_FILTER_TYPE_DEFAULT);
  44
  45         png_bytep *row_pointers = new png_bytep[height];
  46         for (int y = 0; y < height; ++y) {
  47                 row_pointers[y] = screenbuf + (y * width) * 4;
  48         }
  49
  50         png_init_io(png_ptr, fp);
  51         png_set_rows(png_ptr, info_ptr, row_pointers);
  52         png_write_png(png_ptr, info_ptr, PNG_TRANSFORM_BGR, NULL);
  53         png_destroy_write_struct(&png_ptr, &info_ptr);
  54         fclose(fp);
  55
  56         delete[] row_pointers;
  57 }
  58
  59 // [0.0, 1.0)
  60 float uniform_rand()
  61 {
  62         return rand() / float(RAND_MAX + 1.0);
  63 }
  64
  65 Pixel random_pixel()
  66 {
  67         float r = uniform_rand();
  68         float g = uniform_rand();
  69         float b = uniform_rand();
  70         float a = uniform_rand();
  71
  72         r *= a;
  73         g *= a;
  74         b *= a;
  75
  76         return Pixel(r, g, b, a);
  77 }
  78
  79 int assign_to_cluster(const Pixel &pixel, const Pixel *center)
  80 {
  81         int best_cluster = -1;
  82         float best_dist = HUGE_VAL;
  83
  84         for (int k = 0; k < NUM_CLUSTERS; ++k) {
  85                 // TODO: Better perceptual distance than just nonlinear RGB space.
  86                 float dist =
  87                         (pixel.r - center[k].r) * (pixel.r - center[k].r) +
  88                         (pixel.g - center[k].g) * (pixel.g - center[k].g) +
  89                         (pixel.b - center[k].b) * (pixel.b - center[k].b) +
  90                         3.0 * (pixel.a - center[k].a) * (pixel.a - center[k].a);
  91                 if (dist < best_dist) {
  92                         best_dist = dist;
  93                         best_cluster = k;
  94                 }
  95         }
  96
  97         return best_cluster;
  98 }
  99
 100 int main(void)
 101 {
 102         SDL_Surface *foo = IMG_Load("overlay-shadow.png");
 103
 104         vector<Pixel> pixels;
 105         for (int y = 0; y < foo->h; ++y) {
 106                 uint8_t *ptr = ((uint8_t *)foo->pixels + foo->pitch * y);;
 107                 for (int x = 0; x < foo->w; ++x) {
 108                         float r = *ptr++ * (1.0/255.0);
 109                         float g = *ptr++ * (1.0/255.0);
 110                         float b = *ptr++ * (1.0/255.0);
 111                         float a = *ptr++ * (1.0/255.0);
 112
 113                         // Premultiply alpha.
 114                         r *= a;
 115                         g *= a;
 116                         b *= a;
 117
 118                         pixels.push_back(Pixel(r, g, b, a));
 119                 }
 120         }
 121
 122         // K-means: Initialization
 123         Pixel center[NUM_CLUSTERS];
 124         for (int i = 0; i < NUM_CLUSTERS; ++i) {
 125                 center[i] = random_pixel();
 126         }
 127
 128         // K-means: Iteration
 129         for (int i = 0; i < 1000; ++i) {
 130                 Pixel new_center[NUM_CLUSTERS];
 131                 int num_center[NUM_CLUSTERS];
 132
 133                 sort(center, center + NUM_CLUSTERS, CompareByRGBA());
 134
 135                 for (int j = 0; j < NUM_CLUSTERS; ++j) {
 136                         printf("%2d: %.0f, %.0f, %.0f, %.0f\n", j,
 137                                 255.0 * center[j].r,
 138                                 255.0 * center[j].g,
 139                                 255.0 * center[j].b,
 140                                 255.0 * center[j].a);
 141                 }
 142                 printf("\n");
 143
 144                 for (int j = 0; j < NUM_CLUSTERS; ++j) {
 145                         new_center[j] = Pixel(0, 0, 0, 0);
 146                         num_center[j] = 0;
 147                 }
 148
 149                 for (unsigned j = 0; j < pixels.size(); ++j) {
 150                         int best_cluster = assign_to_cluster(pixels[j], center);
 151
 152                         ++num_center[best_cluster];
 153                         new_center[best_cluster].r += pixels[j].r;
 154                         new_center[best_cluster].g += pixels[j].g;
 155                         new_center[best_cluster].b += pixels[j].b;
 156                         new_center[best_cluster].a += pixels[j].a;
 157                 }
 158
 159                 for (int j = 0; j < NUM_CLUSTERS; ++j) {
 160                         if (num_center[j] == 0) {
 161                                 center[j] = random_pixel();
 162                         } else {
 163                                 center[j].r = new_center[j].r / num_center[j];
 164                                 center[j].g = new_center[j].g / num_center[j];
 165                                 center[j].b = new_center[j].b / num_center[j];
 166                                 center[j].a = new_center[j].a / num_center[j];
 167                         }
 168                 }
 169         }
 170
 171         // Write PNG.
 172         unsigned char *out = new unsigned char[foo->w * foo->h * 4];
 173         unsigned char *ptr = out;
 174         for (unsigned i = 0; i < pixels.size(); ++i) {
 175                 Pixel q = pixels[i];
 176
 177 #if 0
 178                 // TPDF dither. (TODO: How do you do dither with non-uniform quantization?)
 179                 q.r += (uniform_rand() + uniform_rand()) * (1.0 / 8.0);
 180                 q.g += (uniform_rand() + uniform_rand()) * (1.0 / 8.0);
 181                 q.b += (uniform_rand() + uniform_rand()) * (1.0 / 8.0);
 182
 183                 printf("POST: %f %f %f\n", q.r, q.g, q.b);
 184 #endif
 185
 186                 const Pixel &p = center[assign_to_cluster(q, center)];
 187                 if (p.a < 1e-6) {
 188                         *ptr++ = 0;
 189                         *ptr++ = 0;
 190                         *ptr++ = 0;
 191                         *ptr++ = 0;
 192                 } else {
 193                         *ptr++ = lrintf((p.r/p.a) * 255.0);
 194                         *ptr++ = lrintf((p.g/p.a) * 255.0);
 195                         *ptr++ = lrintf((p.b/p.a) * 255.0);
 196                         *ptr++ = lrintf(p.a * 255.0);
 197                 }
 198         }
 199         write_png("indexed.png", out, foo->w, foo->h);
 200
 201         // Write indexed stuff.
 202         FILE *fp = fopen("logo.indexed", "wb");
 203         for (unsigned i = 0; i < NUM_CLUSTERS; ++i) {
 204                 const Pixel& p = center[i];
 205                 fputc(lrintf((p.r/p.a) * 255.0), fp);
 206                 fputc(lrintf((p.g/p.a) * 255.0), fp);
 207                 fputc(lrintf((p.b/p.a) * 255.0), fp);
 208                 fputc(lrintf(p.a * 255.0), fp);
 209         }
 210         for (unsigned i = 0; i < pixels.size(); ++i) {
 211                 fputc(assign_to_cluster(pixels[i], center), fp);
 212         }
 213         fclose(fp);
 214 }