com.squareup.gifencoder.OctTreeQuantizer Maven / Gradle / Ivy

Go to download
/*
 * Copyright (C) 2015 Square, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.squareup.gifencoder;

import java.util.HashSet;
import java.util.Set;

/**
 * Implements qct-tree quantization.
 * 
 * The principle of algorithm: http://www.microsoft.com/msj/archive/S3F1.aspx
 * 
 */
public final class OctTreeQuantizer implements ColorQuantizer {
    public static final OctTreeQuantizer INSTANCE = new OctTreeQuantizer();

    private static final char[] mask = new char[]{0x80, 0x40, 0x20, 0x10, 0x08, 0x04, 0x02, 0x01};
    private int leafCount;
    private int inIndex;

    private final Node[] nodeList = new Node[8];

    private OctTreeQuantizer() {
    }

    @Override
    public Set quantize(Multiset originalColors, int maxColorCount) {
        Node node = createNode(0);
        Set distinctElements = originalColors.getDistinctElements();
        for (Color color : distinctElements) {
            addColor(node, color, 0);
            while (leafCount > maxColorCount) {
                reduceTree();
            }
        }
        HashSet colors = new HashSet<>();
        getColorPalette(node, colors);

        leafCount = 0;
        inIndex = 0;
        for (int i = 0; i < 8; i++) {
            nodeList[i] = null;
        }

        return colors;
    }

    private boolean addColor(Node node, Color color, int inLevel) {
        int nIndex, shift;

        if (node == null) {
            node = createNode(inLevel);
        }

        int red = (int) (color.getComponent(0) * 255);
        int green = (int) (color.getComponent(1) * 255);
        int blue = (int) (color.getComponent(2) * 255);

        if (node.isLeaf) {
            node.pixelCount++;
            node.redSum += red;
            node.greenSum += green;
            node.blueSum += blue;
        } else {
            shift = 7 - inLevel;
            nIndex = (((red & mask[inLevel]) >> shift) << 2)
                    | (((green & mask[inLevel]) >> shift) << 1)
                    | ((blue & mask[inLevel]) >> shift);
            Node tmpNode = node.child[nIndex];
            if (tmpNode == null) {
                tmpNode = createNode(inLevel + 1);
            }
            node.child[nIndex] = tmpNode;
            return addColor(node.child[nIndex], color, inLevel + 1);
        }
        return true;
    }

    private Node createNode(int level) {
        Node node = new Node();
        node.level = level;
        node.isLeaf = (level == 8);
        if (node.isLeaf) {
            leafCount++;
        } else {
            node.next = nodeList[level];
            nodeList[level] = node;
        }
        return node;
    }

    private void reduceTree() {
        int i;
        int redSum = 0, greenSum = 0, blueSum = 0, count = 0;

        for (i = 7; i > 0; i--) {
            if (nodeList[i] != null) break;
        }

        Node tmpNode = nodeList[i];
        nodeList[i] = tmpNode.next;

        for (i = 0; i < 8; i++) {
            if (tmpNode.child[i] != null) {
                redSum += tmpNode.child[i].redSum;
                greenSum += tmpNode.child[i].greenSum;
                blueSum += tmpNode.child[i].blueSum;
                count += tmpNode.child[i].pixelCount;
                tmpNode.child[i] = null;
                leafCount--;
            }
        }
        tmpNode.isLeaf = true;
        tmpNode.redSum = redSum;
        tmpNode.greenSum = greenSum;
        tmpNode.blueSum = blueSum;
        tmpNode.pixelCount = count;

        leafCount++;
    }

    private void getColorPalette(Node node, Set colors) {
        if (node.isLeaf) {
            node.colorIndex = inIndex;
            node.redSum = node.redSum / node.pixelCount;
            node.greenSum = node.greenSum / node.pixelCount;
            node.blueSum = node.blueSum / node.pixelCount;
            node.pixelCount = 1;
            inIndex++;
            double red = (double) node.redSum / 255;
            double green = (double) node.greenSum / 255;
            double blue = (double) node.blueSum / 255;
            colors.add(new Color(red, green, blue));
        } else {
            for (int i = 0; i < 8; i++) {
                if (node.child[i] != null) {
                    getColorPalette(node.child[i], colors);
                }
            }
        }
    }

    private static final class Node {
        boolean isLeaf;
        int level;
        int colorIndex;

        int redSum;
        int greenSum;
        int blueSum;
        int pixelCount;

        Node[] child = new Node[8];
        Node next;
    }
}