refactor: separate kmeans from color

Author: fumiama

color.go

@@ -3,8 +3,6 @@ package gg
 import (
 	"image"
 	"image/color"
-	"math"
-	"math/rand"
 	"unsafe"
 )
 
@@ -30,60 +28,15 @@ var (
 //
 // TakeThemeColorsKMeans 使用 k-means 算法从图像中提取 k 个主色。
 func TakeThemeColorsKMeans(img image.Image, k int) []color.RGBA {
-	rgbaimg := ImageToRGBA(img)
-	pixels := unsafe.Slice(
-		(*color.RGBA)(unsafe.Pointer(unsafe.SliceData(rgbaimg.Pix))),
-		uintptr(len(rgbaimg.Pix))/unsafe.Sizeof(color.RGBA{}),
-	)
-
-	// 初始化k个聚类中心
-	clusters := make([]color.RGBA, k)
-	for i := range k {
-		clusters[i] = pixels[rand.Intn(len(pixels))]
-	}
-
-	// 迭代聚类
+	ki := newKMeansImage(img, k) // 初始化k个聚类中心
 	for {
-		// 将每个像素点分配到最近的聚类中心
-		clusterAssignments := make([]int, len(pixels))
-		for i, pixel := range pixels {
-			minDistance := math.MaxFloat64
-			for j, cluster := range clusters {
-				distance := distance(pixel, cluster)
-				if distance < minDistance {
-					minDistance = distance
-					clusterAssignments[i] = j
-				}
-			}
-		}
-
-		// 计算每个聚类的新中心
-		newClusters := make([]color.RGBA, k)
-		for currentCluster := range k {
-			var r, g, b uint32
-			n := 0
-			for j, pixelCluster := range clusterAssignments {
-				if pixelCluster == currentCluster {
-					pixel := pixels[j]
-					r += uint32(pixel.R)
-					g += uint32(pixel.G)
-					b += uint32(pixel.B)
-					n++
-				}
-			}
-			if n != 0 {
-				newClusters[currentCluster] = color.RGBA{uint8(r / uint32(n)), uint8(g / uint32(n)), uint8(b / uint32(n)), 255}
-			}
-		}
-
-		// 如果聚类中心没有变化，则停止迭代
-		if isArrayRGBAEqual(clusters, newClusters) {
+		ki.assign()
+		ki.update()
+		if ki.epilogue() {
 			break
 		}
-		clusters = newClusters
 	}
-
-	return clusters
+	return ki.result()
 }
 
 // isArrayRGBAEqual compares two []color.RGBA is equal fastly.
@@ -115,12 +68,11 @@ func isArrayRGBAEqual(a, b []color.RGBA) bool {
 	return true
 }
 
-// 计算两个颜色之间的距离
-func distance(a, b color.RGBA) float64 {
-	return math.Sqrt(sq(float64(a.R)-float64(b.R)) + sq(float64(a.G)-float64(b.G)) + sq(float64(a.B)-float64(b.B)))
-}
-
-// 计算平方
-func sq(n float64) float64 {
-	return n * n
+// distanceRGBAsq calc between two color.RGBAs (RGB only, no sqrt to speedup)
+//
+// distanceRGBAsq 计算两个 color.RGBA 颜色之间的距离（仅 RGB，忽略开方以加速）
+func distanceRGBAsq(a, b color.RGBA) float64 {
+	return sq(float64(a.R)-float64(b.R)) +
+		sq(float64(a.G)-float64(b.G)) +
+		sq(float64(a.B)-float64(b.B))
 }

color_test.go

@@ -39,7 +39,7 @@ func twoColorImage(w, h int, c1, c2 color.RGBA) image.Image {
 // colorInSlice 判断给定颜色是否在切片中（允许 tolerance 误差）
 func colorInSlice(c color.RGBA, slice []color.RGBA, tolerance float64) bool {
 	for _, s := range slice {
-		if distance(c, s) <= tolerance {
+		if distanceRGBAsq(c, s) <= tolerance {
 			return true
 		}
 	}
@@ -100,15 +100,15 @@ func TestSq_LargeValue(t *testing.T) {
 
 func TestDistance_SameColor(t *testing.T) {
 	c := color.RGBA{100, 150, 200, 255}
-	if got := distance(c, c); got != 0 {
+	if got := distanceRGBAsq(c, c); got != 0 {
 		t.Errorf("distance(same, same) = %v, want 0", got)
 	}
 }
 
 func TestDistance_BlackAndWhite(t *testing.T) {
 	// sqrt(255^2 * 3) = 255 * sqrt(3)
 	want := 255 * math.Sqrt(3)
-	got := distance(Black, White)
+	got := math.Sqrt(distanceRGBAsq(Black, White))
 	if math.Abs(got-want) > 1e-9 {
 		t.Errorf("distance(black, white) = %v, want %v", got, want)
 	}
@@ -119,7 +119,7 @@ func TestDistance_SingleChannel(t *testing.T) {
 	b := color.RGBA{3, 4, 0, 255}
 	// sqrt(9 + 16) = 5
 	want := 5.0
-	got := distance(a, b)
+	got := math.Sqrt(distanceRGBAsq(a, b))
 	if math.Abs(got-want) > 1e-9 {
 		t.Errorf("distance(%v, %v) = %v, want %v", a, b, got, want)
 	}
@@ -128,7 +128,7 @@ func TestDistance_SingleChannel(t *testing.T) {
 func TestDistance_Symmetry(t *testing.T) {
 	a := color.RGBA{10, 20, 30, 255}
 	b := color.RGBA{50, 80, 110, 255}
-	if d1, d2 := distance(a, b), distance(b, a); math.Abs(d1-d2) > 1e-9 {
+	if d1, d2 := distanceRGBAsq(a, b), distanceRGBAsq(b, a); math.Abs(d1-d2) > 1e-9 {
 		t.Errorf("distance not symmetric: d(a,b)=%v, d(b,a)=%v", d1, d2)
 	}
 }
@@ -138,7 +138,7 @@ func TestDistance_NonNegative(t *testing.T) {
 	for range 100 {
 		a := color.RGBA{uint8(rng.Intn(256)), uint8(rng.Intn(256)), uint8(rng.Intn(256)), 255}
 		b := color.RGBA{uint8(rng.Intn(256)), uint8(rng.Intn(256)), uint8(rng.Intn(256)), 255}
-		got := distance(a, b)
+		got := distanceRGBAsq(a, b)
 		if got < 0 {
 			t.Errorf("distance returned negative value %v for %v, %v", got, a, b)
 		}
@@ -150,7 +150,7 @@ func TestDistance_IgnoresAlpha(t *testing.T) {
 	a2 := color.RGBA{100, 150, 200, 255}
 	b := color.RGBA{50, 50, 50, 128}
 	// Alpha 不参与计算，结果应相同
-	if d1, d2 := distance(a1, b), distance(a2, b); math.Abs(d1-d2) > 1e-9 {
+	if d1, d2 := distanceRGBAsq(a1, b), distanceRGBAsq(a2, b); math.Abs(d1-d2) > 1e-9 {
 		t.Errorf("distance should ignore alpha: d(a1,b)=%v, d(a2,b)=%v", d1, d2)
 	}
 }
@@ -320,7 +320,7 @@ func BenchmarkDistance(b *testing.B) {
 	c := color.RGBA{50, 80, 30, 255}
 	b.ResetTimer()
 	for range b.N {
-		distance(a, c)
+		distanceRGBAsq(a, c)
 	}
 }
 

effects.go

@@ -1,37 +1,17 @@
 package gg
 
-func gateN100P100(per int) int {
-	if per > 100 {
-		per = 100
-	}
-	if per < -100 {
-		per = -100
-	}
-	return per
-}
-
-func gate0P255(n int) int {
-	if n < 0 {
-		n = 0
-	}
-	if n > 255 {
-		n = 255
-	}
-	return n
-}
-
 // Brightness 调整亮度 范围：±100%
 func (dc *Context) Brightness(per int) {
 	if per == 0 {
 		return
 	}
-	per = gateN100P100(per)
+	per = clamp(per, -100, 100)
 	gain := 255 * per / 100
 	for i, v := range dc.im.Pix {
 		if i%4 == 3 { // alpha
 			continue
 		}
-		dc.im.Pix[i] = uint8(gate0P255(int(v) + gain))
+		dc.im.Pix[i] = uint8(clamp(int(v)+gain, 0, 255))
 	}
 }
 
@@ -40,15 +20,15 @@ func (dc *Context) Contrast(per int) {
 	if per == 0 {
 		return
 	}
-	per = gateN100P100(per) + 100
+	per = clamp(per, -100, 100) + 100
 	switch {
 	case 0 <= per && per < 100: // 损益
 		gain := per
 		for i, v := range dc.im.Pix {
 			if i%4 == 3 { // alpha
 				continue
 			}
-			dc.im.Pix[i] = uint8(gate0P255(int(v) * gain / 100))
+			dc.im.Pix[i] = uint8(clamp(int(v)*gain/100, 0, 255))
 		}
 	case 100 < per && per <= 200: // 增益
 		gain := 200 - per
@@ -59,7 +39,7 @@ func (dc *Context) Contrast(per int) {
 			if i%4 == 3 { // alpha
 				continue
 			}
-			dc.im.Pix[i] = uint8(gate0P255(int(v) * 100 / gain))
+			dc.im.Pix[i] = uint8(clamp(int(v)*100/gain, 0, 255))
 		}
 	default:
 		panic("unreachable")

examples/theme-colors_test.go

@@ -0,0 +1,98 @@
+package main
+
+import (
+	"image/color"
+	"math"
+	"testing"
+
+	"github.com/FloatTech/gg"
+	"github.com/FloatTech/gg/fio"
+)
+
+func TestThemeColorsGopher(t *testing.T) {
+	im, err := fio.LoadPNG("gopher.png")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	expected := []color.RGBA{
+		{105, 213, 226, 255}, {249, 240, 227, 255}, {4, 6, 6, 255},
+	}
+
+	const (
+		k           = 3
+		maxAttempts = 30
+		tolerance   = 600.0 // 允许的颜色距离平方
+	)
+
+	var result []color.RGBA
+	found := false
+	for range maxAttempts {
+		result = gg.TakeThemeColorsKMeans(im, k)
+		if len(result) != k {
+			t.Fatalf("expected %d colors, got %d", k, len(result))
+		}
+		if allColorsMatch(expected, result, tolerance) {
+			found = true
+			break
+		}
+	}
+	if !found {
+		t.Errorf("theme colors did not match expected in %d attempts\nexpected: %v\ngot:      %v", maxAttempts, expected, result)
+	}
+
+	// 绘制提取到的主题色方块
+	const (
+		blockW = 120
+		blockH = 120
+		pad    = 20
+	)
+	w := pad + k*(blockW+pad)
+	h := pad + blockH + pad
+	dc := gg.NewContext(w, h)
+
+	// 白色背景
+	dc.SetColor(color.White)
+	dc.Clear()
+
+	for i, c := range result {
+		x := float64(pad + i*(blockW+pad))
+		y := float64(pad)
+		dc.SetColor(c)
+		dc.DrawRectangle(x, y, blockW, blockH)
+		dc.Fill()
+	}
+
+	if err := dc.SavePNG(GetFileName() + ".png"); err != nil {
+		t.Fatal(err)
+	}
+}
+
+// allColorsMatch 检查 expected 中每个颜色都能在 result 中找到匹配（无序）
+func allColorsMatch(expected, result []color.RGBA, tolerance float64) bool {
+	if len(expected) != len(result) {
+		return false
+	}
+	used := make([]bool, len(result))
+	for _, e := range expected {
+		matched := false
+		for j, r := range result {
+			if !used[j] && colorDistSq(e, r) <= tolerance {
+				used[j] = true
+				matched = true
+				break
+			}
+		}
+		if !matched {
+			return false
+		}
+	}
+	return true
+}
+
+func colorDistSq(a, b color.RGBA) float64 {
+	dr := float64(a.R) - float64(b.R)
+	dg := float64(a.G) - float64(b.G)
+	db := float64(a.B) - float64(b.B)
+	return math.Abs(dr*dr) + math.Abs(dg*dg) + math.Abs(db*db)
+}

factory/img.go

@@ -11,7 +11,6 @@ import (
 	"image/gif"
 	"io"
 	"net/http"
-	"os"
 	"strings"
 
 	"github.com/FloatTech/gg"
@@ -89,15 +88,7 @@ func LoadAllFrames(path string, w, h int) ([]*Factory, error) {
 			return nil, err
 		}
 	} else {
-		file, err := os.Open(path)
-		if err != nil {
-			return nil, err
-		}
-		im, err = gif.DecodeAll(file)
-		_ = file.Close()
-		if err != nil {
-			return nil, err
-		}
+		im, err = fio.LoadGIF(path)
 	}
 	img, err := Load(path)
 	if err != nil {
@@ -127,15 +118,7 @@ func LoadAllTrueFrames(path string, w, h int) ([]*Factory, error) {
 			return nil, err
 		}
 	} else {
-		file, err := os.Open(path)
-		if err != nil {
-			return nil, err
-		}
-		im, err = gif.DecodeAll(file)
-		_ = file.Close()
-		if err != nil {
-			return nil, err
-		}
+		im, err = fio.LoadGIF(path)
 	}
 	imgWidth, imgHeight := getGifDimensions(im)
 	overpaintImage := image.NewRGBA(image.Rect(0, 0, imgWidth, imgHeight))