001/*******************************************************************************
002 * This software is provided as a supplement to the authors' textbooks on digital
003 * image processing published by Springer-Verlag in various languages and editions.
004 * Permission to use and distribute this software is granted under the BSD 2-Clause
005 * "Simplified" License (see http://opensource.org/licenses/BSD-2-Clause).
006 * Copyright (c) 2006-2023 Wilhelm Burger, Mark J. Burge. All rights reserved.
007 * Visit https://imagingbook.com for additional details.
008 ******************************************************************************/
009
010package imagingbook.common.color.quantize;
011
012
013import imagingbook.common.color.RgbUtils;
014import imagingbook.common.color.statistics.ColorHistogram;
015
016import java.util.AbstractQueue;
017import java.util.Arrays;
018import java.util.Comparator;
019import java.util.Locale;
020import java.util.PriorityQueue;
021
022/**
023 * <p>
024 * This is an implementation of Heckbert's median-cut color quantization algorithm [1]. Unlike in the original
025 * algorithm, no initial uniform (scalar) quantization is used to for reducing the number of image colors. Instead, all
026 * colors contained in the original image are considered in the quantization process. After the set of representative
027 * colors has been found, each image color is mapped to the closest representative in RGB color space using the
028 * Euclidean distance. See Sec. 13.4.2 (Algs. 13.1-3) of [2] for details.
029 * </p>
030 * <p>
031 * The quantization process has two steps: first a {@link ColorQuantizer} instance is created from a given image using
032 * one of the constructor methods provided. This quantizer can then be used to quantize the original image or any other
033 * image using the same set of representative colors (color table).
034 * </p>
035 * <p>
036 * [1] Heckbert P., "Color Image Quantization for Frame Buffer Display", ACM Transactions on Computer Graphics
037 * (SIGGRAPH), pp. 297-307 (1982). <br> [2] W. Burger, M.J. Burge, <em>Digital Image Processing &ndash; An Algorithmic
038 * Introduction</em>, 3rd ed, Springer (2022).
039 * </p>
040 *
041 * @author WB
042 * @version 2022/10/05 revised to use PriorityQueue, made deterministic
043 */
044public class MedianCutQuantizer implements ColorQuantizer {
045        
046        /** The array of original (distinct) colors. */
047        private final ColorBin[] origColors;
048        /** The final color map (array of representative colors). */
049        private final float[][] colorMap;
050        
051        // -------------------------------------------------------------------------------
052
053        /**
054         * Constructor, creates a new {@link MedianCutQuantizer} with up to K colors, but never more than the number of
055         * colors found in the supplied pixel data.
056         *
057         * @param pixels an image as a aRGB-encoded int array
058         * @param K the desired number of colors (1 or more)
059         */
060        public MedianCutQuantizer(int[] pixels, int K) {
061                if (K < 1) 
062                        throw new IllegalArgumentException("K must be at least 1");
063                this.origColors = getDistinctColors(pixels);
064                this.colorMap = (origColors.length > K) ?
065                                makeColorMap(findReferenceColors(K)) :  // more than K colors
066                                makeColorMap(origColors);                               // exactly K or fewer colors
067        }
068        
069        // -------------------------------------------------------------------------------
070        
071        @Override
072        public float[][] getColorMap() {
073                return this.colorMap;
074        }
075        
076        // -------------------------------------------------------------------------------
077
078        /**
079         * Collects and returns the distinct colors found in the specified pixel array as an array of {@link ColorBin}
080         * instances.
081         *
082         * @param pixels an array of aRGB-encoded color values
083         * @return an array of distinct colors
084         */
085        private ColorBin[] getDistinctColors(int[] pixels) {
086                ColorHistogram ch = new ColorHistogram(pixels, true);
087                final int nc = ch.getNumberOfColors();
088                ColorBin[] colors = new ColorBin[nc];
089                for (int i = 0; i < nc; i++) {
090                        int rgb = ch.getColor(i);
091                        int cnt = ch.getFrequency(i);
092                        colors[i] = new ColorBin(cnt, rgb);
093                }
094                return colors;
095        }
096
097        /**
098         * Performs the actual quantization and returns the calculated reference colors as an array of K
099         * {@link ColorBox} instances. Color boxes are maintained in a sorted list (of type {@link PriorityQueue}) B whose
100         * sorting order is determined by the {@link ColorBox#compareTo(ColorBox)}. Boxes are iteratively split until the
101         * specified number of quantized colors is reached. In each step, the first color box (head element) of the sorted
102         * list is removed and two new boxes (the result of splitting) are added. Thus each step increases the size of list
103         * B by 1.
104         *
105         * @return an array of reference colors
106         */
107        private ColorBox[] findReferenceColors(int K) {
108                final int n = origColors.length;
109                ColorBox cb0 = new ColorBox(0, n - 1, 0);
110                AbstractQueue<ColorBox> B = new PriorityQueue<>();      // the (ordered) set of quantized colors
111                
112                B.add(cb0);     
113                int k = 1;                                              // number of quantized color (color boxes)
114                boolean done = false;
115                
116                while (k < K && !done) {                // B.size() < K
117                        ColorBox cb = B.remove();       // take the first box from the queue // findBoxToSplit(B);
118                        ColorBox[] boxes = cb.splitBox();
119                        if (boxes != null) {            // this should never happen
120                                B.add(boxes[0]);
121                                B.add(boxes[1]);
122                                k = k + 1;
123                        }
124                        else {
125                                done = true;
126                        }
127                }
128                // B.size() == K
129                return B.toArray(new ColorBox[0]);      // sort?
130        }
131        
132        // --------------------------------------------------------
133        
134        @SuppressWarnings("unused")
135        private void listColorBoxes(AbstractQueue<ColorBox> boxes) {
136                if (boxes.isEmpty()) {
137                        System.out.println("boxes empty");
138                        return;
139                }
140                System.out.println("*** boxes " + boxes.size());
141                ColorBox[] all = boxes.toArray(new ColorBox[0]);
142                Arrays.sort(all);
143                for (int i = 0; i < all.length; i++) {
144                        System.out.println("    " + i + ": " + all[i].toString());
145                }
146        }
147        
148        // --------------------------------------------------------
149        
150        private float[][] makeColorMap(ColorBin[] colors) {
151                int n = colors.length;
152                float[][] map = new float[n][];
153                int i = 0;
154                for (ColorBin cn : colors) {
155                        map[i] = new float[] { cn.red, cn.grn, cn.blu };
156                        i++;
157                }
158                return map;
159        }
160        
161        private float[][] makeColorMap(ColorBox[] quantColors) {
162                int n = quantColors.length;
163                float[][] map = new float[n][];
164                int i = 0;
165                for (ColorBox cb : quantColors) {
166                        map[i] = cb.getAvgColor();
167                        i++;
168                }
169                return map;
170        }
171
172        // -------------- class ColorBin -------------------------------------------
173
174        private class ColorBin {
175                private final int red, grn, blu;
176                private final int cnt;
177                
178                private ColorBin(int cnt, int rgb) {
179                        this(cnt, RgbUtils.intToRgb(rgb));
180                }
181                
182                private ColorBin (int cnt, int[] rgb) {
183                        this.cnt = cnt;
184                        this.red = rgb[0];
185                        this.grn = rgb[1];
186                        this.blu = rgb[2];
187                }
188                
189                @Override
190                public String toString() {
191                        String s = ColorBin.class.getSimpleName();
192                        s = s + " red=" + red + " green=" + grn + " blue=" + blu + " count=" + cnt;
193                        return s;
194                }
195        }
196        
197        // -------------- non-static inner class ColorBox -----------------------------------
198
199        /**
200         * Represents a 'color box' holding a set of colors (of type {@link ColorBin}), which is implemented as a contiguous
201         * range of elements in array {@link MedianCutQuantizer#origColors}. Instances of {@link ColorBox} reference this
202         * array directly (this is why this class is non-static). Instances of this class are immutable.
203         */
204        private class ColorBox implements Comparable<ColorBox> { 
205                final int loIdx;                // lower index into 'imageColors'
206                final int hiIdx;                // upper index into 'imageColors'
207                final int level;                // split level of this color box
208                
209                final int count;                                // number of pixels represented by this color box
210                final float rmin, rmax;                 // range of contained colors in red dimension
211                final float gmin, gmax;                 // range of contained colors in green dimension
212                final float bmin, bmax;                 // range of contained colors in blue dimension
213                final float maxSpan;                    // largest side length of this color box 
214                final ColorDimension maxDim;    // color dimension of largest length
215                
216                ColorBox(int loIdx, int hiIdx, int level) {
217                        this.loIdx = loIdx;
218                        this.hiIdx = hiIdx;
219                        this.level = level;
220                        
221                        // trim this color box:
222                        int n = 0;      
223                        float rMin = MAX_RGB, gMin = MAX_RGB, bMin = MAX_RGB;
224                        float rMax = 0, gMax = 0, bMax = 0;             
225                        for (int i = loIdx; i <= hiIdx; i++) {
226                                ColorBin cn = origColors[i];
227                                n = n + cn.cnt;
228                                rMax = Math.max(cn.red, rMax);
229                                rMin = Math.min(cn.red, rMin);
230                                gMax = Math.max(cn.grn, gMax);
231                                gMin = Math.min(cn.grn, gMin);
232                                bMax = Math.max(cn.blu, bMax);
233                                bMin = Math.min(cn.blu, bMin);
234                        }
235                        this.rmax = rMax; this.rmin = rMin;
236                        this.gmax = gMax; this.gmin = gMin;
237                        this.bmax = bMax; this.bmin = bMin;
238                        this.count = n;
239                        
240                        // find the largest dimension of this color box
241                        float rd = rmax - rmin;
242                        float gd = gmax - gmin;
243                        float bd = bmax - bmin;
244                        float maxS = rd;
245                        ColorDimension maxD = ColorDimension.RED;
246                        if (bd >= rd && bd >= gd) {
247                                maxS = bd;
248                                maxD = ColorDimension.BLU;
249                        }
250                        else if (gd >= rd && gd >= bd) {
251                                maxS = gd;
252                                maxD = ColorDimension.GRN;
253                        }
254        
255                        this.maxSpan = maxS;
256                        this.maxDim = maxD;
257                }
258                
259                /**
260                 * Returns the number of different colors within this color box.
261                 * 
262                 * @return the number of different colors
263                 */
264                int colorCount() {
265                        return hiIdx - loIdx;
266                }
267
268                /**
269                 * Splits this color box at the median point along its longest color dimension. Modifies the original color box
270                 * and creates a new one, which is returned.
271                 *
272                 * @return A new color box.
273                 */
274                ColorBox[] splitBox() { 
275                        if (this.colorCount() < 2)      // this box cannot be split
276                                throw new RuntimeException("cannot split " + this.toString());
277                        
278                        int m = this.level;
279                        // find longest dimension of this box:
280                        ColorDimension dim = this.maxDim; // this.getMaxBoxDimension();
281                        // find median along dimension d
282                        int medIdx = this.findMedian(dim);
283                        // now split this box at the median return the resulting new boxes
284                        ColorBox cb1 = new ColorBox(loIdx, medIdx, m + 1);
285                        ColorBox cb2 = new ColorBox(medIdx + 1, hiIdx, m + 1);
286                        return new ColorBox[] {cb1, cb2};
287                }
288
289                /**
290                 * Finds the position of the median of this color box in RGB space along the red, green or blue dimension,
291                 * respectively.
292                 *
293                 * @param dim the color dimension
294                 * @return the median value
295                 */
296                int findMedian(ColorDimension dim) {
297                        // sort color in this box along dimension dim:
298                        Arrays.sort(origColors, loIdx, hiIdx + 1, dim.comparator);
299                        // find the median index:
300                        int half = count / 2;
301                        int k, pixCnt;
302                        for (k = loIdx, pixCnt = 0; k < hiIdx; k++) {
303                                pixCnt = pixCnt + origColors[k].cnt;
304                                if (pixCnt >= half)
305                                        break;
306                        }                       
307                        return k;               // k = index of median
308                }
309
310                /**
311                 * Calculates the (linear) average of the colors contained in this color box and returns the result as a
312                 * {@code float} vector. This is the 3D centroid (in RGB color space) of the pixel colors represented by this
313                 * box.
314                 *
315                 * @return the average color for this box
316                 */
317                private float[] getAvgColor() {
318                        double rSum = 0;
319                        double gSum = 0;
320                        double bSum = 0;
321                        int n = 0;                              // pixel count for this color box
322                        for (int i = loIdx; i <= hiIdx; i++) {
323                                ColorBin cn = origColors[i];
324                                int cnt = cn.cnt;
325                                rSum = rSum + cnt * cn.red;
326                                gSum = gSum + cnt * cn.grn;
327                                bSum = bSum + cnt * cn.blu;
328                                n = n + cnt;
329                        }
330                        float avgRed = (float) (rSum / n);
331                        float avgGrn = (float) (gSum / n);
332                        float avgBlu = (float) (bSum / n);
333                        return new float[] {avgRed, avgGrn, avgBlu};
334                }
335
336                @Override
337                public String toString() {
338                        return String.format(Locale.US, 
339                                        "%s[lev=%1d lo=%5d hi=%5d cnt=%5d R={%.1f, %.1f} G={%.1f, %.1f} B={%.1f, %.1f} maxD=%s maxS=%.1f]", 
340                                        getClass().getSimpleName(), level, loIdx, hiIdx, count, 
341                                        rmin, rmax, gmin, gmax, bmin, bmax, maxDim.toString(), maxSpan);
342                }
343
344                /**
345                 * This method is required to satisfy the {@link Comparator} interface. It determines how boxes are selected for
346                 * splitting. The specific strategy is to prefer color boxes at low levels. Boxes at the same level are selected
347                 * based on their maximum color span, i.e., their maximum side length in color space. Finally, the number of
348                 * pixels decides.
349                 */
350                @Override
351                public int compareTo(ColorBox other) {
352                        int v = Integer.compare(this.level, other.level);       // first compare by level (lower levels first)
353                        if (v != 0) return v;
354                        v = Float.compare(other.maxSpan, this.maxSpan);         // then compare by color span (larger spans first)
355                        if (v != 0) return v;
356                        v = Integer.compare(other.count, this.count);           // then compare by number of pixels (more pixels first)
357                        return v;
358                }
359        }
360
361        /**
362         * Associates the color dimensions RED, GRN, BLU with the corresponding comparators (for sorting along selected
363         * color dimensions).
364         */
365        private enum ColorDimension {
366                RED (new Comparator<ColorBin>() {
367                        @Override
368                        public int compare(ColorBin colA, ColorBin colB) {
369                                return Float.compare(colA.red, colB.red);
370                        }}), 
371                GRN (new Comparator<ColorBin>() {
372                        @Override
373                        public int compare(ColorBin colA, ColorBin colB) {
374                                return Float.compare(colA.grn, colB.grn);
375                        }}), 
376                BLU (new Comparator<ColorBin>() {
377                        @Override
378                        public int compare(ColorBin colA, ColorBin colB) {
379                                return Float.compare(colA.blu, colB.blu);
380                        }});
381
382                private final Comparator<ColorBin> comparator;
383
384                // constructor:
385                private ColorDimension(Comparator<ColorBin> cmp) {
386                        this.comparator = cmp;
387                }
388        }
389        
390}
391