summaryrefslogtreecommitdiffstats
path: root/vendor/github.com/Benau/tgsconverter/libtgsconverter/quantize_mediancut.go
blob: 850708b9b4df82bd6819b6e1d7785d09791dd991 (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
package libtgsconverter

import (
	"image"
	"image/color"
	"sync"
)

type bucketPool struct {
	sync.Pool
	maxCap int
	m      sync.Mutex
}

func (p *bucketPool) getBucket(c int) colorBucket {
	p.m.Lock()
	if p.maxCap > c {
		p.maxCap = p.maxCap * 99 / 100
	}
	if p.maxCap < c {
		p.maxCap = c
	}
	maxCap := p.maxCap
	p.m.Unlock()
	val := p.Pool.Get()
	if val == nil || cap(val.(colorBucket)) < c {
		return make(colorBucket, maxCap)[0:c]
	}
	slice := val.(colorBucket)
	slice = slice[0:c]
	for i := range slice {
		slice[i] = colorPriority{}
	}
	return slice
}

var bpool bucketPool

// aggregationType specifies the type of aggregation to be done
type aggregationType uint8

const (
	// Mode - pick the highest priority value
	mode aggregationType = iota
	// Mean - weighted average all values
	mean
)

// medianCutQuantizer implements the go draw.Quantizer interface using the Median Cut method
type medianCutQuantizer struct {
	// The type of aggregation to be used to find final colors
	aggregation aggregationType
	// The weighting function to use on each pixel
	weighting func(image.Image, int, int) uint32
	// Whether need to add a transparent entry after conversion
	reserveTransparent bool
}

//bucketize takes a bucket and performs median cut on it to obtain the target number of grouped buckets
func bucketize(colors colorBucket, num int) (buckets []colorBucket) {
	if len(colors) == 0 || num == 0 {
		return nil
	}
	bucket := colors
	buckets = make([]colorBucket, 1, num*2)
	buckets[0] = bucket

	for len(buckets) < num && len(buckets) < len(colors) { // Limit to palette capacity or number of colors
		bucket, buckets = buckets[0], buckets[1:]
		if len(bucket) < 2 {
			buckets = append(buckets, bucket)
			continue
		} else if len(bucket) == 2 {
			buckets = append(buckets, bucket[:1], bucket[1:])
			continue
		}

		left, right := bucket.partition()
		buckets = append(buckets, left, right)
	}
	return
}

// palettize finds a single color to represent a set of color buckets
func (q* medianCutQuantizer) palettize(p color.Palette, buckets []colorBucket) color.Palette {
	for _, bucket := range buckets {
		switch q.aggregation {
		case mean:
			mean := bucket.mean()
			p = append(p, mean)
		case mode:
			var best colorPriority
			for _, c := range bucket {
				if c.p > best.p {
					best = c
				}
			}
			p = append(p, best.RGBA)
		}
	}
	return p
}

// quantizeSlice expands the provided bucket and then palettizes the result
func (q* medianCutQuantizer) quantizeSlice(p color.Palette, colors []colorPriority) color.Palette {
	numColors := cap(p) - len(p)
	reserveTransparent := q.reserveTransparent
	if reserveTransparent {
		numColors--
	}
	buckets := bucketize(colors, numColors)
	p = q.palettize(p, buckets)
	return p
}

func colorAt(m image.Image, x int, y int) color.RGBA {
	switch i := m.(type) {
	case *image.YCbCr:
		yi := i.YOffset(x, y)
		ci := i.COffset(x, y)
		c := color.YCbCr{
			i.Y[yi],
			i.Cb[ci],
			i.Cr[ci],
		}
		return color.RGBA{c.Y, c.Cb, c.Cr, 255}
	case *image.RGBA:
		ci := i.PixOffset(x, y)
		return color.RGBA{i.Pix[ci+0], i.Pix[ci+1], i.Pix[ci+2], i.Pix[ci+3]}
	default:
		return color.RGBAModel.Convert(i.At(x, y)).(color.RGBA)
	}
}

// buildBucketMultiple creates a prioritized color slice with all the colors in
// the images.
func (q* medianCutQuantizer) buildBucketMultiple(ms []image.Image) (bucket colorBucket) {
	if len(ms) < 1 {
		return colorBucket{}
	}

	bounds := ms[0].Bounds()
	size := (bounds.Max.X - bounds.Min.X) * (bounds.Max.Y - bounds.Min.Y) * 2
	sparseBucket := bpool.getBucket(size)

	for _, m := range ms {
		for y := bounds.Min.Y; y < bounds.Max.Y; y++ {
			for x := bounds.Min.X; x < bounds.Max.X; x++ {
				priority := uint32(1)
				if q.weighting != nil {
					priority = q.weighting(m, x, y)
				}
				c := colorAt(m, x, y)
				if c.A == 0 {
					if !q.reserveTransparent {
						q.reserveTransparent = true
					}
					continue
				}
				if priority != 0 {
					index := int(c.R)<<16 | int(c.G)<<8 | int(c.B)
					for i := 1; ; i++ {
						p := &sparseBucket[index%size]
						if p.p == 0 || p.RGBA == c {
							*p = colorPriority{p.p + priority, c}
							break
						}
						index += 1 + i
					}
				}
			}
		}
	}

	bucket = sparseBucket[:0]
	switch ms[0].(type) {
	case *image.YCbCr:
		for _, p := range sparseBucket {
			if p.p != 0 {
				r, g, b := color.YCbCrToRGB(p.R, p.G, p.B)
				bucket = append(bucket, colorPriority{p.p, color.RGBA{r, g, b, p.A}})
			}
		}
	default:
		for _, p := range sparseBucket {
			if p.p != 0 {
				bucket = append(bucket, p)
			}
		}
	}
	return
}

// Quantize quantizes an image to a palette and returns the palette
func (q* medianCutQuantizer) quantize(p color.Palette, m image.Image) color.Palette {
	// Package quantize offers an implementation of the draw.Quantize interface using an optimized Median Cut method,
	// including advanced functionality for fine-grained control of color priority
	bucket := q.buildBucketMultiple([]image.Image{m})
	defer bpool.Put(bucket)
	return q.quantizeSlice(p, bucket)
}

// QuantizeMultiple quantizes several images at once to a palette and returns
// the palette
func (q* medianCutQuantizer) quantizeMultiple(p color.Palette, m []image.Image) color.Palette {
	bucket := q.buildBucketMultiple(m)
	defer bpool.Put(bucket)
	return q.quantizeSlice(p, bucket)
}