K-Means Color Quantization

K-means 컬러 양자화는 이미지의 화질을 최대한 유지하면서, 이미지에 사용된 고유한 색상의 수를 원하는 개수(K개)로 대폭 줄이는 영상 처리 기법이다. 일반적인 디지털 이미지는 24비트 RGB 인코딩을 사용하여 최대 1,670만 개의 고유 색상을 사용할 수 있는데, K-means 알고리즘을 적용하면 이 색상 공간을 8개, 16개, 64개 등 사용자가 지정한 K개의 대표 색상(컬러 팔레트)으로 압축할 수 있다. 알고리즘은 4단계의 과정으로 구성되어 있다.

1. 이미지의 모든 픽셀을 3차원 색상 공간(Red, Green, Blue 축)의 벡터로 인식
2. Clustering: 임의로 K개의 중심점을 잡은 뒤, 모든 픽셀을 가장 가까운 중심점에 할당
3. 각 그룹에 속한 픽셀들의 평균 색상을 계산하여 새로운 중심점으로 업데이트. 이 과정을 색상 변화가 없을 때까지 반복
4. Quantization: 최종적으로 수렴한 K개의 대표 색상 중, 각 원래 픽셀과 가장 가까운 대표 색상으로 픽셀 값을 치환

K-means 양자화는 잘 알려진 Median-cut이나 Octree 알고리즘에 비해 화질은 매우 우수하지만, 연산량이 많아 속도는 상당히 느린 편이다.

struct ColorItem {
    DWORD key;      // = (r<<16)|(g<<8)|b;
    int count;
    BYTE r, g, b;
    ColorItem(DWORD k, int c, BYTE _r, BYTE _g, BYTE _b)
        : key(k), count(c), r(_r), g(_g), b(_b) {}
};
// K-Means 알고리즘 자체는 단순하지만, 전체 픽셀을 돌면서 군집하므로 속도가 느림;
// 속도를 계선하기 위해서 영상 속에 나타난 unique color에 대한 clustering으로 변환;
CRaster KMeansQuantizer(const CRaster& raster, int nColors) {
    if (raster.GetBPP() != 24) return CRaster();
    CSize sz = raster.GetSize();
	nColors = max(1, nColors);
    
    // image에 포함된 unique color에 대해서 clustering을 시도;
    std::map<DWORD, int> colorMap;
    for (int y = 0; y < sz.cy; y++) {
        BYTE *p = (BYTE *)raster.GetLinePtr(y);
        for (int x = 0; x < sz.cx; x++, p += 3)
            colorMap[(p[2]<<16)|(p[1]<<8)|p[0]]++;
    }
    
    if (colorMap.size() <= nColors) return raster;
    std::vector<ColorItem> colors;
    std::map<DWORD, int>::iterator it = colorMap.begin();
    for (; it != colorMap.end(); ++it) {
        DWORD key = it->first;
        BYTE r = (BYTE)((key >> 16) & 0xFF);
        BYTE g = (BYTE)((key >> 8) & 0xFF);
        BYTE b = (BYTE)(key & 0xFF);
        colors.push_back(ColorItem(key, it->second, r, g, b));
    }

    // 설정;
    const int K = nColors;        // 축소할 컬러 수
    const int maxSteps = 20;      // 최대 반복 횟수
    
    std::vector<Pixel> centroids(K);   // 대표 컬러;
    std::vector<int> labels(colors.size()); 

    // Initial Centroids 설정 (균등 간격 샘플링)
    int step_size = colors.size() / K;
    for (int c = 0; c < K; ++c) {
        int idx = c * step_size;
        centroids[c].r = colors[idx].r;
        centroids[c].g = colors[idx].g;
        centroids[c].b = colors[idx].b;
    }

    // 각 컬러를 가장 가까운 centroid에 할당;
    for (int step = 0; step < maxSteps; ++step) {
        for (int i = colors.size(); i-->0;) {
            double minDist2 = 1e9;
            int winner = 0;
            for (int c = 0; c < K; ++c) {
                const Pixel& p = centroids[c];
                double dr = p.r - colors[i].r;
                double dg = p.g - colors[i].g;
                double db = p.b - colors[i].b;
                double dist2 = dr * dr + dg * dg + db * db;
                if (dist2 < minDist2) {
                    minDist2 = dist2; winner = c;
                }
            }
            labels[i] = winner;
        }

        // 새로운 중심점 계산;
        std::vector<Pixel> newCentroids(K, Pixel(0,0,0));
        std::vector<int> clusterSizes(K, 0);
        for (int i = colors.size(); i-->0;) {
            int c = labels[i];
            int count = colors[i].count;
            newCentroids[c].r += count * colors[i].r;
            newCentroids[c].g += count * colors[i].g;
            newCentroids[c].b += count * colors[i].b;
            clusterSizes[c]   += count;
        }

        // 중심점 업데이트;
        double centroidShift = 0.0;
        for (int c = 0; c < K; ++c) {
            if (clusterSizes[c] > 0) {
                Pixel updated(
                    newCentroids[c].r / clusterSizes[c],
                    newCentroids[c].g / clusterSizes[c],
                    newCentroids[c].b / clusterSizes[c]
                );
                double dr = updated.r - centroids[c].r;
                double dg = updated.g - centroids[c].g;
                double db = updated.b - centroids[c].b;
                centroids[c] = updated;
                centroidShift += sqrt(dr * dr + dg * dg + db * db);
            }
        }
        if ((centroidShift / K) < 0.05) break; // 차이가 작으면;
    }
    // palette 생성;
    std::vector<RGBQUAD> palette(centroids.size());
    for (int i = centroids.size(); i-->0;) {
        const Pixel& cc = centroids[i];
        palette[i].rgbBlue  = max(0, min(255, int(cc.b + 0.5)));
        palette[i].rgbGreen = max(0, min(255, int(cc.g + 0.5)));
        palette[i].rgbRed   = max(0, min(255, int(cc.r + 0.5)));
    }
    // Color(key) to palette index LUT 생성(colorMap 재사용)
    for (int i = colors.size(); i-->0; )
        colorMap[colors[i].key] = labels[i];

    // 줄어든 색상의 24-비트 영상 출력;
    CRaster out;
    out.SetDimensions(sz, 24);
    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) {    
            int pid = colorMap[(p[2]<<16)|(p[1]<<8)|p[0]];
            *q++ = palette[pid].rgbBlue;
            *q++ = palette[pid].rgbGreen;
            *q++ = palette[pid].rgbRed;
        }
    }
    return out;
}