138 lines
3.0 KiB
Go
138 lines
3.0 KiB
Go
package api
|
|
|
|
import (
|
|
"image"
|
|
"image/color"
|
|
"math"
|
|
"math/rand"
|
|
"os"
|
|
)
|
|
|
|
// 定义一个结构来表示 RGB 颜色
|
|
type RGB struct {
|
|
R, G, B int
|
|
}
|
|
|
|
// 根据每个聚类中心包含的颜色数量进行排序
|
|
type cluster struct {
|
|
center RGB
|
|
count 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
|
|
}
|