1 package quantize
2
3 import "image/color"
4
5 type colorAxis uint8
6
7
8 const (
9 red colorAxis = iota
10 green
11 blue
12 )
13
14 type colorPriority struct {
15 p uint32
16 color.RGBA
17 }
18
19 func (c colorPriority) axis(span colorAxis) uint8 {
20 switch span {
21 case red:
22 return c.R
23 case green:
24 return c.G
25 default:
26 return c.B
27 }
28 }
29
30 type colorBucket []colorPriority
31
32 func (cb colorBucket) partition() (colorBucket, colorBucket) {
33 mean, span := cb.span()
34 left, right := 0, len(cb)-1
35 for left < right {
36 cb[left], cb[right] = cb[right], cb[left]
37 for cb[left].axis(span) < mean && left < right {
38 left++
39 }
40 for cb[right].axis(span) >= mean && left < right {
41 right--
42 }
43 }
44 if left == 0 {
45 return cb[:1], cb[1:]
46 }
47 if left == len(cb)-1 {
48 return cb[:len(cb)-1], cb[len(cb)-1:]
49 }
50 return cb[:left], cb[left:]
51 }
52
53 func (cb colorBucket) mean() color.RGBA {
54 var r, g, b uint64
55 var p uint64
56 for _, c := range cb {
57 p += uint64(c.p)
58 r += uint64(c.R) * uint64(c.p)
59 g += uint64(c.G) * uint64(c.p)
60 b += uint64(c.B) * uint64(c.p)
61 }
62 return color.RGBA{uint8(r / p), uint8(g / p), uint8(b / p), 255}
63 }
64
65 type constraint struct {
66 min uint8
67 max uint8
68 vals [256]uint64
69 }
70
71 func (c *constraint) update(index uint8, p uint32) {
72 if index < c.min {
73 c.min = index
74 }
75 if index > c.max {
76 c.max = index
77 }
78 c.vals[index] += uint64(p)
79 }
80
81 func (c *constraint) span() uint8 {
82 return c.max - c.min
83 }
84
85 func (cb colorBucket) span() (uint8, colorAxis) {
86 var R, G, B constraint
87 R.min = 255
88 G.min = 255
89 B.min = 255
90 var p uint64
91 for _, c := range cb {
92 R.update(c.R, c.p)
93 G.update(c.G, c.p)
94 B.update(c.B, c.p)
95 p += uint64(c.p)
96 }
97 var toCount *constraint
98 var span colorAxis
99 if R.span() > G.span() && R.span() > B.span() {
100 span = red
101 toCount = &R
102 } else if G.span() > B.span() {
103 span = green
104 toCount = &G
105 } else {
106 span = blue
107 toCount = &B
108 }
109 var counted uint64
110 var i int
111 var c uint64
112 for i, c = range toCount.vals {
113 if counted > p/2 || counted+c == p {
114 break
115 }
116 counted += c
117 }
118 return uint8(i), span
119 }
120
View as plain text