git.mutantstargoat.com Git - gbajam22/blob - tools/pngdump/quant.c

   1 #include <stdio.h>
   2 #include <stdlib.h>
   3 #include <string.h>
   4 #include <limits.h>
   5 #include <errno.h>
   6 #include <assert.h>
   7 #include "image.h"
   8
   9 #define NUM_LEVELS      8
  10
  11 struct octnode;
  12
  13 struct octree {
  14         struct octnode *root;
  15         struct octnode *redlist[NUM_LEVELS];
  16         int redlev;
  17         int nleaves, maxcol;
  18 };
  19
  20 struct octnode {
  21         int lvl;
  22         struct octree *tree;
  23         int r, g, b, nref;
  24         int palidx;
  25         int nsub, leaf;
  26         struct octnode *sub[8];
  27         struct octnode *next;
  28 };
  29
  30
  31 static void init_octree(struct octree *tree, int maxcol);
  32 static void destroy_octree(struct octree *tree);
  33
  34 static struct octnode *alloc_node(struct octree *tree, int lvl);
  35 static void free_node(struct octnode *n);
  36 static void free_tree(struct octnode *n);
  37
  38 static void add_color(struct octree *tree, int r, int g, int b, int nref);
  39 static void reduce_colors(struct octree *tree);
  40 static int assign_colors(struct octnode *n, int next, struct cmapent *cmap);
  41 static int lookup_color(struct octree *tree, int r, int g, int b);
  42 static int subidx(int bit, int r, int g, int b);
  43 static void print_tree(struct octnode *n, int lvl);
  44
  45 int quantize_image(struct image *img, int maxcol)
  46 {
  47         int i, j, cidx;
  48         unsigned int rgb[3];
  49         struct octree tree;
  50         struct image newimg = *img;
  51
  52         if(maxcol < 2 || maxcol > 256) {
  53                 return -1;
  54         }
  55
  56         if(img->bpp > 8) {
  57                 newimg.bpp = maxcol > 16 ? 8 : 4;
  58                 newimg.nchan = 1;
  59                 newimg.scansz = newimg.width * newimg.bpp / 8;
  60                 newimg.pitch = newimg.scansz;
  61         }
  62
  63         init_octree(&tree, maxcol);
  64
  65         for(i=0; i<img->height; i++) {
  66                 for(j=0; j<img->width; j++) {
  67                         get_pixel_rgb(img, j, i, rgb);
  68                         add_color(&tree, rgb[0], rgb[1], rgb[2], 1);
  69
  70                         while(tree.nleaves > maxcol) {
  71                                 reduce_colors(&tree);
  72                         }
  73                 }
  74         }
  75
  76         /* use created octree to generate the palette */
  77         newimg.cmap_ncolors = assign_colors(tree.root, 0, newimg.cmap);
  78
  79         /* replace image pixels */
  80         for(i=0; i<img->height; i++) {
  81                 for(j=0; j<img->width; j++) {
  82                         get_pixel_rgb(img, j, i, rgb);
  83                         cidx = lookup_color(&tree, rgb[0], rgb[1], rgb[2]);
  84                         assert(cidx >= 0 && cidx < maxcol);
  85                         put_pixel(&newimg, j, i, cidx);
  86                 }
  87         }
  88
  89         *img = newimg;
  90
  91         destroy_octree(&tree);
  92         return 0;
  93 }
  94
  95 int gen_shades(struct image *img, int levels, int maxcol, int *shade_lut)
  96 {
  97         int i, j, cidx, r, g, b;
  98         unsigned int color[3];
  99         struct octree tree;
 100         struct image newimg = *img;
 101
 102         if(maxcol < 2 || maxcol > 256) {
 103                 return -1;
 104         }
 105
 106         init_octree(&tree, maxcol);
 107
 108         for(i=0; i<img->cmap_ncolors; i++) {
 109                 add_color(&tree, img->cmap[i].r, img->cmap[i].g, img->cmap[i].b, 1024);
 110         }
 111
 112         for(i=0; i<img->cmap_ncolors; i++) {
 113                 for(j=0; j<levels - 1; j++) {
 114                         r = img->cmap[i].r * j / (levels - 1);
 115                         g = img->cmap[i].g * j / (levels - 1);
 116                         b = img->cmap[i].b * j / (levels - 1);
 117                         add_color(&tree, r, g, b, 1);
 118
 119                         while(tree.nleaves > maxcol) {
 120                                 reduce_colors(&tree);
 121                         }
 122                 }
 123         }
 124
 125         newimg.cmap_ncolors = assign_colors(tree.root, 0, newimg.cmap);
 126
 127         /* replace image pixels */
 128         for(i=0; i<img->height; i++) {
 129                 for(j=0; j<img->width; j++) {
 130                         get_pixel_rgb(img, j, i, color);
 131                         cidx = lookup_color(&tree, color[0], color[1], color[2]);
 132                         put_pixel(&newimg, j, i, cidx);
 133                 }
 134         }
 135
 136         /* populate shade_lut based on the new palette, can't generate levels only
 137          * for the original colors, because the palette entries will have changed
 138          * and moved around.
 139          */
 140         for(i=0; i<newimg.cmap_ncolors; i++) {
 141                 for(j=0; j<levels; j++) {
 142                         r = newimg.cmap[i].r * j / (levels - 1);
 143                         g = newimg.cmap[i].g * j / (levels - 1);
 144                         b = newimg.cmap[i].b * j / (levels - 1);
 145                         *shade_lut++ = lookup_color(&tree, r, g, b);
 146                 }
 147         }
 148         for(i=0; i<(maxcol - newimg.cmap_ncolors) * levels; i++) {
 149                 *shade_lut++ = maxcol - 1;
 150         }
 151
 152         *img = newimg;
 153
 154         destroy_octree(&tree);
 155         return 0;
 156 }
 157
 158 static void init_octree(struct octree *tree, int maxcol)
 159 {
 160         memset(tree, 0, sizeof *tree);
 161
 162         tree->redlev = NUM_LEVELS - 1;
 163         tree->maxcol = maxcol;
 164
 165         tree->root = alloc_node(tree, 0);
 166 }
 167
 168 static void destroy_octree(struct octree *tree)
 169 {
 170         free_tree(tree->root);
 171 }
 172
 173 static struct octnode *alloc_node(struct octree *tree, int lvl)
 174 {
 175         struct octnode *n;
 176
 177         if(!(n = calloc(1, sizeof *n))) {
 178                 perror("failed to allocate octree node");
 179                 abort();
 180         }
 181
 182         n->lvl = lvl;
 183         n->tree = tree;
 184         n->palidx = -1;
 185
 186         if(lvl < tree->redlev) {
 187                 n->next = tree->redlist[lvl];
 188                 tree->redlist[lvl] = n;
 189         } else {
 190                 n->leaf = 1;
 191                 tree->nleaves++;
 192         }
 193         return n;
 194 }
 195
 196 static void free_node(struct octnode *n)
 197 {
 198         struct octnode *prev, dummy;
 199
 200         dummy.next = n->tree->redlist[n->lvl];
 201         prev = &dummy;
 202         while(prev->next) {
 203                 if(prev->next == n) {
 204                         prev->next = n->next;
 205                         break;
 206                 }
 207                 prev = prev->next;
 208         }
 209         n->tree->redlist[n->lvl] = dummy.next;
 210
 211         if(n->leaf) {
 212                 n->tree->nleaves--;
 213                 assert(n->tree->nleaves >= 0);
 214         }
 215         free(n);
 216 }
 217
 218 static void free_tree(struct octnode *n)
 219 {
 220         int i;
 221
 222         if(!n) return;
 223
 224         for(i=0; i<8; i++) {
 225                 free_tree(n->sub[i]);
 226         }
 227         free_node(n);
 228 }
 229
 230 static void add_color(struct octree *tree, int r, int g, int b, int nref)
 231 {
 232         int i, idx, rr, gg, bb;
 233         struct octnode *n;
 234
 235         rr = r * nref;
 236         gg = g * nref;
 237         bb = b * nref;
 238
 239         n = tree->root;
 240         n->r += rr;
 241         n->g += gg;
 242         n->b += bb;
 243         n->nref += nref;
 244
 245         for(i=0; i<NUM_LEVELS; i++) {
 246                 if(n->leaf) break;
 247
 248                 idx = subidx(i, r, g, b);
 249
 250                 if(!n->sub[idx]) {
 251                         n->sub[idx] = alloc_node(tree, i + 1);
 252                 }
 253                 n->nsub++;
 254                 n = n->sub[idx];
 255
 256                 n->r += rr;
 257                 n->g += gg;
 258                 n->b += bb;
 259                 n->nref += nref;
 260         }
 261 }
 262
 263 static struct octnode *get_reducible(struct octree *tree)
 264 {
 265         int best_nref;
 266         struct octnode dummy, *n, *prev, *best_prev, *best = 0;
 267
 268         while(tree->redlev >= 0) {
 269                 best_nref = INT_MAX;
 270                 best = 0;
 271                 dummy.next = tree->redlist[tree->redlev];
 272                 prev = &dummy;
 273                 while(prev->next) {
 274                         n = prev->next;
 275                         if(n->nref < best_nref) {
 276                                 best = n;
 277                                 best_nref = n->nref;
 278                                 best_prev = prev;
 279                         }
 280                         prev = prev->next;
 281                 }
 282                 if(best) {
 283                         best_prev->next = best->next;
 284                         tree->redlist[tree->redlev] = dummy.next;
 285                         break;
 286                 }
 287                 tree->redlev--;
 288         }
 289
 290         return best;
 291 }
 292
 293 static void reduce_colors(struct octree *tree)
 294 {
 295         int i;
 296         struct octnode *n;
 297
 298         if(!(n = get_reducible(tree))) {
 299                 fprintf(stderr, "warning: no reducible nodes!\n");
 300                 return;
 301         }
 302         for(i=0; i<8; i++) {
 303                 if(n->sub[i]) {
 304                         free_node(n->sub[i]);
 305                         n->sub[i] = 0;
 306                 }
 307         }
 308         n->leaf = 1;
 309         tree->nleaves++;
 310 }
 311
 312 static int assign_colors(struct octnode *n, int next, struct cmapent *cmap)
 313 {
 314         int i;
 315
 316         if(!n) return next;
 317
 318         if(n->leaf) {
 319                 assert(next < n->tree->maxcol);
 320                 assert(n->nref);
 321                 cmap[next].r = n->r / n->nref;
 322                 cmap[next].g = n->g / n->nref;
 323                 cmap[next].b = n->b / n->nref;
 324                 n->palidx = next;
 325                 return next + 1;
 326         }
 327
 328         for(i=0; i<8; i++) {
 329                 next = assign_colors(n->sub[i], next, cmap);
 330         }
 331         return next;
 332 }
 333
 334 static int lookup_color(struct octree *tree, int r, int g, int b)
 335 {
 336         int i, j, idx;
 337         struct octnode *n;
 338
 339         n = tree->root;
 340         for(i=0; i<NUM_LEVELS; i++) {
 341                 if(n->leaf) {
 342                         assert(n->palidx >= 0);
 343                         return n->palidx;
 344                 }
 345
 346                 idx = subidx(i, r, g, b);
 347                 for(j=0; j<8; j++) {
 348                         if(n->sub[idx]) break;
 349                         idx = (idx + 1) & 7;
 350                 }
 351
 352                 assert(n->sub[idx]);
 353                 n = n->sub[idx];
 354         }
 355
 356         fprintf(stderr, "lookup_color(%d, %d, %d) failed!\n", r, g, b);
 357         abort();
 358         return -1;
 359 }
 360
 361 static int subidx(int bit, int r, int g, int b)
 362 {
 363         assert(bit >= 0 && bit < NUM_LEVELS);
 364         bit = NUM_LEVELS - 1 - bit;
 365         return ((r >> bit) & 1) | ((g >> (bit - 1)) & 2) | ((b >> (bit - 2)) & 4);
 366 }
 367
 368 static void print_tree(struct octnode *n, int lvl)
 369 {
 370         int i;
 371         char ptrbuf[32], *p;
 372
 373         if(!n) return;
 374
 375         for(i=0; i<lvl; i++) {
 376                 fputs("|  ", stdout);
 377         }
 378
 379         sprintf(ptrbuf, "%p", (void*)n);
 380         p = ptrbuf + strlen(ptrbuf) - 4;
 381
 382         if(n->nref) {
 383                 printf("+-(%d) %s: <%d %d %d> #%d", n->lvl, p, n->r / n->nref, n->g / n->nref,
 384                                 n->b / n->nref, n->nref);
 385         } else {
 386                 printf("+-(%d) %s: <- - -> #0", n->lvl, p);
 387         }
 388         if(n->palidx >= 0) printf(" [%d]", n->palidx);
 389         if(n->leaf) printf(" LEAF");
 390         putchar('\n');
 391
 392         for(i=0; i<8; i++) {
 393                 print_tree(n->sub[i], lvl + 1);
 394         }
 395 }
 396