Neural Network Quantization

$$ \text{New_color}_j = \text{Old_color}_j + \alpha \times H(i, j) (\text{Color} - \text{Old_color}_j)$$

$$ H(i, j)  = \left\{   \begin{matrix} 1 & \text{if} ~| i - j  | \le R \\ 0 & \text{if} ~| i- j |>R   \end{matrix}\right.$$

/*
 * NeuQuant: Anthony Dekker's neural-network color quantizer.
 * Call ProcessImage() once, then QuantizeColor().
 */
class NeuQuant {
public:
    NeuQuant(int nColors, int sampleFactor);
    std::vector<RGBQUAD> ProcessImage(const CRaster& raster);
    int QuantizeColor(BYTE r, BYTE g, BYTE b) const;
private:
    enum {
        CYCLE_COUNT = 100,
    };
    struct Neuron {
        float b, g, r;
        int paletteIndex;
    };
    void InitNetwork();
    void Learn();
    int  Contest(int b, int g, int r);
    void AlterSingle(float alpha, int index, int b, int g, int r);
    void AlterNeighbours(int radius, int index, int b, int g, int r);
    void BuildIndex();
    std::vector<RGBQUAD> GetPalette() const;

    int                  m_ColorCount;
    int                  m_sampleFactor;
    int                  m_netIndex[256];
    std::vector<Neuron>  m_network;			
    std::vector<float>   m_bias;			
    std::vector<float>   m_frequency;		
    std::vector<float>   m_radiusPower;		
    std::vector<BYTE>    m_image;
    bool                 m_processed;
};

namespace {
const int PRIME1 = 499;
const int PRIME2 = 491;
const int PRIME3 = 487;
const int PRIME4 = 503;
const int   MIN_PICTURE_BYTES = 3 * PRIME4;
const float INITIAL_ALPHA     = 1.0f;
const float FREQUENCY_DECAY   = 1.0f / 1024.0f;
const float BETA              = 1.0f / 1024.0f;
const float GAMMA             = 1024.0f;
} // namespace
NeuQuant::NeuQuant(int nColors, int sampleFactor)
    : m_sampleFactor(sampleFactor < 1 ? 1 : sampleFactor),
      m_processed(false) {    
    m_ColorCount = max(1, min(256, nColors));
    
    m_network.resize(m_ColorCount);
    m_bias.resize(m_ColorCount);		
    m_frequency.resize(m_ColorCount);	
    m_radiusPower.resize(m_ColorCount >> 3);
}
// 다른 image에서 재사용할 수 있게;
std::vector<RGBQUAD> NeuQuant::ProcessImage(const CRaster& raster) {
    if (raster.IsEmpty() || raster.GetBPP() != 24) 
        return std::vector<RGBQUAD> ();
    
    m_processed = false;
    CSize sz = raster.GetSize();
    m_image.clear();
    m_image.reserve(sz.cx * sz.cy * 3);
    for (int y = 0; y < sz.cy; ++y) {
        const BYTE *q = (BYTE *)raster.GetLinePtr(y);
        for (int x = 0; x < sz.cx; x++) {
            m_image.push_back(*q++);
            m_image.push_back(*q++);
            m_image.push_back(*q++);
        }
    }
    
    InitNetwork();
    Learn();
    BuildIndex();
    m_processed = true;
    return GetPalette();
}
std::vector<RGBQUAD> NeuQuant::GetPalette() const { 
    std::vector<int> index(m_ColorCount);   
    std::vector<RGBQUAD> palette(m_ColorCount);

    for (int i = m_ColorCount; i-->0;) index[m_network[i].paletteIndex] = i;
    for (int i = m_ColorCount; i-->0;) {
        const Neuron& n = m_network[index[i]];
        palette[i].rgbBlue  = BYTE(max(0, min(255, n.b + 0.5f)));
        palette[i].rgbGreen = BYTE(max(0, min(255, n.g + 0.5f)));
        palette[i].rgbRed   = BYTE(max(0, min(255, n.r + 0.5f)));
    }	
    return palette; 
}

int NeuQuant::QuantizeColor(BYTE r, BYTE g, BYTE b) const {
    if (!m_processed) return -1;

    int best = -1;
    float bestDist = FLT_MAX;
    int i = m_netIndex[g];
    int j = i - 1;
    while (i < m_ColorCount || j >= 0) {
        if (i < m_ColorCount) {
            const Neuron& n = m_network[i];
            float dist = n.g - g;
            if (dist >= bestDist) i = m_ColorCount;
            else {
                ++i;
                dist = fabs(dist) + fabs(n.b - b);
                if (dist < bestDist) {
                    dist += fabs(n.r - r);
                    if (dist < bestDist) {
                        bestDist = dist;
                        best = n.paletteIndex;
                    }
                }
            }
        }
        if (j >= 0) {
            const Neuron& n = m_network[j];
            float dist = g - n.g;
            if (dist >= bestDist) j = -1;
            else {
                --j;
                dist = fabs(dist) + fabs(n.b - b);
                if (dist < bestDist) {
                    dist += fabs(n.r - r);
                    if (dist < bestDist) {
                        bestDist = dist;
                        best = n.paletteIndex;
                    }
                }
            }
        }
    }
    return best;
}
void NeuQuant::InitNetwork() {
    for (int i = m_ColorCount; i-->0;) {
    	float v = 255.0F * i / (m_ColorCount - 1);
        m_network[i].b = v;
        m_network[i].g = v;
        m_network[i].r = v;
        m_network[i].paletteIndex = i;
        m_frequency[i] = 1.0f / m_ColorCount;
        m_bias[i] = 0.0f;
    }
    m_processed = false;
}
void NeuQuant::Learn() {
    const int length = int(m_image.size());
    int factor = m_sampleFactor;
    float alpha = INITIAL_ALPHA;
    float radius = float(m_ColorCount>>3);
    int rad = int(radius);

    if (length < MIN_PICTURE_BYTES) factor = 1;
    int alphaDecay = 30 + (factor - 1) / 3;
    int nSamples = length / (3 * factor);
    int delta = nSamples / CYCLE_COUNT;
    if (delta == 0) delta = 1;

    if (rad <= 1) rad = 0;
    for (int k = 0; k < rad; ++k)
        m_radiusPower[k] = alpha * (1.0f - float(k * k) / (rad * rad));

    int step = 3;
    if (length < MIN_PICTURE_BYTES) step = 3;
    else if (length % PRIME1 != 0)  step = 3 * PRIME1;
    else if (length % PRIME2 != 0)  step = 3 * PRIME2;
    else if (length % PRIME3 != 0)  step = 3 * PRIME3;
    else                            step = 3 * PRIME4;

    for (int count = 0, pixel = 0; count < nSamples; ++count) {
        const int b = m_image[pixel + 0];
        const int g = m_image[pixel + 1];
        const int r = m_image[pixel + 2];
        const int winner = Contest(b, g, r);
        AlterSingle(alpha, winner, b, g, r);
        if (rad != 0) AlterNeighbours(rad, winner, b, g, r);

        pixel += step;
        if (pixel >= length) pixel -= length;

        if ((count + 1) % delta == 0) {
            alpha -= alpha / alphaDecay;
            radius -= radius / 30.0f;
            rad = int(radius);
            if (rad <= 1) rad = 0;
            for (int j = 0; j < rad; ++j)
                m_radiusPower[j] = alpha * (1.0f - float(j * j) / (rad * rad));
        }
    }
}
int NeuQuant::Contest(int b, int g, int r) {
    int bestPos = -1;
    int bestBiasedPos = -1;
    float bestDist = FLT_MAX;
    float bestBiasedDist = FLT_MAX;
    for (int i = 0; i < m_ColorCount; ++i) {
        const Neuron& n = m_network[i];
        const float dist = fabs(n.b - b) + fabs(n.g - g) + fabs(n.r - r);
        const float biasedDist = dist - m_bias[i];
        if (dist < bestDist) {
            bestDist = dist;
            bestPos = i;
        }
        if (biasedDist < bestBiasedDist) {
            bestBiasedDist = biasedDist;
            bestBiasedPos = i;
        }
        const float freqChange = BETA * m_frequency[i];
        m_frequency[i] -= freqChange;
        m_bias[i] += GAMMA * freqChange;
    }
    m_frequency[bestPos] += BETA;
    m_bias[bestPos] -= BETA * GAMMA;
    return bestBiasedPos;
}
void NeuQuant::AlterSingle(float alpha, int index, int b, int g, int r) {
    Neuron& n = m_network[index];
    n.b += alpha * (float(b) - n.b);
    n.g += alpha * (float(g) - n.g);
    n.r += alpha * (float(r) - n.r);
}
void NeuQuant::AlterNeighbours(int radius, int index, int b, int g, int r) {
    int lo = max(-1, index - radius);
    int hi = min(m_ColorCount, index + radius);
    int j = index + 1;
    int k = index - 1;
    int m = 1;
    while (j < hi || k > lo) {
        if (m >= m_radiusPower.size()) break;
        const float amount = m_radiusPower[m++];
        if (j < hi) {
            Neuron& n = m_network[j++];
            n.b += amount * (b - n.b);
            n.g += amount * (g - n.g);
            n.r += amount * (r - n.r);
        }
        if (k > lo) {
            Neuron& n = m_network[k--];
            n.b += amount * (b - n.b);
            n.g += amount * (g - n.g);
            n.r += amount * (r - n.r);
        }
    }
}
// 어느 색을 기준으로 해도 무방;
void NeuQuant::BuildIndex() {
    int prevColor = 0, startPos = 0;
    for (int i = 0; i < m_ColorCount; ++i) {
        int smallestPos = i;
        float smallestGreen = m_network[i].g;
        for (int j = i + 1; j < m_ColorCount; ++j) {
            if (m_network[j].g < smallestGreen) {
                smallestPos = j;
                smallestGreen = m_network[j].g;
            }
        }
        if (smallestPos != i) {
            const Neuron temp = m_network[i];
            m_network[i] = m_network[smallestPos];
            m_network[smallestPos] = temp;
        }
        const int green = max(0, min(255, int(m_network[i].g)));
        if (green != prevColor) {
            m_netIndex[prevColor] = (startPos + i) >> 1;
            for (int j = prevColor + 1; j < green; ++j) m_netIndex[j] = i;
            prevColor = green;
            startPos = i;
        }
    }
    m_netIndex[prevColor] = (startPos + m_ColorCount - 1) >> 1;
    for (int i = prevColor + 1; i < 256; ++i) m_netIndex[i] = m_ColorCount - 1;
}
CRaster NeuQuantizeTest(const CRaster& raster, int nColors, int sampleFactor) {
    CRaster out;
    NeuQuant nq(nColors, sampleFactor);
    const std::vector<RGBQUAD> palette = nq.ProcessImage(raster);
    CSize sz = raster.GetSize();
    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) {
            int idx = nq.QuantizeColor(p[2], p[1], p[0]);
            *q++ = palette[idx].rgbBlue;
            *q++ = palette[idx].rgbGreen;
            *q++ = palette[idx].rgbRed;
            p += 3;
        }
    }
    return out;
}