mediancut.go

Documentation: github.com/ericpauley/go-quantize/quantize

     1  // Package quantize offers an implementation of the draw.Quantize interface using an optimized Median Cut method,
     2  // including advanced functionality for fine-grained control of color priority
     3  package quantize
     4  
     5  import (
     6  	"image"
     7  	"image/color"
     8  	"sync"
     9  )
    10  
    11  type bucketPool struct {
    12  	sync.Pool
    13  	maxCap int
    14  	m      sync.Mutex
    15  }
    16  
    17  func (p *bucketPool) getBucket(c int) colorBucket {
    18  	p.m.Lock()
    19  	if p.maxCap > c {
    20  		p.maxCap = p.maxCap * 99 / 100
    21  	}
    22  	if p.maxCap < c {
    23  		p.maxCap = c
    24  	}
    25  	maxCap := p.maxCap
    26  	p.m.Unlock()
    27  	val := p.Pool.Get()
    28  	if val == nil || cap(val.(colorBucket)) < c {
    29  		return make(colorBucket, maxCap)[0:c]
    30  	}
    31  	slice := val.(colorBucket)
    32  	slice = slice[0:c]
    33  	for i := range slice {
    34  		slice[i] = colorPriority{}
    35  	}
    36  	return slice
    37  }
    38  
    39  var bpool bucketPool
    40  
    41  // AggregationType specifies the type of aggregation to be done
    42  type AggregationType uint8
    43  
    44  const (
    45  	// Mode - pick the highest priority value
    46  	Mode AggregationType = iota
    47  	// Mean - weighted average all values
    48  	Mean
    49  )
    50  
    51  // MedianCutQuantizer implements the go draw.Quantizer interface using the Median Cut method
    52  type MedianCutQuantizer struct {
    53  	// The type of aggregation to be used to find final colors
    54  	Aggregation AggregationType
    55  	// The weighting function to use on each pixel
    56  	Weighting func(image.Image, int, int) uint32
    57  	// Whether to create a transparent entry
    58  	AddTransparent bool
    59  }
    60  
    61  //bucketize takes a bucket and performs median cut on it to obtain the target number of grouped buckets
    62  func bucketize(colors colorBucket, num int) (buckets []colorBucket) {
    63  	if len(colors) == 0 || num == 0 {
    64  		return nil
    65  	}
    66  	bucket := colors
    67  	buckets = make([]colorBucket, 1, num*2)
    68  	buckets[0] = bucket
    69  
    70  	for len(buckets) < num && len(buckets) < len(colors) { // Limit to palette capacity or number of colors
    71  		bucket, buckets = buckets[0], buckets[1:]
    72  		if len(bucket) < 2 {
    73  			buckets = append(buckets, bucket)
    74  			continue
    75  		} else if len(bucket) == 2 {
    76  			buckets = append(buckets, bucket[:1], bucket[1:])
    77  			continue
    78  		}
    79  
    80  		left, right := bucket.partition()
    81  		buckets = append(buckets, left, right)
    82  	}
    83  	return
    84  }
    85  
    86  // palettize finds a single color to represent a set of color buckets
    87  func (q MedianCutQuantizer) palettize(p color.Palette, buckets []colorBucket) color.Palette {
    88  	for _, bucket := range buckets {
    89  		switch q.Aggregation {
    90  		case Mean:
    91  			mean := bucket.mean()
    92  			p = append(p, mean)
    93  		case Mode:
    94  			var best colorPriority
    95  			for _, c := range bucket {
    96  				if c.p > best.p {
    97  					best = c
    98  				}
    99  			}
   100  			p = append(p, best.RGBA)
   101  		}
   102  	}
   103  	return p
   104  }
   105  
   106  // quantizeSlice expands the provided bucket and then palettizes the result
   107  func (q MedianCutQuantizer) quantizeSlice(p color.Palette, colors []colorPriority) color.Palette {
   108  	numColors := cap(p) - len(p)
   109  	addTransparent := q.AddTransparent
   110  	if addTransparent {
   111  		for _, c := range p {
   112  			if _, _, _, a := c.RGBA(); a == 0 {
   113  				addTransparent = false
   114  			}
   115  		}
   116  		if addTransparent {
   117  			numColors--
   118  		}
   119  	}
   120  	buckets := bucketize(colors, numColors)
   121  	p = q.palettize(p, buckets)
   122  	if addTransparent {
   123  		p = append(p, color.RGBA{0, 0, 0, 0})
   124  	}
   125  	return p
   126  }
   127  
   128  func colorAt(m image.Image, x int, y int) color.RGBA {
   129  	switch i := m.(type) {
   130  	case *image.YCbCr:
   131  		yi := i.YOffset(x, y)
   132  		ci := i.COffset(x, y)
   133  		c := color.YCbCr{
   134  			i.Y[yi],
   135  			i.Cb[ci],
   136  			i.Cr[ci],
   137  		}
   138  		return color.RGBA{c.Y, c.Cb, c.Cr, 255}
   139  	case *image.RGBA:
   140  		ci := i.PixOffset(x, y)
   141  		return color.RGBA{i.Pix[ci+0], i.Pix[ci+1], i.Pix[ci+2], i.Pix[ci+3]}
   142  	default:
   143  		return color.RGBAModel.Convert(i.At(x, y)).(color.RGBA)
   144  	}
   145  }
   146  
   147  // buildBucket creates a prioritized color slice with all the colors in the image
   148  func (q MedianCutQuantizer) buildBucket(m image.Image) (bucket colorBucket) {
   149  	bounds := m.Bounds()
   150  	size := (bounds.Max.X - bounds.Min.X) * (bounds.Max.Y - bounds.Min.Y) * 2
   151  	sparseBucket := bpool.getBucket(size)
   152  
   153  	for y := bounds.Min.Y; y < bounds.Max.Y; y++ {
   154  		for x := bounds.Min.X; x < bounds.Max.X; x++ {
   155  			priority := uint32(1)
   156  			if q.Weighting != nil {
   157  				priority = q.Weighting(m, x, y)
   158  			}
   159  			if priority != 0 {
   160  				c := colorAt(m, x, y)
   161  				index := int(c.R)<<16 | int(c.G)<<8 | int(c.B)
   162  				for i := 1; ; i++ {
   163  					p := &sparseBucket[index%size]
   164  					if p.p == 0 || p.RGBA == c {
   165  						*p = colorPriority{p.p + priority, c}
   166  						break
   167  					}
   168  					index += 1 + i
   169  				}
   170  			}
   171  		}
   172  	}
   173  	bucket = sparseBucket[:0]
   174  	switch m.(type) {
   175  	case *image.YCbCr:
   176  		for _, p := range sparseBucket {
   177  			if p.p != 0 {
   178  				r, g, b := color.YCbCrToRGB(p.R, p.G, p.B)
   179  				bucket = append(bucket, colorPriority{p.p, color.RGBA{r, g, b, p.A}})
   180  			}
   181  		}
   182  	default:
   183  		for _, p := range sparseBucket {
   184  			if p.p != 0 {
   185  				bucket = append(bucket, p)
   186  			}
   187  		}
   188  	}
   189  	return
   190  }
   191  
   192  // Quantize quantizes an image to a palette and returns the palette
   193  func (q MedianCutQuantizer) Quantize(p color.Palette, m image.Image) color.Palette {
   194  	bucket := q.buildBucket(m)
   195  	defer bpool.Put(bucket)
   196  	return q.quantizeSlice(p, bucket)
   197  }
   198
View as plain text