Wu Color Quantization

Wu color quantization 알고리즘은 $\tt RGB$ 이미지를 제한된 수의 대표색 팔레트로 줄이되, 전체 색상 근사 오차의 제곱합($\tt SSE:$ Sum of Squared Errors)을 작게 만드는 방법이다. 3차원 $\tt RGB$ 공간을 작은 상자의 모임으로 분할할 때,  각 $\tt RGB$ 색을 벡터 $\mathbf{x}_i=(r_i,g_i,b_i)$, 한 색상 상자(box)의 대표색을 $\boldsymbol\mu$라 하면, 그 상자의 오차는 $$\operatorname{SSE}(B) = \sum_{i\in B}\left\|\mathbf{x}_i-\boldsymbol\mu\right\|^2$$

이다. 평균색을 대표색으로 쓰면 이 값은 최소가 되며,

$$\operatorname{SSE}(B) = \sum_{i\in B}(r_i^2+g_i^2+b_i^2) - \frac{(\sum r_i)^2+(\sum g_i)^2+(\sum b_i)^2}{N}$$로 계산된다.

1. $\tt RGB$ 색공간을 보통 채널당 5비트 정도로 양자화하여 3차원 히스토그램 생성
2. 각 색상 bin에 대해 픽셀 수 $N$, 각 채널의 합 $\sum r_i$, $\sum g_i $, $\sum b_i$, 그리고 제곱합 $\sum(r_i^2+g_i^2+b_i^2)$을 저장
3. 이 값들에 대한 3차원 누적합을 구하여 $\tt RGB$ 공간의 임의 직육면체 box에 포함된 픽셀들의 통계를 빠르게 구할 수 있게 함
4. 처음에는 전체 $\tt RGB$ 공간을 하나의 box로 둔 후 각 box를 $\tt R·G·B$ 세 축 중 한 방향에서 둘로 나누는 모든 후보를 조사
5. 현재 박스들 중 $\tt SSE$가 가장 큰 박스를 선택하고, 그 박스에 대해 분할 후 두 박스의 $\tt SSE$ 합이 최소가 되는 절단 축과 위치를 선택한다. 
6. 이 과정을 원하는 팔레트 색 수 $K$개가 될 때까지 반복하고, 최종 box 각각의 평균 RGB 값이 팔레트의 대표색이 된다.

Wu 알고리즘이 실제로 최소화하는 것은 단순한 분산이 아니라, 픽셀 수까지 반영한 가중 제곱오차합이다. 픽셀 수가 서로 다른 box의 분할 우선순위를 정해야 하므로 이 차이가 중요하게 된다. Wu 알고리즘은 3차원 누적합으로 $\tt RGB$ box의 SSE를 빠르게 계산하고, SSE 감소량이 가장 큰 box 분할을 반복하여 팔레트를 구성하는 최소제곱 기반 색상 양자화 알고리즘이다.

 

SSE( ):  ColorBox 영역 내 픽셀들이 대표 색상으로부터 벗어난 SSE를 계산. 픽셀 수(weight)로 나누어 평균 오차(MSE)를 구하지 않고 오차의 총량(Total Error)을 유지하므로 픽셀 수가 많고 색상 차이가 큰 박스일수록 높은 SSE 값을 가지며, 이는 전체 화질 오차 지분을 나타낸다.

 

bestCut( ): 선택된 박스를 $\tt R$, $\tt G$, $\tt B$ 3개 축 전체에 대해 조사하여 어느 축에서 평면으로 자르는 것이 오차 감소를 최대화할지 탐색한다. 주어진 박스 내의 컬러 벡터 제곱의 합을 $||S||^2$으로 놓으면, 부모 박스 B를 두 서브박스 $L$과 $R$로 분할할 때 $$ \operatorname{SSE}(L) + \operatorname{SSE}(R) = \sum_{B} (r_i^2 + g_i^2 + b_i^2) - \frac{||S_L||^2}{w_L} - \frac{||S_R||^2}{w_R}  $$이므로 $$\text{score} = \frac{||S_L ||^2 }{w_L}+\frac{||S_R|| ^2}{w_R}$$가 가장 커지는 절단점을 찾는 것이 분할 후 두 박스의 SSE 합을 최소화(= 오차 감소량 $\Delta$)하는 지점을 찾는 것과 같다.

const int boxSide  = 33;
const int tableSize = boxSide * boxSide * boxSide;

struct ColorBox {
    int   r0, r1;
    int   g0, g1;
    int   b0, b1;
    bool  splittable;
    ColorBox() : r0(0), r1(0), g0(0), g1(0), b0(0), b1(0), splittable(true) {}
    ColorBox(int r0_, int r1_, int g0_, int g1_, int b0_, int b1_)
      : r0(r0_), r1(r1_), g0(g0_), g1(g1_), b0(b0_), b1(b1_), splittable(true) {}
};
enum Axis { Red, Green, Blue };

class WuQuantizer {
public:
    WuQuantizer()
        : weights(tableSize), rMoments(tableSize), gMoments(tableSize),
          bMoments(tableSize), sqMoments(tableSize) {}
          
    std::vector<RGBQUAD> processImage(const CRaster& raster, int maxColors) {
        if (raster.IsEmpty() || raster.GetBPP() != 24)
            return std::vector<RGBQUAD>();
        maxColors = max(1, min(256, maxColors));
        clearMoments();
        buildHistogram(raster);
        buildCumulativeMoments();

        std::vector<ColorBox> boxes;
        boxes.reserve(maxColors);
        boxes.push_back(ColorBox(0, 32, 0, 32, 0, 32));
        while (boxes.size() < maxColors) {
            const int index = findBoxWithMaxSSE(boxes);
            if (index < 0) break;
            ColorBox newBox;
            if (splitBox(boxes[index], newBox))
                boxes.push_back(newBox);
            else
                boxes[index].splittable = false;
        }
        std::vector<RGBQUAD> palette = getPalette(boxes);
        setInverseMap(palette);
        return palette;
    }

    BYTE quantizeColor(BYTE r, BYTE g, BYTE b) const {
        return inverseMap[(r >> 3) + 1][(g >> 3) + 1][(b >> 3) + 1];
    }

private:
    std::vector<int64_t> weights;
    std::vector<int64_t> rMoments;
    std::vector<int64_t> gMoments;
    std::vector<int64_t> bMoments;
    std::vector<double>  sqMoments;
    BYTE inverseMap[boxSide][boxSide][boxSide];

    inline int index(int r, int g, int b) const {
        return (r * boxSide + g) * boxSide + b;
    }

    void clearMoments() {
        std::fill(weights.begin(), weights.end(), 0);
        std::fill(rMoments.begin(), rMoments.end(), 0);
        std::fill(gMoments.begin(), gMoments.end(), 0);
        std::fill(bMoments.begin(), bMoments.end(), 0);
        std::fill(sqMoments.begin(), sqMoments.end(), 0.0);
    }

    void buildHistogram(const CRaster& raster) {
        const CSize size = raster.GetSize();
        for (int y = 0; y < size.cy; ++y) {
            const BYTE* p = (BYTE *)raster.GetLinePtr(y);
            for (int x = 0; x < size.cx; ++x) {
                const BYTE b = *p++;
                const BYTE g = *p++;
                const BYTE r = *p++;
                int bin = index((r >> 3) + 1, (g >> 3) + 1, (b >> 3) + 1);

                ++weights[bin];
                rMoments[bin]  += r;
                gMoments[bin]  += g;
                bMoments[bin]  += b;
                sqMoments[bin] += r * r + g * g + b * b;
            }
        }
    }
    // 3차원 누적히스토그램;
    void buildCumulativeMoments() {
        std::vector<int64_t> wArea(boxSide);
        std::vector<int64_t> rArea(boxSide);
        std::vector<int64_t> gArea(boxSide);
        std::vector<int64_t> bArea(boxSide);
        std::vector<double>  sqArea(boxSide); 
        for (int r = 1; r <= 32; ++r) {
            std::fill(wArea.begin(), wArea.end(), 0);
            std::fill(rArea.begin(), rArea.end(), 0);
            std::fill(gArea.begin(), gArea.end(), 0);
            std::fill(bArea.begin(), bArea.end(), 0);
            std::fill(sqArea.begin(), sqArea.end(), 0.0);

            for (int g = 1; g <= 32; ++g) {
                int64_t wLine = 0;
                int64_t rLine = 0;
                int64_t gLine = 0;
                int64_t bLine = 0;
                double sqLine = 0.0;

                for (int b = 1; b <= 32; ++b) {
                    const int curr = index(r, g, b);
                    wLine  += weights[curr];
                    rLine  += rMoments[curr];
                    gLine  += gMoments[curr];
                    bLine  += bMoments[curr];
                    sqLine += sqMoments[curr];

                    wArea[b]  += wLine;
                    rArea[b]  += rLine;
                    gArea[b]  += gLine;
                    bArea[b]  += bLine;
                    sqArea[b] += sqLine;

                    int rPrevPlane = index(r - 1, g, b);
                    weights[curr]   = weights[rPrevPlane] + wArea[b];
                    rMoments[curr]  = rMoments[rPrevPlane] + rArea[b];
                    gMoments[curr]  = gMoments[rPrevPlane] + gArea[b];
                    bMoments[curr]  = bMoments[rPrevPlane] + bArea[b];
                    sqMoments[curr] = sqMoments[rPrevPlane] + sqArea[b];
                }
            }
        }
    }
    // 3차원 적분영상(tab)을 이용한 box 내의 r, g, b합 및 (r^2+g^2+b^2) 합을 O(1) 내에 구함
    template <typename T>
    T volume(const ColorBox& box, const std::vector<T>& tab) const {
        return tab[index(box.r1, box.g1, box.b1)] - tab[index(box.r1, box.g1, box.b0)]
             - tab[index(box.r1, box.g0, box.b1)] + tab[index(box.r1, box.g0, box.b0)]
             - tab[index(box.r0, box.g1, box.b1)] + tab[index(box.r0, box.g1, box.b0)]
             + tab[index(box.r0, box.g0, box.b1)] - tab[index(box.r0, box.g0, box.b0)];
    }

    double SSE(const ColorBox& box) const {
        const int64_t weight = volume(box, weights);
        if (weight == 0) return 0.0;

        const int64_t sr = volume(box, rMoments);
        const int64_t sg = volume(box, gMoments);
        const int64_t sb = volume(box, bMoments);
        const double sum2 = double(sr) * double(sr) 
                          + double(sg) * double(sg)
                          + double(sb) * double(sb);
        return volume(box, sqMoments) - sum2 / double(weight);
    }
   
    int findBoxWithMaxSSE(std::vector<ColorBox>& boxes) const {
        int best = -1;
        double maxSSE = 0.0; 
         for (int i = 0; i < boxes.size(); ++i) {
            if (!boxes[i].splittable) continue;
            const double sse = SSE(boxes[i]);
            if (sse > maxSSE) {
                maxSSE = sse;
                best = i;
            }
        }
        return best;
    }
    // 주어진 박스를 분할 후 전체 SSE를 가장 줄여주는 절단면으로 분할;
    bool splitBox(ColorBox& box, ColorBox& other) const {
        const std::pair<Axis, int> cut = bestCut(box);
        if (cut.second < 0) return false;

        other = box;
        switch (cut.first) {
        case Red:   box.r1 = cut.second; other.r0 = cut.second; break;
        case Green: box.g1 = cut.second; other.g0 = cut.second; break;
        case Blue:  box.b1 = cut.second; other.b0 = cut.second; break;
        }
        return true;
    }

    std::pair<Axis, int> bestCut(const ColorBox& box) const {
        Axis bestAxis = Red;
        int bestPosition = -1;
        double bestScore = 0.0;
        const int64_t wTotal = volume(box, weights);

        const Axis axes[] = {Red, Green, Blue};
        for (int axisIdx = 0; axisIdx < 3; ++axisIdx) {
            const Axis axis = axes[axisIdx];
            const int first = axis == Red ? box.r0 : axis == Green ? box.g0 : box.b0;
            const int last = axis == Red ? box.r1 : axis == Green ? box.g1 : box.b1;

            for (int cut = first + 1; cut < last; ++cut) {
                ColorBox lower = box;
                if (axis == Red) lower.r1 = cut;
                else if (axis == Green) lower.g1 = cut;
                else lower.b1 = cut;

                const int64_t wLower = volume(lower, weights);
                const int64_t wUpper = wTotal - wLower;
                if (wLower == 0 || wUpper == 0) continue;

                const int64_t rL = volume(lower, rMoments);
                const int64_t gL = volume(lower, gMoments);
                const int64_t bL = volume(lower, bMoments);
                const int64_t rU = volume(box, rMoments) - rL;
                const int64_t gU = volume(box, gMoments) - gL;
                const int64_t bU = volume(box, bMoments) - bL;
                double score = double(rL * rL + gL * gL + bL * bL) / wLower
                    + double(rU * rU + gU * gU + bU * bU) / wUpper;

                if (score > bestScore) {
                    bestScore = score;
                    bestAxis = axis;
                    bestPosition = cut;
                }
            }
        }
        return std::make_pair(bestAxis, bestPosition);
    }

    std::vector<RGBQUAD> getPalette(const std::vector<ColorBox>& boxes) const {
        std::vector<RGBQUAD> palette;
        palette.reserve(boxes.size());

        for (int i = 0; i < boxes.size(); ++i) {
            const ColorBox& box = boxes[i];
            const int64_t weight = volume(box, weights);
            if (weight == 0) continue;

            RGBQUAD color  = {0, 0, 0, 0};
            color.rgbRed   = BYTE(volume(box, rMoments) / weight);
            color.rgbGreen = BYTE(volume(box, gMoments) / weight);
            color.rgbBlue  = BYTE(volume(box, bMoments) / weight);
            palette.push_back(color);
        }
        return palette;
    }
    void setInverseMap(const std::vector<RGBQUAD>& palette) {
        memset(inverseMap, 0, sizeof(inverseMap));
        if (palette.empty()) return;

        for (int r = 1; r <= 32; ++r) 
            for (int g = 1; g <= 32; ++g) 
                for (int b = 1; b <= 32; ++b) 
                    inverseMap[r][g][b] = getPaletteIndex(r, g, b, palette);
    }
    BYTE getPaletteIndex(int rBin, int gBin, int bBin, 
                         const std::vector<RGBQUAD>& palette) const {
        const int r = (rBin << 3) - 4;
        const int g = (gBin << 3) - 4;
        const int b = (bBin << 3) - 4;
        int winner = 0;
        int minDist = INT_MAX;

        for (int i = 0; i < palette.size(); ++i) {
            const int dr = r - palette[i].rgbRed;
            const int dg = g - palette[i].rgbGreen;
            const int db = b - palette[i].rgbBlue;
            const int dist = dr * dr + dg * dg + db * db;
            if (dist < minDist) {
                minDist = dist; winner = i;
            }
        }
        return BYTE(winner);
    }
};
CRaster WuQuantizeTest(CRaster& raster, int nColors) {
    if (raster.IsEmpty() || raster.GetBPP() != 24) return CRaster();
    
    WuQuantizer wq;
    std::vector<RGBQUAD> palette = wq.processImage(raster, nColors);
    
    CRaster out;
    out.SetDimensions(raster.GetSize(), 24);
    CSize sz = out.GetSize();
    for (int y = 0; y < sz.cy; y++) {
        BYTE *p = (BYTE *)raster.GetLinePtr(y);
        BYTE *q = (BYTE *)out.GetLinePtr(y);
        for (int x = 0; x < sz.cx; x++, p += 3) {
            BYTE idx = wq.quantizeColor(p[2], p[1], p[0]);
            *q++ = palette[idx].rgbBlue;
            *q++ = palette[idx].rgbGreen;
            *q++ = palette[idx].rgbRed;
        }
    }
    return out;
}