001/*
002Copyright 2006 Jerry Huxtable
003
004Licensed under the Apache License, Version 2.0 (the "License");
005you may not use this file except in compliance with the License.
006You may obtain a copy of the License at
007
008   http://www.apache.org/licenses/LICENSE-2.0
009
010Unless required by applicable law or agreed to in writing, software
011distributed under the License is distributed on an "AS IS" BASIS,
012WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013See the License for the specific language governing permissions and
014limitations under the License.
015*/
016
017package com.jhlabs.image;
018
019import java.util.*;
020import java.io.*;
021import java.awt.*;
022import java.awt.image.*;
023
024/**
025 * An image Quantizer based on the Octree algorithm. This is a very basic implementation
026 * at present and could be much improved by picking the nodes to reduce more carefully 
027 * (i.e. not completely at random) when I get the time.
028 */
029public class OctTreeQuantizer implements Quantizer {
030
031        /**
032         * The greatest depth the tree is allowed to reach
033         */
034        final static int MAX_LEVEL = 5;
035
036        /**
037         * An Octtree node.
038         */
039        class OctTreeNode {
040                int children;
041                int level;
042                OctTreeNode parent;
043                OctTreeNode leaf[] = new OctTreeNode[8];
044                boolean isLeaf;
045                int count;
046                int     totalRed;
047                int     totalGreen;
048                int     totalBlue;
049                int index;
050                
051                /**
052                 * A debugging method which prints the tree out.
053                 */
054                public void list(PrintStream s, int level) {
055                        for (int i = 0; i < level; i++)
056                                System.out.print(' ');
057                        if (count == 0)
058                                System.out.println(index+": count="+count);
059                        else
060                                System.out.println(index+": count="+count+" red="+(totalRed/count)+" green="+(totalGreen/count)+" blue="+(totalBlue/count));
061                        for (int i = 0; i < 8; i++)
062                                if (leaf[i] != null)
063                                        leaf[i].list(s, level+2);
064                }
065        }
066
067        private int nodes = 0;
068        private OctTreeNode root;
069        private int reduceColors;
070        private int maximumColors;
071        private int colors = 0;
072        private Vector[] colorList;
073        
074        public OctTreeQuantizer() {
075                setup(256);
076                colorList = new Vector[MAX_LEVEL+1];
077                for (int i = 0; i < MAX_LEVEL+1; i++)
078                        colorList[i] = new Vector();
079                root = new OctTreeNode();
080        }
081
082        /**
083         * Initialize the quantizer. This should be called before adding any pixels.
084         * @param numColors the number of colors we're quantizing to.
085         */
086        public void setup(int numColors) {
087                maximumColors = numColors;
088                reduceColors = Math.max(512, numColors * 2);
089        }
090        
091        /**
092         * Add pixels to the quantizer.
093         * @param pixels the array of ARGB pixels
094         * @param offset the offset into the array
095         * @param count the count of pixels
096         */
097        public void addPixels(int[] pixels, int offset, int count) {
098                for (int i = 0; i < count; i++) {
099                        insertColor(pixels[i+offset]);
100                        if (colors > reduceColors)
101                                reduceTree(reduceColors);
102                }
103        }       
104
105    /**
106     * Get the color table index for a color.
107     * @param rgb the color
108     * @return the index
109     */
110        public int getIndexForColor(int rgb) {
111                int red = (rgb >> 16) & 0xff;
112                int green = (rgb >> 8) & 0xff;
113                int blue = rgb & 0xff;
114
115                OctTreeNode node = root;
116
117                for (int level = 0; level <= MAX_LEVEL; level++) {
118                        OctTreeNode child;
119                        int bit = 0x80 >> level;
120
121                        int index = 0;
122                        if ((red & bit) != 0)
123                                index += 4;
124                        if ((green & bit) != 0)
125                                index += 2;
126                        if ((blue & bit) != 0)
127                                index += 1;
128
129                        child = node.leaf[index];
130
131                        if (child == null)
132                                return node.index;
133                        else if (child.isLeaf)
134                                return child.index;
135                        else
136                                node = child;
137                }
138                System.out.println("getIndexForColor failed");
139                return 0;
140        }
141
142        private void insertColor(int rgb) {
143                int red = (rgb >> 16) & 0xff;
144                int green = (rgb >> 8) & 0xff;
145                int blue = rgb & 0xff;
146
147                OctTreeNode node = root;
148
149//              System.out.println("insertColor="+Integer.toHexString(rgb));
150                for (int level = 0; level <= MAX_LEVEL; level++) {
151                        OctTreeNode child;
152                        int bit = 0x80 >> level;
153
154                        int index = 0;
155                        if ((red & bit) != 0)
156                                index += 4;
157                        if ((green & bit) != 0)
158                                index += 2;
159                        if ((blue & bit) != 0)
160                                index += 1;
161
162                        child = node.leaf[index];
163
164                        if (child == null) {
165                                node.children++;
166
167                                child = new OctTreeNode();
168                                child.parent = node;
169                                node.leaf[index] = child;
170                                node.isLeaf = false;
171                                nodes++;
172                                colorList[level].addElement(child);
173
174                                if (level == MAX_LEVEL) {
175                                        child.isLeaf = true;
176                                        child.count = 1;
177                                        child.totalRed = red;
178                                        child.totalGreen = green;
179                                        child.totalBlue = blue;
180                                        child.level = level;
181                                        colors++;
182                                        return;
183                                }
184
185                                node = child;
186                        } else if (child.isLeaf) {
187                                child.count++;
188                                child.totalRed += red;
189                                child.totalGreen += green;
190                                child.totalBlue += blue;
191                                return;
192                        } else
193                                node = child;
194                }
195                System.out.println("insertColor failed");
196        }
197
198        private void reduceTree(int numColors) {
199                for (int level = MAX_LEVEL-1; level >= 0; level--) {
200                        Vector v = colorList[level];
201                        if (v != null && v.size() > 0) {
202                                for (int j = 0; j < v.size(); j++) {
203                                        OctTreeNode node = (OctTreeNode)v.elementAt(j);
204                                        if (node.children > 0) {
205                                                for (int i = 0; i < 8; i++) {
206                                                        OctTreeNode child = node.leaf[i];
207                                                        if (child != null) {
208                                                                if (!child.isLeaf)
209                                                                        System.out.println("not a leaf!");
210                                                                node.count += child.count;
211                                                                node.totalRed += child.totalRed;
212                                                                node.totalGreen += child.totalGreen;
213                                                                node.totalBlue += child.totalBlue;
214                                                                node.leaf[i] = null;
215                                                                node.children--;
216                                                                colors--;
217                                                                nodes--;
218                                                                colorList[level+1].removeElement(child);
219                                                        }
220                                                }
221                                                node.isLeaf = true;
222                                                colors++;
223                                                if (colors <= numColors)
224                                                        return;
225                                        }
226                                }
227                        }
228                }
229
230                System.out.println("Unable to reduce the OctTree");
231        }
232
233    /**
234     * Build the color table.
235     * @return the color table
236     */
237        public int[] buildColorTable() {
238                int[] table = new int[colors];
239                buildColorTable(root, table, 0);
240                return table;
241        }
242        
243        /**
244         * A quick way to use the quantizer. Just create a table the right size and pass in the pixels.
245     * @param inPixels the input colors
246     * @param table the output color table
247     */
248        public void buildColorTable(int[] inPixels, int[] table) {
249                int count = inPixels.length;
250                maximumColors = table.length;
251                for (int i = 0; i < count; i++) {
252                        insertColor(inPixels[i]);
253                        if (colors > reduceColors)
254                                reduceTree(reduceColors);
255                }
256                if (colors > maximumColors)
257                        reduceTree(maximumColors);
258                buildColorTable(root, table, 0);
259        }
260
261        private int buildColorTable(OctTreeNode node, int[] table, int index) {
262                if (colors > maximumColors)
263                        reduceTree(maximumColors);
264
265                if (node.isLeaf) {
266                        int count = node.count;
267                        table[index] = 0xff000000 |
268                                ((node.totalRed/count) << 16) |
269                                ((node.totalGreen/count) << 8) |
270                                node.totalBlue/count;
271                        node.index = index++;
272                } else {
273                        for (int i = 0; i < 8; i++) {
274                                if (node.leaf[i] != null) {
275                                        node.index = index;
276                                        index = buildColorTable(node.leaf[i], table, index);
277                                }
278                        }
279                }
280                return index;
281        }
282        
283}
284