Task: Color quantization




Go: package main

import (
    "container/heap"
    "image"
    "image/color"
    "image/png"
    "log"
    "math"
    "os"
    "sort"
)

func main() {
    f, err := os.Open("Quantum_frog.png")
    if err != nil {
        log.Fatal(err)
    }
    img, err := png.Decode(f)
    if ec := f.Close(); err != nil {
        log.Fatal(err)
    } else if ec != nil {
        log.Fatal(ec)
    }
    fq, err := os.Create("frog16.png")
    if err != nil {
        log.Fatal(err)
    }
    if err = png.Encode(fq, quant(img, 16)); err != nil {
        log.Fatal(err)
    }
}


func quant(img image.Image, nq int) image.Image {
    qz := newQuantizer(img, nq) 
    qz.cluster()                
    return qz.Paletted()        
}


type quantizer struct {
    img image.Image 
    cs  []cluster   
    px  []point     
    ch  chValues    
    eq  []point     
}

type cluster struct {
    px       []point 
    widestCh int     
    chRange  uint32  
}

type point struct{ x, y int }
type chValues []uint32
type queue []*cluster

const (
    rx = iota
    gx
    bx
)

func newQuantizer(img image.Image, nq int) *quantizer {
    b := img.Bounds()
    npx := (b.Max.X - b.Min.X) * (b.Max.Y - b.Min.Y)
    
    qz := &quantizer{
        img: img,
        ch:  make(chValues, npx),
        cs:  make([]cluster, nq),
    }
    
    c := &qz.cs[0]
    px := make([]point, npx)
    c.px = px
    i := 0
    for y := b.Min.Y; y < b.Max.Y; y++ {
        for x := b.Min.X; x < b.Max.X; x++ {
            px[i].x = x
            px[i].y = y
            i++
        }
    }
    return qz
}

func (qz *quantizer) cluster() {
    
    
    
    
    
    pq := new(queue)
    
    c := &qz.cs[0]
    for i := 1; ; {
        qz.setColorRange(c)
        
        
        if c.chRange > 0 {
            heap.Push(pq, c) 
        }
        
        
        if len(*pq) == 0 {
            qz.cs = qz.cs[:i]
            break
        }
        s := heap.Pop(pq).(*cluster) 
        c = &qz.cs[i]                
        i++
        m := qz.Median(s)
        qz.Split(s, c, m) 
        
        if i == len(qz.cs) {
            break
        }
        qz.setColorRange(s)
        if s.chRange > 0 {
            heap.Push(pq, s) 
        }
    }
}
    
func (q *quantizer) setColorRange(c *cluster) {
    
    var maxR, maxG, maxB uint32
    minR := uint32(math.MaxUint32)
    minG := uint32(math.MaxUint32)
    minB := uint32(math.MaxUint32) 
    for _, p := range c.px {
        r, g, b, _ := q.img.At(p.x, p.y).RGBA()
        if r < minR { 
            minR = r
        }
        if r > maxR {
            maxR = r
        }
        if g < minG {
            minG = g 
        }
        if g > maxG {
            maxG = g
        }
        if b < minB {
            minB = b
        }
        if b > maxB {
            maxB = b
        }
    }
    
    s := gx
    min := minG
    max := maxG
    if maxR-minR > max-min {
        s = rx
        min = minR
        max = maxR
    }
    if maxB-minB > max-min {
        s = bx
        min = minB
        max = maxB
    }
    c.widestCh = s
    c.chRange = max - min 
}

func (q *quantizer) Median(c *cluster) uint32 {
    px := c.px
    ch := q.ch[:len(px)]
    
    switch c.widestCh {
    case rx:
        for i, p := range c.px {
            ch[i], _, _, _ = q.img.At(p.x, p.y).RGBA()
        }
    case gx:
        for i, p := range c.px {
            _, ch[i], _, _ = q.img.At(p.x, p.y).RGBA()
        }
    case bx:
        for i, p := range c.px {
            _, _, ch[i], _ = q.img.At(p.x, p.y).RGBA()
        }
    }
    
    sort.Sort(ch)
    half := len(ch) / 2
    m := ch[half]
    if len(ch)%2 == 0 {
        m = (m + ch[half-1]) / 2
    }
    return m
}

func (q *quantizer) Split(s, c *cluster, m uint32) {
    px := s.px
    var v uint32
    i := 0
    lt := 0
    gt := len(px) - 1
    eq := q.eq[:0] 
    for i <= gt {
        
        r, g, b, _ := q.img.At(px[i].x, px[i].y).RGBA()
        switch s.widestCh {
        case rx:
            v = r
        case gx:
            v = g
        case bx:
            v = b
        } 
        
        switch {
        case v < m:
            px[lt] = px[i]
            lt++
            i++
        case v > m:
            px[gt], px[i] = px[i], px[gt]
            gt--
        default:
            eq = append(eq, px[i])
            i++
        }
    }
    
    if len(eq) > 0 {
        copy(px[lt:], eq) 
        
        
        
        if len(px)-i < lt {
            i = lt
        }
        q.eq = eq 
    }
    
    s.px = px[:i]
    c.px = px[i:]
}   
    
func (qz *quantizer) Paletted() *image.Paletted {
    cp := make(color.Palette, len(qz.cs))
    pi := image.NewPaletted(qz.img.Bounds(), cp)
    for i := range qz.cs {
        px := qz.cs[i].px
        
        var rsum, gsum, bsum int64
        for _, p := range px {
            r, g, b, _ := qz.img.At(p.x, p.y).RGBA()
            rsum += int64(r)
            gsum += int64(g)
            bsum += int64(b)
        } 
        n64 := int64(len(px))
        cp[i] = color.NRGBA64{
            uint16(rsum / n64),
            uint16(gsum / n64),
            uint16(bsum / n64),
            0xffff,
        }
        
        for _, p := range px {
            pi.SetColorIndex(p.x, p.y, uint8(i))
        }
    }
    return pi
}


func (c chValues) Len() int           { return len(c) }
func (c chValues) Less(i, j int) bool { return c[i] < c[j] }
func (c chValues) Swap(i, j int)      { c[i], c[j] = c[j], c[i] }


func (q queue) Len() int { return len(q) }


func (q queue) Less(i, j int) bool {
    return len(q[j].px) < len(q[i].px)
}

func (q queue) Swap(i, j int) {
    q[i], q[j] = q[j], q[i]
}
func (pq *queue) Push(x interface{}) {
    c := x.(*cluster)
    *pq = append(*pq, c)
}
func (pq *queue) Pop() interface{} {
    q := *pq
    n := len(q) - 1
    c := q[n]
    *pq = q[:n]
    return c
}