Pillow

1
/*
2
 * The Python Imaging Library
3
 * $Id$
4
 *
5
 * image quantizer
6
 *
7
 * history:
8
 * 1998-09-10 tjs  Contributed
9
 * 1998-12-29 fl   Added to PIL 1.0b1
10
 * 2004-02-21 fl   Fixed bogus free() on quantization error
11
 * 2005-02-07 fl   Limit number of colors to 256
12
 *
13
 * Written by Toby J Sargeant <tjs@longford.cs.monash.edu.au>.
14
 *
15
 * Copyright (c) 1998 by Toby J Sargeant
16
 * Copyright (c) 1998-2004 by Secret Labs AB.  All rights reserved.
17
 *
18
 * See the README file for information on usage and redistribution.
19
 */
20

21
#include "Imaging.h"
22

23
#include <stdio.h>
24
#include <stdlib.h>
25
#include <memory.h>
26
#include <time.h>
27

28
#include "QuantTypes.h"
29
#include "QuantOctree.h"
30
#include "QuantPngQuant.h"
31
#include "QuantHash.h"
32
#include "QuantHeap.h"
33

34
/* MSVC9.0 */
35
#ifndef UINT32_MAX
36
#define UINT32_MAX 0xffffffff
37
#endif
38

39
// #define DEBUG
40
// #define TEST_NEAREST_NEIGHBOUR
41

42
typedef struct {
43
    uint32_t scale;
44
} PixelHashData;
45

46
typedef struct _PixelList {
47
    struct _PixelList *next[3], *prev[3];
48
    Pixel p;
49
    unsigned int flag : 1;
50
    int count;
51
} PixelList;
52

53
typedef struct _BoxNode {
54
    struct _BoxNode *l, *r;
55
    PixelList *head[3], *tail[3];
56
    int axis;
57
    int volume;
58
    uint32_t pixelCount;
59
} BoxNode;
60

61
#define _SQR(x) ((x) * (x))
62
#define _DISTSQR(p1, p2)                            \
63
    _SQR((int)((p1)->c.r) - (int)((p2)->c.r)) +     \
64
        _SQR((int)((p1)->c.g) - (int)((p2)->c.g)) + \
65
        _SQR((int)((p1)->c.b) - (int)((p2)->c.b))
66

67
#define MAX_HASH_ENTRIES 65536
68

69
#define PIXEL_HASH(r, g, b)                                         \
70
    (((unsigned int)(r)) * 463 ^ ((unsigned int)(g) << 8) * 10069 ^ \
71
     ((unsigned int)(b) << 16) * 64997)
72

73
#define PIXEL_UNSCALE(p, q, s)                                  \
74
    ((q)->c.r = (p)->c.r << (s)), ((q)->c.g = (p)->c.g << (s)), \
75
        ((q)->c.b = (p)->c.b << (s))
76

77
#define PIXEL_SCALE(p, q, s)                                    \
78
    ((q)->c.r = (p)->c.r >> (s)), ((q)->c.g = (p)->c.g >> (s)), \
79
        ((q)->c.b = (p)->c.b >> (s))
80

81
static uint32_t
82
unshifted_pixel_hash(const HashTable *h, const Pixel pixel) {
83
    return PIXEL_HASH(pixel.c.r, pixel.c.g, pixel.c.b);
84
}
85

86
static int
87
unshifted_pixel_cmp(const HashTable *h, const Pixel pixel1, const Pixel pixel2) {
88
    if (pixel1.c.r == pixel2.c.r) {
89
        if (pixel1.c.g == pixel2.c.g) {
90
            if (pixel1.c.b == pixel2.c.b) {
91
                return 0;
92
            } else {
93
                return (int)(pixel1.c.b) - (int)(pixel2.c.b);
94
            }
95
        } else {
96
            return (int)(pixel1.c.g) - (int)(pixel2.c.g);
97
        }
98
    } else {
99
        return (int)(pixel1.c.r) - (int)(pixel2.c.r);
100
    }
101
}
102

103
static uint32_t
104
pixel_hash(const HashTable *h, const Pixel pixel) {
105
    PixelHashData *d = (PixelHashData *)hashtable_get_user_data(h);
106
    return PIXEL_HASH(
107
        pixel.c.r >> d->scale, pixel.c.g >> d->scale, pixel.c.b >> d->scale
108
    );
109
}
110

111
static int
112
pixel_cmp(const HashTable *h, const Pixel pixel1, const Pixel pixel2) {
113
    PixelHashData *d = (PixelHashData *)hashtable_get_user_data(h);
114
    uint32_t A, B;
115
    A = PIXEL_HASH(
116
        pixel1.c.r >> d->scale, pixel1.c.g >> d->scale, pixel1.c.b >> d->scale
117
    );
118
    B = PIXEL_HASH(
119
        pixel2.c.r >> d->scale, pixel2.c.g >> d->scale, pixel2.c.b >> d->scale
120
    );
121
    return (A == B) ? 0 : ((A < B) ? -1 : 1);
122
}
123

124
static void
125
exists_count_func(const HashTable *h, const Pixel key, uint32_t *val) {
126
    *val += 1;
127
}
128

129
static void
130
new_count_func(const HashTable *h, const Pixel key, uint32_t *val) {
131
    *val = 1;
132
}
133

134
static void
135
rehash_collide(
136
    const HashTable *h, Pixel *keyp, uint32_t *valp, Pixel newkey, uint32_t newval
137
) {
138
    *valp += newval;
139
}
140

141
/* %% */
142

143
static HashTable *
144
create_pixel_hash(Pixel *pixelData, uint32_t nPixels) {
145
    PixelHashData *d;
146
    HashTable *hash;
147
    uint32_t i;
148
#ifdef DEBUG
149
    uint32_t timer, timer2, timer3;
150
#endif
151

152
    /* malloc check ok, small constant allocation */
153
    d = malloc(sizeof(PixelHashData));
154
    if (!d) {
155
        return NULL;
156
    }
157
    hash = hashtable_new(pixel_hash, pixel_cmp);
158
    hashtable_set_user_data(hash, d);
159
    d->scale = 0;
160
#ifdef DEBUG
161
    timer = timer3 = clock();
162
#endif
163
    for (i = 0; i < nPixels; i++) {
164
        if (!hashtable_insert_or_update_computed(
165
                hash, pixelData[i], new_count_func, exists_count_func
166
            )) {
167
            ;
168
        }
169
        while (hashtable_get_count(hash) > MAX_HASH_ENTRIES) {
170
            d->scale++;
171
#ifdef DEBUG
172
            printf("rehashing - new scale: %d\n", (int)d->scale);
173
            timer2 = clock();
174
#endif
175
            hashtable_rehash_compute(hash, rehash_collide);
176
#ifdef DEBUG
177
            timer2 = clock() - timer2;
178
            printf("rehash took %f sec\n", timer2 / (double)CLOCKS_PER_SEC);
179
            timer += timer2;
180
#endif
181
        }
182
    }
183
#ifdef DEBUG
184
    printf("inserts took %f sec\n", (clock() - timer) / (double)CLOCKS_PER_SEC);
185
#endif
186
#ifdef DEBUG
187
    printf("total        %f sec\n", (clock() - timer3) / (double)CLOCKS_PER_SEC);
188
#endif
189
    return hash;
190
}
191

192
static void
193
destroy_pixel_hash(HashTable *hash) {
194
    PixelHashData *d = (PixelHashData *)hashtable_get_user_data(hash);
195
    if (d) {
196
        free(d);
197
    }
198
    hashtable_free(hash);
199
}
200

201
/* 1. hash quantized pixels.                                                  */
202
/* 2. create R,G,B lists of sorted quantized pixels.                          */
203
/* 3. median cut.                                                             */
204
/* 4. build hash table from median cut boxes.                                 */
205
/* 5. for each pixel, compute entry in hash table, and hence median cut box.  */
206
/* 6. compute median cut box pixel averages.                                  */
207
/* 7. map each pixel to nearest average.                                      */
208

209
static int
210
compute_box_volume(BoxNode *b) {
211
    unsigned char rl, rh, gl, gh, bl, bh;
212
    if (b->volume >= 0) {
213
        return b->volume;
214
    }
215
    if (!b->head[0]) {
216
        b->volume = 0;
217
    } else {
218
        rh = b->head[0]->p.c.r;
219
        rl = b->tail[0]->p.c.r;
220
        gh = b->head[1]->p.c.g;
221
        gl = b->tail[1]->p.c.g;
222
        bh = b->head[2]->p.c.b;
223
        bl = b->tail[2]->p.c.b;
224
        b->volume = (rh - rl + 1) * (gh - gl + 1) * (bh - bl + 1);
225
    }
226
    return b->volume;
227
}
228

229
static void
230
hash_to_list(const HashTable *h, const Pixel pixel, const uint32_t count, void *u) {
231
    PixelHashData *d = (PixelHashData *)hashtable_get_user_data(h);
232
    PixelList **pl = (PixelList **)u;
233
    PixelList *p;
234
    int i;
235
    Pixel q;
236

237
    PIXEL_SCALE(&pixel, &q, d->scale);
238

239
    /* malloc check ok, small constant allocation */
240
    p = malloc(sizeof(PixelList));
241
    if (!p) {
242
        return;
243
    }
244

245
    p->flag = 0;
246
    p->p = q;
247
    p->count = count;
248
    for (i = 0; i < 3; i++) {
249
        p->next[i] = pl[i];
250
        p->prev[i] = NULL;
251
        if (pl[i]) {
252
            pl[i]->prev[i] = p;
253
        }
254
        pl[i] = p;
255
    }
256
}
257

258
static PixelList *
259
mergesort_pixels(PixelList *head, int i) {
260
    PixelList *c, *t, *a, *b, *p;
261
    if (!head || !head->next[i]) {
262
        if (head) {
263
            head->next[i] = NULL;
264
            head->prev[i] = NULL;
265
        }
266
        return head;
267
    }
268
    for (c = t = head; c && t;
269
         c = c->next[i], t = (t->next[i]) ? t->next[i]->next[i] : NULL);
270
    if (c) {
271
        if (c->prev[i]) {
272
            c->prev[i]->next[i] = NULL;
273
        }
274
        c->prev[i] = NULL;
275
    }
276
    a = mergesort_pixels(head, i);
277
    b = mergesort_pixels(c, i);
278
    head = NULL;
279
    p = NULL;
280
    while (a && b) {
281
        if (a->p.a.v[i] > b->p.a.v[i]) {
282
            c = a;
283
            a = a->next[i];
284
        } else {
285
            c = b;
286
            b = b->next[i];
287
        }
288
        c->prev[i] = p;
289
        c->next[i] = NULL;
290
        if (p) {
291
            p->next[i] = c;
292
        }
293
        p = c;
294
        if (!head) {
295
            head = c;
296
        }
297
    }
298
    if (a) {
299
        c->next[i] = a;
300
        a->prev[i] = c;
301
    } else if (b) {
302
        c->next[i] = b;
303
        b->prev[i] = c;
304
    }
305
    return head;
306
}
307

308
#ifdef DEBUG
309
static int
310
test_sorted(PixelList *pl[3]) {
311
    int i, l;
312
    PixelList *t;
313

314
    for (i = 0; i < 3; i++) {
315
        l = 256;
316
        for (t = pl[i]; t; t = t->next[i]) {
317
            if (l < t->p.a.v[i]) {
318
                return 0;
319
            }
320
            l = t->p.a.v[i];
321
        }
322
    }
323
    return 1;
324
}
325
#endif
326

327
static int
328
box_heap_cmp(const Heap *h, const void *A, const void *B) {
329
    BoxNode *a = (BoxNode *)A;
330
    BoxNode *b = (BoxNode *)B;
331
    return (int)a->pixelCount - (int)b->pixelCount;
332
}
333

334
#define LUMINANCE(p) (77 * (p)->c.r + 150 * (p)->c.g + 29 * (p)->c.b)
335

336
static int
337
splitlists(
338
    PixelList *h[3],
339
    PixelList *t[3],
340
    PixelList *nh[2][3],
341
    PixelList *nt[2][3],
342
    uint32_t nCount[2],
343
    int axis,
344
    uint32_t pixelCount
345
) {
346
    uint32_t left;
347

348
    PixelList *l, *r, *c, *n;
349
    int i;
350
    int nRight;
351
#ifdef DEBUG
352
    int nLeft;
353
#endif
354
    int splitColourVal;
355

356
#ifdef DEBUG
357
    {
358
        PixelList *_prevTest, *_nextTest;
359
        int _i, _nextCount[3], _prevCount[3];
360
        for (_i = 0; _i < 3; _i++) {
361
            for (_nextCount[_i] = 0, _nextTest = h[_i];
362
                 _nextTest && _nextTest->next[_i];
363
                 _nextTest = _nextTest->next[_i], _nextCount[_i]++);
364
            for (_prevCount[_i] = 0, _prevTest = t[_i];
365
                 _prevTest && _prevTest->prev[_i];
366
                 _prevTest = _prevTest->prev[_i], _prevCount[_i]++);
367
            if (_nextTest != t[_i]) {
368
                printf("next-list of axis %d does not end at tail\n", _i);
369
                exit(1);
370
            }
371
            if (_prevTest != h[_i]) {
372
                printf("prev-list of axis %d does not end at head\n", _i);
373
                exit(1);
374
            }
375
            for (; _nextTest && _nextTest->prev[_i]; _nextTest = _nextTest->prev[_i]);
376
            for (; _prevTest && _prevTest->next[_i]; _prevTest = _prevTest->next[_i]);
377
            if (_nextTest != h[_i]) {
378
                printf("next-list of axis %d does not loop back to head\n", _i);
379
                exit(1);
380
            }
381
            if (_prevTest != t[_i]) {
382
                printf("prev-list of axis %d does not loop back to tail\n", _i);
383
                exit(1);
384
            }
385
        }
386
        for (_i = 1; _i < 3; _i++) {
387
            if (_prevCount[_i] != _prevCount[_i - 1] ||
388
                _nextCount[_i] != _nextCount[_i - 1] ||
389
                _prevCount[_i] != _nextCount[_i]) {
390
                printf(
391
                    "{%d %d %d} {%d %d %d}\n",
392
                    _prevCount[0],
393
                    _prevCount[1],
394
                    _prevCount[2],
395
                    _nextCount[0],
396
                    _nextCount[1],
397
                    _nextCount[2]
398
                );
399
                exit(1);
400
            }
401
        }
402
    }
403
#endif
404
    nCount[0] = nCount[1] = 0;
405
    nRight = 0;
406
#ifdef DEBUG
407
    nLeft = 0;
408
#endif
409
    for (left = 0, c = h[axis]; c;) {
410
        left = left + c->count;
411
        nCount[0] += c->count;
412
        c->flag = 0;
413
#ifdef DEBUG
414
        nLeft++;
415
#endif
416
        c = c->next[axis];
417
        if (left * 2 > pixelCount) {
418
            break;
419
        }
420
    }
421
    if (c) {
422
        splitColourVal = c->prev[axis]->p.a.v[axis];
423
        for (; c; c = c->next[axis]) {
424
            if (splitColourVal != c->p.a.v[axis]) {
425
                break;
426
            }
427
            c->flag = 0;
428
#ifdef DEBUG
429
            nLeft++;
430
#endif
431
            nCount[0] += c->count;
432
        }
433
    }
434
    for (; c; c = c->next[axis]) {
435
        c->flag = 1;
436
        nRight++;
437
        nCount[1] += c->count;
438
    }
439
    if (!nRight) {
440
        for (c = t[axis], splitColourVal = t[axis]->p.a.v[axis]; c; c = c->prev[axis]) {
441
            if (splitColourVal != c->p.a.v[axis]) {
442
                break;
443
            }
444
            c->flag = 1;
445
            nRight++;
446
#ifdef DEBUG
447
            nLeft--;
448
#endif
449
            nCount[0] -= c->count;
450
            nCount[1] += c->count;
451
        }
452
    }
453
#ifdef DEBUG
454
    if (!nLeft) {
455
        for (c = h[axis]; c; c = c->next[axis]) {
456
            printf("[%d %d %d]\n", c->p.c.r, c->p.c.g, c->p.c.b);
457
        }
458
        printf("warning... trivial split\n");
459
    }
460
#endif
461

462
    for (i = 0; i < 3; i++) {
463
        l = r = NULL;
464
        nh[0][i] = nt[0][i] = NULL;
465
        nh[1][i] = nt[1][i] = NULL;
466
        for (c = h[i]; c; c = n) {
467
            n = c->next[i];
468
            if (c->flag) { /* move pixel to right  list*/
469
                if (r) {
470
                    r->next[i] = c;
471
                } else {
472
                    nh[1][i] = c;
473
                }
474
                c->prev[i] = r;
475
                r = c;
476
            } else { /* move pixel to left list */
477
                if (l) {
478
                    l->next[i] = c;
479
                } else {
480
                    nh[0][i] = c;
481
                }
482
                c->prev[i] = l;
483
                l = c;
484
            }
485
        }
486
        if (l) {
487
            l->next[i] = NULL;
488
        }
489
        if (r) {
490
            r->next[i] = NULL;
491
        }
492
        nt[0][i] = l;
493
        nt[1][i] = r;
494
    }
495
    return 1;
496
}
497

498
static int
499
split(BoxNode *node) {
500
    unsigned char rl, rh, gl, gh, bl, bh;
501
    int f[3];
502
    int best, axis;
503
    int i;
504
    PixelList *heads[2][3];
505
    PixelList *tails[2][3];
506
    uint32_t newCounts[2];
507
    BoxNode *left, *right;
508

509
    rh = node->head[0]->p.c.r;
510
    rl = node->tail[0]->p.c.r;
511
    gh = node->head[1]->p.c.g;
512
    gl = node->tail[1]->p.c.g;
513
    bh = node->head[2]->p.c.b;
514
    bl = node->tail[2]->p.c.b;
515
#ifdef DEBUG
516
    printf("splitting node [%d %d %d] [%d %d %d] ", rl, gl, bl, rh, gh, bh);
517
#endif
518
    f[0] = (rh - rl) * 77;
519
    f[1] = (gh - gl) * 150;
520
    f[2] = (bh - bl) * 29;
521

522
    best = f[0];
523
    axis = 0;
524
    for (i = 1; i < 3; i++) {
525
        if (best < f[i]) {
526
            best = f[i];
527
            axis = i;
528
        }
529
    }
530
#ifdef DEBUG
531
    printf("along axis %d\n", axis + 1);
532
    {
533
        PixelList *_prevTest, *_nextTest;
534
        int _i, _nextCount[3], _prevCount[3];
535
        for (_i = 0; _i < 3; _i++) {
536
            if (node->tail[_i]->next[_i]) {
537
                printf("tail is not tail\n");
538
                printf(
539
                    "node->tail[%d]->next[%d]=%p\n", _i, _i, node->tail[_i]->next[_i]
540
                );
541
            }
542
            if (node->head[_i]->prev[_i]) {
543
                printf("head is not head\n");
544
                printf(
545
                    "node->head[%d]->prev[%d]=%p\n", _i, _i, node->head[_i]->prev[_i]
546
                );
547
            }
548
        }
549

550
        for (_i = 0; _i < 3; _i++) {
551
            for (_nextCount[_i] = 0, _nextTest = node->head[_i];
552
                 _nextTest && _nextTest->next[_i];
553
                 _nextTest = _nextTest->next[_i], _nextCount[_i]++);
554
            for (_prevCount[_i] = 0, _prevTest = node->tail[_i];
555
                 _prevTest && _prevTest->prev[_i];
556
                 _prevTest = _prevTest->prev[_i], _prevCount[_i]++);
557
            if (_nextTest != node->tail[_i]) {
558
                printf("next-list of axis %d does not end at tail\n", _i);
559
            }
560
            if (_prevTest != node->head[_i]) {
561
                printf("prev-list of axis %d does not end at head\n", _i);
562
            }
563
            for (; _nextTest && _nextTest->prev[_i]; _nextTest = _nextTest->prev[_i]);
564
            for (; _prevTest && _prevTest->next[_i]; _prevTest = _prevTest->next[_i]);
565
            if (_nextTest != node->head[_i]) {
566
                printf("next-list of axis %d does not loop back to head\n", _i);
567
            }
568
            if (_prevTest != node->tail[_i]) {
569
                printf("prev-list of axis %d does not loop back to tail\n", _i);
570
            }
571
        }
572
        for (_i = 1; _i < 3; _i++) {
573
            if (_prevCount[_i] != _prevCount[_i - 1] ||
574
                _nextCount[_i] != _nextCount[_i - 1] ||
575
                _prevCount[_i] != _nextCount[_i]) {
576
                printf(
577
                    "{%d %d %d} {%d %d %d}\n",
578
                    _prevCount[0],
579
                    _prevCount[1],
580
                    _prevCount[2],
581
                    _nextCount[0],
582
                    _nextCount[1],
583
                    _nextCount[2]
584
                );
585
            }
586
        }
587
    }
588
#endif
589
    node->axis = axis;
590
    if (!splitlists(
591
            node->head, node->tail, heads, tails, newCounts, axis, node->pixelCount
592
        )) {
593
#ifdef DEBUG
594
        printf("list split failed.\n");
595
#endif
596
        return 0;
597
    }
598
#ifdef DEBUG
599
    if (!test_sorted(heads[0])) {
600
        printf("bug in split");
601
        exit(1);
602
    }
603
    if (!test_sorted(heads[1])) {
604
        printf("bug in split");
605
        exit(1);
606
    }
607
#endif
608
    /* malloc check ok, small constant allocation */
609
    left = malloc(sizeof(BoxNode));
610
    right = malloc(sizeof(BoxNode));
611
    if (!left || !right) {
612
        free(left);
613
        free(right);
614
        return 0;
615
    }
616
    for (i = 0; i < 3; i++) {
617
        left->head[i] = heads[0][i];
618
        left->tail[i] = tails[0][i];
619
        right->head[i] = heads[1][i];
620
        right->tail[i] = tails[1][i];
621
        node->head[i] = NULL;
622
        node->tail[i] = NULL;
623
    }
624
#ifdef DEBUG
625
    if (left->head[0]) {
626
        rh = left->head[0]->p.c.r;
627
        rl = left->tail[0]->p.c.r;
628
        gh = left->head[1]->p.c.g;
629
        gl = left->tail[1]->p.c.g;
630
        bh = left->head[2]->p.c.b;
631
        bl = left->tail[2]->p.c.b;
632
        printf("   left node  [%3d %3d %3d] [%3d %3d %3d]\n", rl, gl, bl, rh, gh, bh);
633
    }
634
    if (right->head[0]) {
635
        rh = right->head[0]->p.c.r;
636
        rl = right->tail[0]->p.c.r;
637
        gh = right->head[1]->p.c.g;
638
        gl = right->tail[1]->p.c.g;
639
        bh = right->head[2]->p.c.b;
640
        bl = right->tail[2]->p.c.b;
641
        printf("   right node [%3d %3d %3d] [%3d %3d %3d]\n", rl, gl, bl, rh, gh, bh);
642
    }
643
#endif
644
    left->l = left->r = NULL;
645
    right->l = right->r = NULL;
646
    left->axis = right->axis = -1;
647
    left->volume = right->volume = -1;
648
    left->pixelCount = newCounts[0];
649
    right->pixelCount = newCounts[1];
650
    node->l = left;
651
    node->r = right;
652
    return 1;
653
}
654

655
static BoxNode *
656
median_cut(PixelList *hl[3], uint32_t imPixelCount, int nPixels) {
657
    PixelList *tl[3];
658
    int i;
659
    BoxNode *root;
660
    Heap *h;
661
    BoxNode *thisNode;
662

663
    h = ImagingQuantHeapNew(box_heap_cmp);
664
    /* malloc check ok, small constant allocation */
665
    root = malloc(sizeof(BoxNode));
666
    if (!root) {
667
        ImagingQuantHeapFree(h);
668
        return NULL;
669
    }
670
    for (i = 0; i < 3; i++) {
671
        for (tl[i] = hl[i]; tl[i] && tl[i]->next[i]; tl[i] = tl[i]->next[i]);
672
        root->head[i] = hl[i];
673
        root->tail[i] = tl[i];
674
    }
675
    root->l = root->r = NULL;
676
    root->axis = -1;
677
    root->volume = -1;
678
    root->pixelCount = imPixelCount;
679

680
    ImagingQuantHeapAdd(h, (void *)root);
681
    while (--nPixels) {
682
        do {
683
            if (!ImagingQuantHeapRemove(h, (void **)&thisNode)) {
684
                goto done;
685
            }
686
        } while (compute_box_volume(thisNode) == 1);
687
        if (!split(thisNode)) {
688
#ifdef DEBUG
689
            printf("Oops, split failed...\n");
690
#endif
691
            exit(1);
692
        }
693
        ImagingQuantHeapAdd(h, (void *)(thisNode->l));
694
        ImagingQuantHeapAdd(h, (void *)(thisNode->r));
695
    }
696
done:
697
    ImagingQuantHeapFree(h);
698
    return root;
699
}
700

701
static void
702
free_box_tree(BoxNode *n) {
703
    PixelList *p, *pp;
704
    if (n->l) {
705
        free_box_tree(n->l);
706
    }
707
    if (n->r) {
708
        free_box_tree(n->r);
709
    }
710
    for (p = n->head[0]; p; p = pp) {
711
        pp = p->next[0];
712
        free(p);
713
    }
714
    free(n);
715
}
716

717
#ifdef DEBUG
718
static int
719
checkContained(BoxNode *n, Pixel *pp) {
720
    if (n->l && n->r) {
721
        return checkContained(n->l, pp) + checkContained(n->r, pp);
722
    }
723
    if (n->l || n->r) {
724
        printf("box tree is dead\n");
725
        return 0;
726
    }
727
    if (pp->c.r <= n->head[0]->p.c.r && pp->c.r >= n->tail[0]->p.c.r &&
728
        pp->c.g <= n->head[1]->p.c.g && pp->c.g >= n->tail[1]->p.c.g &&
729
        pp->c.b <= n->head[2]->p.c.b && pp->c.b >= n->tail[2]->p.c.b) {
730
        return 1;
731
    }
732
    return 0;
733
}
734
#endif
735

736
static int
737
annotate_hash_table(BoxNode *n, HashTable *h, uint32_t *box) {
738
    PixelList *p;
739
    PixelHashData *d = (PixelHashData *)hashtable_get_user_data(h);
740
    Pixel q;
741
    if (n->l && n->r) {
742
        return annotate_hash_table(n->l, h, box) && annotate_hash_table(n->r, h, box);
743
    }
744
    if (n->l || n->r) {
745
#ifdef DEBUG
746
        printf("box tree is dead\n");
747
#endif
748
        return 0;
749
    }
750
    for (p = n->head[0]; p; p = p->next[0]) {
751
        PIXEL_UNSCALE(&(p->p), &q, d->scale);
752
        if (!hashtable_insert(h, q, *box)) {
753
#ifdef DEBUG
754
            printf("hashtable insert failed\n");
755
#endif
756
            return 0;
757
        }
758
    }
759
    if (n->head[0]) {
760
        (*box)++;
761
    }
762
    return 1;
763
}
764

765
typedef struct {
766
    uint32_t *distance;
767
    uint32_t index;
768
} DistanceWithIndex;
769

770
static int
771
_distance_index_cmp(const void *a, const void *b) {
772
    DistanceWithIndex *A = (DistanceWithIndex *)a;
773
    DistanceWithIndex *B = (DistanceWithIndex *)b;
774
    if (*A->distance == *B->distance) {
775
        return A->index < B->index ? -1 : +1;
776
    }
777
    return *A->distance < *B->distance ? -1 : +1;
778
}
779

780
static int
781
resort_distance_tables(
782
    uint32_t *avgDist, uint32_t **avgDistSortKey, Pixel *p, uint32_t nEntries
783
) {
784
    uint32_t i, j, k;
785
    uint32_t **skRow;
786
    uint32_t *skElt;
787

788
    for (i = 0; i < nEntries; i++) {
789
        avgDist[i * nEntries + i] = 0;
790
        for (j = 0; j < i; j++) {
791
            avgDist[j * nEntries + i] = avgDist[i * nEntries + j] =
792
                _DISTSQR(p + i, p + j);
793
        }
794
    }
795
    for (i = 0; i < nEntries; i++) {
796
        skRow = avgDistSortKey + i * nEntries;
797
        for (j = 1; j < nEntries; j++) {
798
            skElt = skRow[j];
799
            for (k = j; k && (*(skRow[k - 1]) > *(skRow[k])); k--) {
800
                skRow[k] = skRow[k - 1];
801
            }
802
            if (k != j) {
803
                skRow[k] = skElt;
804
            }
805
        }
806
    }
807
    return 1;
808
}
809

810
static int
811
build_distance_tables(
812
    uint32_t *avgDist, uint32_t **avgDistSortKey, Pixel *p, uint32_t nEntries
813
) {
814
    uint32_t i, j;
815
    DistanceWithIndex *dwi;
816

817
    for (i = 0; i < nEntries; i++) {
818
        avgDist[i * nEntries + i] = 0;
819
        avgDistSortKey[i * nEntries + i] = &(avgDist[i * nEntries + i]);
820
        for (j = 0; j < i; j++) {
821
            avgDist[j * nEntries + i] = avgDist[i * nEntries + j] =
822
                _DISTSQR(p + i, p + j);
823
            avgDistSortKey[j * nEntries + i] = &(avgDist[j * nEntries + i]);
824
            avgDistSortKey[i * nEntries + j] = &(avgDist[i * nEntries + j]);
825
        }
826
    }
827

828
    dwi = calloc(nEntries, sizeof(DistanceWithIndex));
829
    if (!dwi) {
830
        return 0;
831
    }
832
    for (i = 0; i < nEntries; i++) {
833
        for (j = 0; j < nEntries; j++) {
834
            dwi[j] = (DistanceWithIndex){&(avgDist[i * nEntries + j]), j};
835
        }
836
        qsort(dwi, nEntries, sizeof(DistanceWithIndex), _distance_index_cmp);
837
        for (j = 0; j < nEntries; j++) {
838
            avgDistSortKey[i * nEntries + j] = dwi[j].distance;
839
        }
840
    }
841
    free(dwi);
842
    return 1;
843
}
844

845
static int
846
map_image_pixels(
847
    Pixel *pixelData,
848
    uint32_t nPixels,
849
    Pixel *paletteData,
850
    uint32_t nPaletteEntries,
851
    uint32_t *avgDist,
852
    uint32_t **avgDistSortKey,
853
    uint32_t *pixelArray
854
) {
855
    uint32_t *aD, **aDSK;
856
    uint32_t idx;
857
    uint32_t i, j;
858
    uint32_t bestdist, bestmatch, dist;
859
    uint32_t initialdist;
860
    HashTable *h2;
861

862
    h2 = hashtable_new(unshifted_pixel_hash, unshifted_pixel_cmp);
863
    for (i = 0; i < nPixels; i++) {
864
        if (!hashtable_lookup(h2, pixelData[i], &bestmatch)) {
865
            bestmatch = 0;
866
            initialdist = _DISTSQR(paletteData + bestmatch, pixelData + i);
867
            bestdist = initialdist;
868
            initialdist <<= 2;
869
            aDSK = avgDistSortKey + bestmatch * nPaletteEntries;
870
            aD = avgDist + bestmatch * nPaletteEntries;
871
            for (j = 0; j < nPaletteEntries; j++) {
872
                idx = aDSK[j] - aD;
873
                if (*(aDSK[j]) <= initialdist) {
874
                    dist = _DISTSQR(paletteData + idx, pixelData + i);
875
                    if (dist < bestdist) {
876
                        bestdist = dist;
877
                        bestmatch = idx;
878
                    }
879
                } else {
880
                    break;
881
                }
882
            }
883
            hashtable_insert(h2, pixelData[i], bestmatch);
884
        }
885
        pixelArray[i] = bestmatch;
886
    }
887
    hashtable_free(h2);
888
    return 1;
889
}
890

891
static int
892
map_image_pixels_from_quantized_pixels(
893
    Pixel *pixelData,
894
    uint32_t nPixels,
895
    Pixel *paletteData,
896
    uint32_t nPaletteEntries,
897
    uint32_t *avgDist,
898
    uint32_t **avgDistSortKey,
899
    uint32_t *pixelArray,
900
    uint32_t *avg[3],
901
    uint32_t *count
902
) {
903
    uint32_t *aD, **aDSK;
904
    uint32_t idx;
905
    uint32_t i, j;
906
    uint32_t bestdist, bestmatch, dist;
907
    uint32_t initialdist;
908
    HashTable *h2;
909
    int changes = 0;
910

911
    h2 = hashtable_new(unshifted_pixel_hash, unshifted_pixel_cmp);
912
    for (i = 0; i < nPixels; i++) {
913
        if (!hashtable_lookup(h2, pixelData[i], &bestmatch)) {
914
            bestmatch = pixelArray[i];
915
            initialdist = _DISTSQR(paletteData + bestmatch, pixelData + i);
916
            bestdist = initialdist;
917
            initialdist <<= 2;
918
            aDSK = avgDistSortKey + bestmatch * nPaletteEntries;
919
            aD = avgDist + bestmatch * nPaletteEntries;
920
            for (j = 0; j < nPaletteEntries; j++) {
921
                idx = aDSK[j] - aD;
922
                if (*(aDSK[j]) <= initialdist) {
923
                    dist = _DISTSQR(paletteData + idx, pixelData + i);
924
                    if (dist < bestdist) {
925
                        bestdist = dist;
926
                        bestmatch = idx;
927
                    }
928
                } else {
929
                    break;
930
                }
931
            }
932
            hashtable_insert(h2, pixelData[i], bestmatch);
933
        }
934
        if (pixelArray[i] != bestmatch) {
935
            changes++;
936
            avg[0][bestmatch] += pixelData[i].c.r;
937
            avg[1][bestmatch] += pixelData[i].c.g;
938
            avg[2][bestmatch] += pixelData[i].c.b;
939
            avg[0][pixelArray[i]] -= pixelData[i].c.r;
940
            avg[1][pixelArray[i]] -= pixelData[i].c.g;
941
            avg[2][pixelArray[i]] -= pixelData[i].c.b;
942
            count[bestmatch]++;
943
            count[pixelArray[i]]--;
944
            pixelArray[i] = bestmatch;
945
        }
946
    }
947
    hashtable_free(h2);
948
    return changes;
949
}
950

951
static int
952
map_image_pixels_from_median_box(
953
    Pixel *pixelData,
954
    uint32_t nPixels,
955
    Pixel *paletteData,
956
    uint32_t nPaletteEntries,
957
    HashTable *medianBoxHash,
958
    uint32_t *avgDist,
959
    uint32_t **avgDistSortKey,
960
    uint32_t *pixelArray
961
) {
962
    uint32_t *aD, **aDSK;
963
    uint32_t idx;
964
    uint32_t i, j;
965
    uint32_t bestdist, bestmatch, dist;
966
    uint32_t initialdist;
967
    HashTable *h2;
968
    uint32_t pixelVal;
969

970
    h2 = hashtable_new(unshifted_pixel_hash, unshifted_pixel_cmp);
971
    for (i = 0; i < nPixels; i++) {
972
        if (hashtable_lookup(h2, pixelData[i], &pixelVal)) {
973
            pixelArray[i] = pixelVal;
974
            continue;
975
        }
976
        if (!hashtable_lookup(medianBoxHash, pixelData[i], &pixelVal)) {
977
#ifdef DEBUG
978
            printf("pixel lookup failed\n");
979
#endif
980
            return 0;
981
        }
982
        initialdist = _DISTSQR(paletteData + pixelVal, pixelData + i);
983
        bestdist = initialdist;
984
        bestmatch = pixelVal;
985
        initialdist <<= 2;
986
        aDSK = avgDistSortKey + pixelVal * nPaletteEntries;
987
        aD = avgDist + pixelVal * nPaletteEntries;
988
        for (j = 0; j < nPaletteEntries; j++) {
989
            idx = aDSK[j] - aD;
990
            if (*(aDSK[j]) <= initialdist) {
991
                dist = _DISTSQR(paletteData + idx, pixelData + i);
992
                if (dist < bestdist) {
993
                    bestdist = dist;
994
                    bestmatch = idx;
995
                }
996
            } else {
997
                break;
998
            }
999
        }
1000
        pixelArray[i] = bestmatch;
1001
        hashtable_insert(h2, pixelData[i], bestmatch);
1002
    }
1003
    hashtable_free(h2);
1004
    return 1;
1005
}
1006

1007
static int
1008
compute_palette_from_median_cut(
1009
    Pixel *pixelData,
1010
    uint32_t nPixels,
1011
    HashTable *medianBoxHash,
1012
    Pixel **palette,
1013
    uint32_t nPaletteEntries,
1014
    BoxNode *root
1015
) {
1016
    uint32_t i;
1017
    uint32_t paletteEntry;
1018
    Pixel *p;
1019
    uint32_t *avg[3];
1020
    uint32_t *count;
1021

1022
    *palette = NULL;
1023
    /* malloc check ok, using calloc */
1024
    if (!(count = calloc(nPaletteEntries, sizeof(uint32_t)))) {
1025
        return 0;
1026
    }
1027
    for (i = 0; i < 3; i++) {
1028
        avg[i] = NULL;
1029
    }
1030
    for (i = 0; i < 3; i++) {
1031
        /* malloc check ok, using calloc */
1032
        if (!(avg[i] = calloc(nPaletteEntries, sizeof(uint32_t)))) {
1033
            for (i = 0; i < 3; i++) {
1034
                if (avg[i]) {
1035
                    free(avg[i]);
1036
                }
1037
            }
1038
            free(count);
1039
            return 0;
1040
        }
1041
    }
1042
    for (i = 0; i < nPixels; i++) {
1043
#ifdef DEBUG
1044
        if (!(i % 100)) {
1045
            printf("%05d\r", i);
1046
            fflush(stdout);
1047
        }
1048
        if (checkContained(root, pixelData + i) > 1) {
1049
            printf("pixel in two boxes\n");
1050
            for (i = 0; i < 3; i++) {
1051
                free(avg[i]);
1052
            }
1053
            free(count);
1054
            return 0;
1055
        }
1056
#endif
1057
        if (!hashtable_lookup(medianBoxHash, pixelData[i], &paletteEntry)) {
1058
#ifdef DEBUG
1059
            printf("pixel lookup failed\n");
1060
#endif
1061
            for (i = 0; i < 3; i++) {
1062
                free(avg[i]);
1063
            }
1064
            free(count);
1065
            return 0;
1066
        }
1067
        if (paletteEntry >= nPaletteEntries) {
1068
#ifdef DEBUG
1069
            printf(
1070
                "panic - paletteEntry>=nPaletteEntries (%d>=%d)\n",
1071
                (int)paletteEntry,
1072
                (int)nPaletteEntries
1073
            );
1074
#endif
1075
            for (i = 0; i < 3; i++) {
1076
                free(avg[i]);
1077
            }
1078
            free(count);
1079
            return 0;
1080
        }
1081
        avg[0][paletteEntry] += pixelData[i].c.r;
1082
        avg[1][paletteEntry] += pixelData[i].c.g;
1083
        avg[2][paletteEntry] += pixelData[i].c.b;
1084
        count[paletteEntry]++;
1085
    }
1086
    /* malloc check ok, using calloc */
1087
    p = calloc(nPaletteEntries, sizeof(Pixel));
1088
    if (!p) {
1089
        for (i = 0; i < 3; i++) {
1090
            free(avg[i]);
1091
        }
1092
        free(count);
1093
        return 0;
1094
    }
1095
    for (i = 0; i < nPaletteEntries; i++) {
1096
        p[i].c.r = (int)(.5 + (double)avg[0][i] / (double)count[i]);
1097
        p[i].c.g = (int)(.5 + (double)avg[1][i] / (double)count[i]);
1098
        p[i].c.b = (int)(.5 + (double)avg[2][i] / (double)count[i]);
1099
    }
1100
    *palette = p;
1101
    for (i = 0; i < 3; i++) {
1102
        free(avg[i]);
1103
    }
1104
    free(count);
1105
    return 1;
1106
}
1107

1108
static int
1109
recompute_palette_from_averages(
1110
    Pixel *palette, uint32_t nPaletteEntries, uint32_t *avg[3], uint32_t *count
1111
) {
1112
    uint32_t i;
1113

1114
    for (i = 0; i < nPaletteEntries; i++) {
1115
        palette[i].c.r = (int)(.5 + (double)avg[0][i] / (double)count[i]);
1116
        palette[i].c.g = (int)(.5 + (double)avg[1][i] / (double)count[i]);
1117
        palette[i].c.b = (int)(.5 + (double)avg[2][i] / (double)count[i]);
1118
    }
1119
    return 1;
1120
}
1121

1122
static int
1123
compute_palette_from_quantized_pixels(
1124
    Pixel *pixelData,
1125
    uint32_t nPixels,
1126
    Pixel *palette,
1127
    uint32_t nPaletteEntries,
1128
    uint32_t *avg[3],
1129
    uint32_t *count,
1130
    uint32_t *qp
1131
) {
1132
    uint32_t i;
1133

1134
    memset(count, 0, sizeof(uint32_t) * nPaletteEntries);
1135
    for (i = 0; i < 3; i++) {
1136
        memset(avg[i], 0, sizeof(uint32_t) * nPaletteEntries);
1137
    }
1138
    for (i = 0; i < nPixels; i++) {
1139
        if (qp[i] >= nPaletteEntries) {
1140
#ifdef DEBUG
1141
            printf("scream\n");
1142
#endif
1143
            return 0;
1144
        }
1145
        avg[0][qp[i]] += pixelData[i].c.r;
1146
        avg[1][qp[i]] += pixelData[i].c.g;
1147
        avg[2][qp[i]] += pixelData[i].c.b;
1148
        count[qp[i]]++;
1149
    }
1150
    for (i = 0; i < nPaletteEntries; i++) {
1151
        palette[i].c.r = (int)(.5 + (double)avg[0][i] / (double)count[i]);
1152
        palette[i].c.g = (int)(.5 + (double)avg[1][i] / (double)count[i]);
1153
        palette[i].c.b = (int)(.5 + (double)avg[2][i] / (double)count[i]);
1154
    }
1155
    return 1;
1156
}
1157

1158
static int
1159
k_means(
1160
    Pixel *pixelData,
1161
    uint32_t nPixels,
1162
    Pixel *paletteData,
1163
    uint32_t nPaletteEntries,
1164
    uint32_t *qp,
1165
    int threshold
1166
) {
1167
    uint32_t *avg[3];
1168
    uint32_t *count;
1169
    uint32_t i;
1170
    uint32_t *avgDist;
1171
    uint32_t **avgDistSortKey;
1172
    int changes;
1173
    int built = 0;
1174

1175
    if (nPaletteEntries > UINT32_MAX / (sizeof(uint32_t))) {
1176
        return 0;
1177
    }
1178
    /* malloc check ok, using calloc */
1179
    if (!(count = calloc(nPaletteEntries, sizeof(uint32_t)))) {
1180
        return 0;
1181
    }
1182
    for (i = 0; i < 3; i++) {
1183
        avg[i] = NULL;
1184
    }
1185
    for (i = 0; i < 3; i++) {
1186
        /* malloc check ok, using calloc */
1187
        if (!(avg[i] = calloc(nPaletteEntries, sizeof(uint32_t)))) {
1188
            goto error_1;
1189
        }
1190
    }
1191

1192
    /* this is enough of a check, since the multiplication n*size is done above */
1193
    if (nPaletteEntries > UINT32_MAX / nPaletteEntries) {
1194
        goto error_1;
1195
    }
1196
    /* malloc check ok, using calloc, checking n*n above */
1197
    avgDist = calloc(nPaletteEntries * nPaletteEntries, sizeof(uint32_t));
1198
    if (!avgDist) {
1199
        goto error_1;
1200
    }
1201

1202
    /* malloc check ok, using calloc, checking n*n above */
1203
    avgDistSortKey = calloc(nPaletteEntries * nPaletteEntries, sizeof(uint32_t *));
1204
    if (!avgDistSortKey) {
1205
        goto error_2;
1206
    }
1207

1208
#ifdef DEBUG
1209
    printf("[");
1210
    fflush(stdout);
1211
#endif
1212
    while (1) {
1213
        if (!built) {
1214
            compute_palette_from_quantized_pixels(
1215
                pixelData, nPixels, paletteData, nPaletteEntries, avg, count, qp
1216
            );
1217
            if (!build_distance_tables(
1218
                    avgDist, avgDistSortKey, paletteData, nPaletteEntries
1219
                )) {
1220
                goto error_3;
1221
            }
1222
            built = 1;
1223
        } else {
1224
            recompute_palette_from_averages(paletteData, nPaletteEntries, avg, count);
1225
            resort_distance_tables(
1226
                avgDist, avgDistSortKey, paletteData, nPaletteEntries
1227
            );
1228
        }
1229
        changes = map_image_pixels_from_quantized_pixels(
1230
            pixelData,
1231
            nPixels,
1232
            paletteData,
1233
            nPaletteEntries,
1234
            avgDist,
1235
            avgDistSortKey,
1236
            qp,
1237
            avg,
1238
            count
1239
        );
1240
        if (changes < 0) {
1241
            goto error_3;
1242
        }
1243
#ifdef DEBUG
1244
        printf(".(%d)", changes);
1245
        fflush(stdout);
1246
#endif
1247
        if (changes <= threshold) {
1248
            break;
1249
        }
1250
    }
1251
#ifdef DEBUG
1252
    printf("]\n");
1253
#endif
1254
    if (avgDistSortKey) {
1255
        free(avgDistSortKey);
1256
    }
1257
    if (avgDist) {
1258
        free(avgDist);
1259
    }
1260
    for (i = 0; i < 3; i++) {
1261
        if (avg[i]) {
1262
            free(avg[i]);
1263
        }
1264
    }
1265
    if (count) {
1266
        free(count);
1267
    }
1268
    return 1;
1269

1270
error_3:
1271
    if (avgDistSortKey) {
1272
        free(avgDistSortKey);
1273
    }
1274
error_2:
1275
    if (avgDist) {
1276
        free(avgDist);
1277
    }
1278
error_1:
1279
    for (i = 0; i < 3; i++) {
1280
        if (avg[i]) {
1281
            free(avg[i]);
1282
        }
1283
    }
1284
    if (count) {
1285
        free(count);
1286
    }
1287
    return 0;
1288
}
1289

1290
static int
1291
quantize(
1292
    Pixel *pixelData,
1293
    uint32_t nPixels,
1294
    uint32_t nQuantPixels,
1295
    Pixel **palette,
1296
    uint32_t *paletteLength,
1297
    uint32_t **quantizedPixels,
1298
    int kmeans
1299
) {
1300
    PixelList *hl[3];
1301
    HashTable *h;
1302
    BoxNode *root;
1303
    uint32_t i;
1304
    uint32_t *qp;
1305
    uint32_t nPaletteEntries;
1306

1307
    uint32_t *avgDist;
1308
    uint32_t **avgDistSortKey;
1309
    Pixel *p;
1310

1311
#ifdef DEBUG
1312
    uint32_t timer, timer2;
1313
#endif
1314

1315
#ifdef DEBUG
1316
    timer2 = clock();
1317
    printf("create hash table...");
1318
    fflush(stdout);
1319
    timer = clock();
1320
#endif
1321
    h = create_pixel_hash(pixelData, nPixels);
1322
#ifdef DEBUG
1323
    printf("done (%f)\n", (clock() - timer) / (double)CLOCKS_PER_SEC);
1324
#endif
1325
    if (!h) {
1326
        goto error_0;
1327
    }
1328

1329
#ifdef DEBUG
1330
    printf("create lists from hash table...");
1331
    fflush(stdout);
1332
    timer = clock();
1333
#endif
1334
    hl[0] = hl[1] = hl[2] = NULL;
1335
    hashtable_foreach(h, hash_to_list, hl);
1336
#ifdef DEBUG
1337
    printf("done (%f)\n", (clock() - timer) / (double)CLOCKS_PER_SEC);
1338
#endif
1339

1340
    if (!hl[0]) {
1341
        goto error_1;
1342
    }
1343

1344
#ifdef DEBUG
1345
    printf("mergesort lists...");
1346
    fflush(stdout);
1347
    timer = clock();
1348
#endif
1349
    for (i = 0; i < 3; i++) {
1350
        hl[i] = mergesort_pixels(hl[i], i);
1351
    }
1352
#ifdef DEBUG
1353
    if (!test_sorted(hl)) {
1354
        printf("bug in mergesort\n");
1355
        goto error_1;
1356
    }
1357
    printf("done (%f)\n", (clock() - timer) / (double)CLOCKS_PER_SEC);
1358
#endif
1359

1360
#ifdef DEBUG
1361
    printf("median cut...");
1362
    fflush(stdout);
1363
    timer = clock();
1364
#endif
1365
    root = median_cut(hl, nPixels, nQuantPixels);
1366
#ifdef DEBUG
1367
    printf("done (%f)\n", (clock() - timer) / (double)CLOCKS_PER_SEC);
1368
#endif
1369
    if (!root) {
1370
        goto error_1;
1371
    }
1372
    nPaletteEntries = 0;
1373
#ifdef DEBUG
1374
    printf("median cut tree to hash table...");
1375
    fflush(stdout);
1376
    timer = clock();
1377
#endif
1378
    annotate_hash_table(root, h, &nPaletteEntries);
1379
#ifdef DEBUG
1380
    printf("done (%f)\n", (clock() - timer) / (double)CLOCKS_PER_SEC);
1381
#endif
1382
#ifdef DEBUG
1383
    printf("compute palette...\n");
1384
    fflush(stdout);
1385
    timer = clock();
1386
#endif
1387
    if (!compute_palette_from_median_cut(
1388
            pixelData, nPixels, h, &p, nPaletteEntries, root
1389
        )) {
1390
        goto error_3;
1391
    }
1392
#ifdef DEBUG
1393
    printf("done (%f)\n", (clock() - timer) / (double)CLOCKS_PER_SEC);
1394
#endif
1395

1396
    free_box_tree(root);
1397
    root = NULL;
1398

1399
    /* malloc check ok, using calloc for overflow */
1400
    qp = calloc(nPixels, sizeof(uint32_t));
1401
    if (!qp) {
1402
        goto error_4;
1403
    }
1404

1405
    if (nPaletteEntries > UINT32_MAX / nPaletteEntries) {
1406
        goto error_5;
1407
    }
1408
    /* malloc check ok, using calloc for overflow, check of n*n above */
1409
    avgDist = calloc(nPaletteEntries * nPaletteEntries, sizeof(uint32_t));
1410
    if (!avgDist) {
1411
        goto error_5;
1412
    }
1413

1414
    /* malloc check ok, using calloc for overflow, check of n*n above */
1415
    avgDistSortKey = calloc(nPaletteEntries * nPaletteEntries, sizeof(uint32_t *));
1416
    if (!avgDistSortKey) {
1417
        goto error_6;
1418
    }
1419

1420
    if (!build_distance_tables(avgDist, avgDistSortKey, p, nPaletteEntries)) {
1421
        goto error_7;
1422
    }
1423

1424
    if (!map_image_pixels_from_median_box(
1425
            pixelData, nPixels, p, nPaletteEntries, h, avgDist, avgDistSortKey, qp
1426
        )) {
1427
        goto error_7;
1428
    }
1429

1430
#ifdef TEST_NEAREST_NEIGHBOUR
1431
#include <math.h>
1432
    {
1433
        uint32_t bestmatch, bestdist, dist;
1434
        HashTable *h2;
1435
        printf("nearest neighbour search (full search)...");
1436
        fflush(stdout);
1437
        timer = clock();
1438
        h2 = hashtable_new(unshifted_pixel_hash, unshifted_pixel_cmp);
1439
        for (i = 0; i < nPixels; i++) {
1440
            if (hashtable_lookup(h2, pixelData[i], &paletteEntry)) {
1441
                bestmatch = paletteEntry;
1442
            } else {
1443
                bestmatch = 0;
1444
                bestdist = _SQR(pixelData[i].c.r - p[0].c.r) +
1445
                           _SQR(pixelData[i].c.g - p[0].c.g) +
1446
                           _SQR(pixelData[i].c.b - p[0].c.b);
1447
                for (j = 1; j < nPaletteEntries; j++) {
1448
                    dist = _SQR(pixelData[i].c.r - p[j].c.r) +
1449
                           _SQR(pixelData[i].c.g - p[j].c.g) +
1450
                           _SQR(pixelData[i].c.b - p[j].c.b);
1451
                    if (dist == bestdist && j == qp[i]) {
1452
                        bestmatch = j;
1453
                    }
1454
                    if (dist < bestdist) {
1455
                        bestdist = dist;
1456
                        bestmatch = j;
1457
                    }
1458
                }
1459
                hashtable_insert(h2, pixelData[i], bestmatch);
1460
            }
1461
            if (qp[i] != bestmatch) {
1462
                printf(
1463
                    "discrepancy in matching algorithms pixel %d [%d %d] %f %f\n",
1464
                    i,
1465
                    qp[i],
1466
                    bestmatch,
1467
                    sqrt((double)(_SQR(pixelData[i].c.r - p[qp[i]].c.r) +
1468
                                  _SQR(pixelData[i].c.g - p[qp[i]].c.g) +
1469
                                  _SQR(pixelData[i].c.b - p[qp[i]].c.b))),
1470
                    sqrt((double)(_SQR(pixelData[i].c.r - p[bestmatch].c.r) +
1471
                                  _SQR(pixelData[i].c.g - p[bestmatch].c.g) +
1472
                                  _SQR(pixelData[i].c.b - p[bestmatch].c.b)))
1473
                );
1474
            }
1475
        }
1476
        hashtable_free(h2);
1477
    }
1478
#endif
1479
#ifdef DEBUG
1480
    printf("k means...\n");
1481
    fflush(stdout);
1482
    timer = clock();
1483
#endif
1484
    if (kmeans > 0) {
1485
        k_means(pixelData, nPixels, p, nPaletteEntries, qp, kmeans - 1);
1486
    }
1487
#ifdef DEBUG
1488
    printf("done (%f)\n", (clock() - timer) / (double)CLOCKS_PER_SEC);
1489
#endif
1490

1491
    *quantizedPixels = qp;
1492
    *palette = p;
1493
    *paletteLength = nPaletteEntries;
1494

1495
#ifdef DEBUG
1496
    printf("cleanup...");
1497
    fflush(stdout);
1498
    timer = clock();
1499
#endif
1500
    if (avgDist) {
1501
        free(avgDist);
1502
    }
1503
    if (avgDistSortKey) {
1504
        free(avgDistSortKey);
1505
    }
1506
    destroy_pixel_hash(h);
1507
#ifdef DEBUG
1508
    printf("done (%f)\n", (clock() - timer) / (double)CLOCKS_PER_SEC);
1509
    printf("-----\ntotal time %f\n", (clock() - timer2) / (double)CLOCKS_PER_SEC);
1510
#endif
1511
    return 1;
1512

1513
error_7:
1514
    if (avgDistSortKey) {
1515
        free(avgDistSortKey);
1516
    }
1517
error_6:
1518
    if (avgDist) {
1519
        free(avgDist);
1520
    }
1521
error_5:
1522
    if (qp) {
1523
        free(qp);
1524
    }
1525
error_4:
1526
    if (p) {
1527
        free(p);
1528
    }
1529
error_3:
1530
    if (root) {
1531
        free_box_tree(root);
1532
    }
1533
error_1:
1534
    destroy_pixel_hash(h);
1535
error_0:
1536
    *quantizedPixels = NULL;
1537
    *paletteLength = 0;
1538
    *palette = NULL;
1539
    return 0;
1540
}
1541

1542
typedef struct {
1543
    Pixel new;
1544
    uint32_t furthestV;
1545
    uint32_t furthestDistance;
1546
    int secondPixel;
1547
} DistanceData;
1548

1549
static void
1550
compute_distances(const HashTable *h, const Pixel pixel, uint32_t *dist, void *u) {
1551
    DistanceData *data = (DistanceData *)u;
1552
    uint32_t oldDist = *dist;
1553
    uint32_t newDist;
1554
    newDist = _DISTSQR(&(data->new), &pixel);
1555
    if (data->secondPixel || newDist < oldDist) {
1556
        *dist = newDist;
1557
        oldDist = newDist;
1558
    }
1559
    if (oldDist > data->furthestDistance) {
1560
        data->furthestDistance = oldDist;
1561
        data->furthestV = pixel.v;
1562
    }
1563
}
1564

1565
static int
1566
quantize2(
1567
    Pixel *pixelData,
1568
    uint32_t nPixels,
1569
    uint32_t nQuantPixels,
1570
    Pixel **palette,
1571
    uint32_t *paletteLength,
1572
    uint32_t **quantizedPixels,
1573
    int kmeans
1574
) {
1575
    HashTable *h;
1576
    uint32_t i;
1577
    uint32_t mean[3];
1578
    Pixel *p;
1579
    DistanceData data;
1580

1581
    uint32_t *qp;
1582
    uint32_t *avgDist;
1583
    uint32_t **avgDistSortKey;
1584

1585
    /* malloc check ok, using calloc */
1586
    p = calloc(nQuantPixels, sizeof(Pixel));
1587
    if (!p) {
1588
        return 0;
1589
    }
1590
    mean[0] = mean[1] = mean[2] = 0;
1591
    h = hashtable_new(unshifted_pixel_hash, unshifted_pixel_cmp);
1592
    for (i = 0; i < nPixels; i++) {
1593
        hashtable_insert(h, pixelData[i], 0xffffffff);
1594
        mean[0] += pixelData[i].c.r;
1595
        mean[1] += pixelData[i].c.g;
1596
        mean[2] += pixelData[i].c.b;
1597
    }
1598
    data.new.c.r = (int)(.5 + (double)mean[0] / (double)nPixels);
1599
    data.new.c.g = (int)(.5 + (double)mean[1] / (double)nPixels);
1600
    data.new.c.b = (int)(.5 + (double)mean[2] / (double)nPixels);
1601
    for (i = 0; i < nQuantPixels; i++) {
1602
        data.furthestDistance = 0;
1603
        data.furthestV = pixelData[0].v;
1604
        data.secondPixel = (i == 1) ? 1 : 0;
1605
        hashtable_foreach_update(h, compute_distances, &data);
1606
        p[i].v = data.furthestV;
1607
        data.new.v = data.furthestV;
1608
    }
1609
    hashtable_free(h);
1610

1611
    /* malloc check ok, using calloc */
1612
    qp = calloc(nPixels, sizeof(uint32_t));
1613
    if (!qp) {
1614
        goto error_1;
1615
    }
1616

1617
    if (nQuantPixels > UINT32_MAX / nQuantPixels) {
1618
        goto error_2;
1619
    }
1620

1621
    /* malloc check ok, using calloc for overflow, check of n*n above */
1622
    avgDist = calloc(nQuantPixels * nQuantPixels, sizeof(uint32_t));
1623
    if (!avgDist) {
1624
        goto error_2;
1625
    }
1626

1627
    /* malloc check ok, using calloc for overflow, check of n*n above */
1628
    avgDistSortKey = calloc(nQuantPixels * nQuantPixels, sizeof(uint32_t *));
1629
    if (!avgDistSortKey) {
1630
        goto error_3;
1631
    }
1632

1633
    if (!build_distance_tables(avgDist, avgDistSortKey, p, nQuantPixels)) {
1634
        goto error_4;
1635
    }
1636

1637
    if (!map_image_pixels(
1638
            pixelData, nPixels, p, nQuantPixels, avgDist, avgDistSortKey, qp
1639
        )) {
1640
        goto error_4;
1641
    }
1642
    if (kmeans > 0) {
1643
        k_means(pixelData, nPixels, p, nQuantPixels, qp, kmeans - 1);
1644
    }
1645

1646
    *paletteLength = nQuantPixels;
1647
    *palette = p;
1648
    *quantizedPixels = qp;
1649
    free(avgDistSortKey);
1650
    free(avgDist);
1651
    return 1;
1652

1653
error_4:
1654
    free(avgDistSortKey);
1655
error_3:
1656
    free(avgDist);
1657
error_2:
1658
    free(qp);
1659
error_1:
1660
    free(p);
1661
    return 0;
1662
}
1663

1664
Imaging
1665
ImagingQuantize(Imaging im, int colors, int mode, int kmeans) {
1666
    int i, j;
1667
    int x, y, v;
1668
    UINT8 *pp;
1669
    Pixel *p;
1670
    Pixel *palette;
1671
    uint32_t paletteLength;
1672
    int result;
1673
    uint32_t *newData;
1674
    Imaging imOut;
1675
    int withAlpha = 0;
1676
    ImagingSectionCookie cookie;
1677

1678
    if (!im) {
1679
        return ImagingError_ModeError();
1680
    }
1681
    if (colors < 1 || colors > 256) {
1682
        /* FIXME: for colors > 256, consider returning an RGB image
1683
           instead (see @PIL205) */
1684
        return (Imaging)ImagingError_ValueError("bad number of colors");
1685
    }
1686

1687
    if (strcmp(im->mode, "L") != 0 && strcmp(im->mode, "P") != 0 &&
1688
        strcmp(im->mode, "RGB") != 0 && strcmp(im->mode, "RGBA") != 0) {
1689
        return ImagingError_ModeError();
1690
    }
1691

1692
    /* only octree and imagequant supports RGBA */
1693
    if (!strcmp(im->mode, "RGBA") && mode != 2 && mode != 3) {
1694
        return ImagingError_ModeError();
1695
    }
1696

1697
    if (im->xsize > INT_MAX / im->ysize) {
1698
        return ImagingError_MemoryError();
1699
    }
1700
    /* malloc check ok, using calloc for final overflow, x*y above */
1701
    p = calloc(im->xsize * im->ysize, sizeof(Pixel));
1702
    if (!p) {
1703
        return ImagingError_MemoryError();
1704
    }
1705

1706
    /* collect statistics */
1707

1708
    /* FIXME: maybe we could load the hash tables directly from the
1709
       image data? */
1710

1711
    if (!strcmp(im->mode, "L")) {
1712
        /* grayscale */
1713

1714
        /* FIXME: converting a "L" image to "P" with 256 colors
1715
           should be done by a simple copy... */
1716

1717
        for (i = y = 0; y < im->ysize; y++) {
1718
            for (x = 0; x < im->xsize; x++, i++) {
1719
                p[i].c.r = p[i].c.g = p[i].c.b = im->image8[y][x];
1720
                p[i].c.a = 255;
1721
            }
1722
        }
1723

1724
    } else if (!strcmp(im->mode, "P")) {
1725
        /* palette */
1726

1727
        pp = im->palette->palette;
1728

1729
        for (i = y = 0; y < im->ysize; y++) {
1730
            for (x = 0; x < im->xsize; x++, i++) {
1731
                v = im->image8[y][x];
1732
                p[i].c.r = pp[v * 4 + 0];
1733
                p[i].c.g = pp[v * 4 + 1];
1734
                p[i].c.b = pp[v * 4 + 2];
1735
                p[i].c.a = pp[v * 4 + 3];
1736
            }
1737
        }
1738

1739
    } else if (!strcmp(im->mode, "RGB") || !strcmp(im->mode, "RGBA")) {
1740
        /* true colour */
1741

1742
        withAlpha = !strcmp(im->mode, "RGBA");
1743
        int transparency = 0;
1744
        unsigned char r = 0, g = 0, b = 0;
1745
        for (i = y = 0; y < im->ysize; y++) {
1746
            for (x = 0; x < im->xsize; x++, i++) {
1747
                p[i].v = im->image32[y][x];
1748
                if (withAlpha && p[i].c.a == 0) {
1749
                    if (transparency == 0) {
1750
                        transparency = 1;
1751
                        r = p[i].c.r;
1752
                        g = p[i].c.g;
1753
                        b = p[i].c.b;
1754
                    } else {
1755
                        /* Set all subsequent transparent pixels
1756
                        to the same colour as the first */
1757
                        p[i].c.r = r;
1758
                        p[i].c.g = g;
1759
                        p[i].c.b = b;
1760
                    }
1761
                }
1762
            }
1763
        }
1764

1765
    } else {
1766
        free(p);
1767
        return (Imaging)ImagingError_ValueError("internal error");
1768
    }
1769

1770
    ImagingSectionEnter(&cookie);
1771

1772
    switch (mode) {
1773
        case 0:
1774
            /* median cut */
1775
            result = quantize(
1776
                p,
1777
                im->xsize * im->ysize,
1778
                colors,
1779
                &palette,
1780
                &paletteLength,
1781
                &newData,
1782
                kmeans
1783
            );
1784
            break;
1785
        case 1:
1786
            /* maximum coverage */
1787
            result = quantize2(
1788
                p,
1789
                im->xsize * im->ysize,
1790
                colors,
1791
                &palette,
1792
                &paletteLength,
1793
                &newData,
1794
                kmeans
1795
            );
1796
            break;
1797
        case 2:
1798
            result = quantize_octree(
1799
                p,
1800
                im->xsize * im->ysize,
1801
                colors,
1802
                &palette,
1803
                &paletteLength,
1804
                &newData,
1805
                withAlpha
1806
            );
1807
            break;
1808
        case 3:
1809
#ifdef HAVE_LIBIMAGEQUANT
1810
            result = quantize_pngquant(
1811
                p,
1812
                im->xsize,
1813
                im->ysize,
1814
                colors,
1815
                &palette,
1816
                &paletteLength,
1817
                &newData,
1818
                withAlpha
1819
            );
1820
#else
1821
            result = -1;
1822
#endif
1823
            break;
1824
        default:
1825
            result = 0;
1826
            break;
1827
    }
1828

1829
    free(p);
1830
    ImagingSectionLeave(&cookie);
1831

1832
    if (result > 0) {
1833
        imOut = ImagingNewDirty("P", im->xsize, im->ysize);
1834
        ImagingSectionEnter(&cookie);
1835

1836
        for (i = y = 0; y < im->ysize; y++) {
1837
            for (x = 0; x < im->xsize; x++) {
1838
                imOut->image8[y][x] = (unsigned char)newData[i++];
1839
            }
1840
        }
1841

1842
        free(newData);
1843

1844
        imOut->palette->size = (int)paletteLength;
1845
        pp = imOut->palette->palette;
1846

1847
        for (i = j = 0; i < (int)paletteLength; i++) {
1848
            *pp++ = palette[i].c.r;
1849
            *pp++ = palette[i].c.g;
1850
            *pp++ = palette[i].c.b;
1851
            if (withAlpha) {
1852
                *pp = palette[i].c.a;
1853
            }
1854
            pp++;
1855
        }
1856

1857
        if (withAlpha) {
1858
            strcpy(imOut->palette->mode, "RGBA");
1859
        }
1860

1861
        free(palette);
1862
        ImagingSectionLeave(&cookie);
1863

1864
        return imOut;
1865

1866
    } else {
1867
        if (result == -1) {
1868
            return (Imaging)ImagingError_ValueError(
1869
                "dependency required by this method was not "
1870
                "enabled at compile time"
1871
            );
1872
        }
1873

1874
        return (Imaging)ImagingError_ValueError("quantization error");
1875
    }
1876
}
1877