Files
webp/api/kMeans.go
2024-11-14 19:31:59 +08:00

132 lines
2.9 KiB
Go

package api
import (
"image"
"image/color"
"math"
"math/rand"
"os"
)
// 定义一个结构来表示 RGB 颜色
type RGB struct {
R, G, B int
}
// 计算两个 RGB 颜色的欧氏距离
func distance(c1, c2 RGB) float64 {
return math.Sqrt(float64((c1.R-c2.R)*(c1.R-c2.R) + (c1.G-c2.G)*(c1.G-c2.G) + (c1.B-c2.B)*(c1.B-c2.B)))
}
// 对图像颜色进行 KMeans 聚类
func KMeans(colors []RGB, k int) ([]RGB, []int) {
// 随机选择初始中心
var centers []RGB
for i := 0; i < k; i++ {
centers = append(centers, colors[rand.Intn(len(colors))])
}
// 聚类算法
labels := make([]int, len(colors))
for iter := 0; iter < 10; iter++ { // 设置最大迭代次数为 10
// 为每个颜色分配最近的聚类中心
for i, color := range colors {
minDist := distance(color, centers[0])
labels[i] = 0
for j := 1; j < k; j++ {
dist := distance(color, centers[j])
if dist < minDist {
minDist = dist
labels[i] = j
}
}
}
// 计算新的聚类中心
newCenters := make([]RGB, k)
counts := make([]int, k)
for i, label := range labels {
newCenters[label].R += colors[i].R
newCenters[label].G += colors[i].G
newCenters[label].B += colors[i].B
counts[label]++
}
// 计算平均值作为新的中心
for i := 0; i < k; i++ {
if counts[i] > 0 {
newCenters[i].R /= counts[i]
newCenters[i].G /= counts[i]
newCenters[i].B /= counts[i]
}
}
// 如果中心没有变化,停止迭代
changed := false
for i := 0; i < k; i++ {
if newCenters[i] != centers[i] {
changed = true
break
}
}
if !changed {
break
}
centers = newCenters
}
return centers, labels
}
// 加载图像并提取颜色
func extractColors(img image.Image) (colors []RGB) {
bounds := img.Bounds()
for y := bounds.Min.Y; y < bounds.Max.Y; y++ {
for x := bounds.Min.X; x < bounds.Max.X; x++ {
r, g, b, _ := img.At(x, y).RGBA()
colors = append(colors, RGB{int(r >> 8), int(g >> 8), int(b >> 8)})
}
}
return colors
}
// 将聚类后的颜色应用到图像上
func recolorImage(img image.Image, centers []RGB, labels []int) image.Image {
bounds := img.Bounds()
newImage := image.NewRGBA(bounds)
for y := bounds.Min.Y; y < bounds.Max.Y; y++ {
for x := bounds.Min.X; x < bounds.Max.X; x++ {
label := labels[(y-bounds.Min.Y)*bounds.Max.X+(x-bounds.Min.X)]
// 将 RGB 颜色转换为 RGBA 类型
newColor := color.RGBA{
R: uint8(centers[label].R),
G: uint8(centers[label].G),
B: uint8(centers[label].B),
A: 255, // 不透明
}
// 设置新图像像素颜色
newImage.Set(x, y, newColor)
}
}
return newImage
}
func ImageToColors(str string) (colors []RGB, err error) {
// 打开图像文件
file, err := os.Open(str)
if err != nil {
return nil, err
}
defer file.Close()
// 解码 JPEG 图像
img, _, err := image.Decode(file)
if err != nil {
return nil, err
}
// 提取图像颜色
centers, _ := KMeans(extractColors(img), 8)
return centers, nil
}