From b5430d90359fb508cc77f61919efdf7664ef9517 Mon Sep 17 00:00:00 2001
From: nuknal <linktolane@gmail.com>
Date: Mon, 11 Nov 2024 17:16:50 +0800
Subject: [PATCH] cloud cover by kmeans

---
 pkg/cloud-cover/k_means.go | 152 +++++++++++++++++++++++++++++++++++++
 pkg/fusion/fusion.go       |   4 +
 2 files changed, 156 insertions(+)
 create mode 100644 pkg/cloud-cover/k_means.go

diff --git a/pkg/cloud-cover/k_means.go b/pkg/cloud-cover/k_means.go
new file mode 100644
index 0000000..6102101
--- /dev/null
+++ b/pkg/cloud-cover/k_means.go
@@ -0,0 +1,152 @@
+package cloudcover
+
+import (
+	"image"
+	"image/color"
+
+	"math"
+	"math/rand"
+	"os"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+)
+
+const (
+	k         = 2   // 聚类数目，2表示云和非云
+	maxIters  = 100 // 最大迭代次数
+	tolerance = 1.0 // 聚类中心变化容忍度
+)
+
+type Pixel struct {
+	r, g, b float64
+}
+
+func CloudPercentByKMeans(imagePath string) float64 {
+	file, err := os.Open(imagePath)
+	if err != nil {
+		log.Errorf("无法打开图像: %v", err)
+		return 0.0
+	}
+	defer file.Close()
+
+	img, _, err := image.Decode(file)
+	if err != nil {
+		log.Errorf("无法解码图像: %v", err)
+		return 0.0
+	}
+
+	bounds := img.Bounds()
+	width, height := bounds.Max.X, bounds.Max.Y
+	pixels := make([]Pixel, 0, width*height)
+
+	for y := 0; y < height; y++ {
+		for x := 0; x < width; x++ {
+			r, g, b, _ := img.At(x, y).RGBA()
+			pixels = append(pixels, Pixel{
+				r: float64(r >> 8),
+				g: float64(g >> 8),
+				b: float64(b >> 8),
+			})
+		}
+	}
+
+	rand.Seed(time.Now().UnixNano())
+	centroids := initializeCentroids(pixels)
+	assignments := make([]int, len(pixels))
+
+	for iter := 0; iter < maxIters; iter++ {
+		for i, pixel := range pixels {
+			assignments[i] = nearestCentroid(pixel, centroids)
+		}
+
+		newCentroids := updateCentroids(pixels, assignments)
+		if hasConverged(centroids, newCentroids) {
+			log.Printf("在第 %d 次迭代后收敛\n", iter+1)
+			break
+		}
+		centroids = newCentroids
+	}
+
+	cloudPixels := 0
+	totalPixels := width * height
+
+	// 新增：检查哪个聚类的平均亮度更高，确保白色表示云
+	cloudCluster := 0
+	if centroids[1].r+centroids[1].g+centroids[1].b > centroids[0].r+centroids[0].g+centroids[0].b {
+		cloudCluster = 1
+	}
+
+	outputImg := image.NewGray(bounds)
+	for y := 0; y < height; y++ {
+		for x := 0; x < width; x++ {
+			i := y*width + x
+			if assignments[i] == cloudCluster {
+				outputImg.SetGray(x, y, color.Gray{Y: 255}) // 云区域
+				cloudPixels++
+			} else {
+				outputImg.SetGray(x, y, color.Gray{Y: 0}) // 非云区域
+			}
+		}
+	}
+
+	cloudRatio := float64(cloudPixels) / float64(totalPixels)
+	return cloudRatio
+}
+
+func initializeCentroids(pixels []Pixel) []Pixel {
+	centroids := make([]Pixel, k)
+	for i := range centroids {
+		centroids[i] = pixels[rand.Intn(len(pixels))]
+	}
+	return centroids
+}
+
+func nearestCentroid(pixel Pixel, centroids []Pixel) int {
+	minDist := math.Inf(1)
+	bestIndex := 0
+	for i, c := range centroids {
+		dist := math.Sqrt(math.Pow(pixel.r-c.r, 2) + math.Pow(pixel.g-c.g, 2) + math.Pow(pixel.b-c.b, 2))
+		if dist < minDist {
+			minDist = dist
+			bestIndex = i
+		}
+	}
+	return bestIndex
+}
+
+func updateCentroids(pixels []Pixel, assignments []int) []Pixel {
+	sums := make([]Pixel, k)
+	counts := make([]int, k)
+
+	for i, assignment := range assignments {
+		sums[assignment].r += pixels[i].r
+		sums[assignment].g += pixels[i].g
+		sums[assignment].b += pixels[i].b
+		counts[assignment]++
+	}
+
+	newCentroids := make([]Pixel, k)
+	for i := range newCentroids {
+		if counts[i] > 0 {
+			newCentroids[i].r = sums[i].r / float64(counts[i])
+			newCentroids[i].g = sums[i].g / float64(counts[i])
+			newCentroids[i].b = sums[i].b / float64(counts[i])
+		} else {
+			newCentroids[i] = sums[i]
+		}
+	}
+
+	return newCentroids
+}
+
+func hasConverged(oldCentroids, newCentroids []Pixel) bool {
+	for i := range oldCentroids {
+		if math.Abs(oldCentroids[i].r-newCentroids[i].r) > tolerance ||
+			math.Abs(oldCentroids[i].g-newCentroids[i].g) > tolerance ||
+			math.Abs(oldCentroids[i].b-newCentroids[i].b) > tolerance {
+			return false
+		}
+	}
+	return true
+}
diff --git a/pkg/fusion/fusion.go b/pkg/fusion/fusion.go
index 8974cc9..0bd9a68 100644
--- a/pkg/fusion/fusion.go
+++ b/pkg/fusion/fusion.go
@@ -24,6 +24,10 @@ const (
 	GDALPansharpen
 	IHS
 	GIHS
+	WeightR   = 1
+	WeightG   = 0.75
+	WeightB   = 0.55
+	WeightNIR = 0.85
 )
 
 func Pansharpen(panTif, mssTif, fusTif string, method PansharpenMethod, weight float32) error {