Use a median cut for better gifs

1 files changed, 146 insertions, 2 deletions

diff --git a/convert.go b/convert.go
index 2d79be2..acb9747 100644
--- a/convert.go
+++ b/convert.go
@@ -5,10 +5,13 @@ import (
 	"flag"
 	"fmt"
 	"image"
+	"image/color"
 	"image/gif"
 	"image/jpeg"
 	"image/png"
+	"math"
 	"os"
+	"slices"
 	"strings"
 
 	"golang.org/x/image/bmp"
@@ -43,10 +46,151 @@ func load(filename string) (image.Image, error) {
 	return img, nil
 }
 
-func divRound(n, d int) int {
+func divRound[N ~int | ~int64](n, d N) N {
 	return (n + d/2) / d
 }
 
+func sign(n int) float64 {
+	if n < 0 {
+		return -1
+	}
+	if n > 0 {
+		return 1
+	}
+	return 0
+}
+
+func floydRivest[S ~[]E, E any](array S, left, right, k int, cmp func(E, E) int) {
+	for right > left {
+		if right-left > 600 {
+			n := right - left + 1
+			i := k - left + 1
+			z := math.Log(float64(n))
+			s := .5 * math.Exp(2*z/3)
+			sd := .5 * math.Sqrt(z*s*(float64(n)-s)/float64(n)) * sign(i-n/2)
+			newLeft := max(left, int(float64(k)-float64(i)*s/float64(n)+sd))
+			newRight := min(right, int(float64(k)+float64(n-i)*s/float64(n)+sd))
+			floydRivest(array, newLeft, newRight, k, cmp)
+		}
+		t := array[k]
+		i := left
+		j := right
+		array[left], array[k] = array[k], array[left]
+		if cmp(array[right], t) > 0 {
+			array[right], array[left] = array[left], array[right]
+		}
+		for i < j {
+			array[i], array[j] = array[j], array[i]
+			i++
+			j--
+			for ; cmp(array[i], t) < 0; i++ {
+			}
+			for ; cmp(array[j], t) > 0; j-- {
+			}
+		}
+		if cmp(array[left], t) == 0 {
+			array[left], array[j] = array[j], array[left]
+		} else {
+			j++
+			array[j], array[right] = array[right], array[j]
+		}
+		if j <= k {
+			left = j + 1
+		}
+		if k <= j {
+			right = j - 1
+		}
+	}
+}
+
+func quickSelect[S ~[]E, E any](list S, k int, cmp func(E, E) int) {
+	floydRivest(list, 0, len(list)-1, k, cmp)
+}
+
+func bucketRange(colors []color.RGBA) color.RGBA {
+	if len(colors) == 0 {
+		return color.RGBA{}
+	}
+	var minR, minG, minB uint8 = math.MaxUint8, math.MaxUint8, math.MaxUint8
+	var maxR, maxG, maxB uint8
+	for _, c := range colors {
+		minR, maxR = min(minR, c.R), max(maxR, c.R)
+		minG, maxG = min(minG, c.G), max(maxG, c.G)
+		minB, maxB = min(minB, c.B), max(maxB, c.B)
+	}
+	return color.RGBA{R: maxR - minR, G: maxG - minG, B: maxB - minB}
+}
+
+func cutOnce(colors []color.RGBA, bucketRange color.RGBA) [2][]color.RGBA {
+	if len(colors) == 0 {
+		return [...][]color.RGBA{colors, colors}
+	}
+	rRange, gRange, bRange := bucketRange.R, bucketRange.G, bucketRange.B
+	if rRange >= gRange && rRange >= bRange {
+		quickSelect(colors, len(colors)/2, func(x, y color.RGBA) int { return int(x.R) - int(y.R) })
+	} else if gRange >= rRange && gRange >= bRange {
+		quickSelect(colors, len(colors)/2, func(x, y color.RGBA) int { return int(x.G) - int(y.G) })
+	} else {
+		quickSelect(colors, len(colors)/2, func(x, y color.RGBA) int { return int(x.B) - int(y.B) })
+	}
+	return [...][]color.RGBA{colors[:len(colors)/2], colors[len(colors)/2:]}
+}
+
+func colorAvg(colors []color.RGBA) color.RGBA {
+	var r, g, b int64
+	for _, c := range colors {
+		r += int64(c.R)
+		g += int64(c.G)
+		b += int64(c.B)
+	}
+	n := int64(len(colors))
+	return color.RGBA{R: uint8(divRound(r, n)), G: uint8(divRound(g, n)), B: uint8(divRound(b, n)), A: 0xff}
+}
+
+func medianCut(palette color.Palette, img image.Image, n int) color.Palette {
+	var colors []color.RGBA
+	for y := img.Bounds().Min.Y; y < img.Bounds().Max.Y; y++ {
+		for x := img.Bounds().Min.X; x < img.Bounds().Max.X; x++ {
+			r, g, b, a := img.At(x, y).RGBA()
+			if a > 0 {
+				colors = append(colors, color.RGBA{R: uint8(r >> 8), G: uint8(g >> 8), B: uint8(b >> 8), A: 0xff})
+			}
+		}
+	}
+	buckets := [][]color.RGBA{colors}
+	bucketRanges := []color.RGBA{{}}
+	for {
+		var bestRange uint8
+		var bestIdx int
+		for i, rng := range bucketRanges {
+			r := max(rng.R, rng.G, rng.B)
+			if r >= bestRange {
+				bestRange = r
+				bestIdx = i
+			}
+		}
+		split := cutOnce(buckets[bestIdx], bucketRanges[bestIdx])
+		buckets = slices.Replace(buckets, bestIdx, bestIdx+1, split[:]...)
+		if len(buckets) == n-1 {
+			break
+		}
+		bucketRanges = slices.Replace(bucketRanges, bestIdx, bestIdx+1, bucketRange(split[0]), bucketRange(split[1]))
+	}
+	palette = append(palette, color.Transparent)
+	for _, b := range buckets {
+		if len(b) > 0 {
+			palette = append(palette, colorAvg(b))
+		}
+	}
+	return palette
+}
+
+type medianCutQuantizer struct{}
+
+func (medianCutQuantizer) Quantize(p color.Palette, m image.Image) color.Palette {
+	return medianCut(p, m, cap(p)-len(p))
+}
+
 func resize(img image.Image) image.Image {
 	x, y := *x, *y
 	if x == 0 {
@@ -111,7 +255,7 @@ func run() error {
 	defer outputFile.Close()
 	switch ext {
 	case "gif":
-		err = gif.Encode(outputFile, input, nil)
+		err = gif.Encode(outputFile, input, &gif.Options{Quantizer: medianCutQuantizer{}})
 	case "jpeg", "jpg":
 		err = jpeg.Encode(outputFile, input, &jpeg.Options{Quality: 100})
 	case "png":