kenaito-media-picker/internal/probe/brightness.go

package probe

import (
	"image"
	"os"

	_ "image/gif"
	_ "image/jpeg"
	_ "image/png"

	_ "golang.org/x/image/bmp"
	"golang.org/x/image/draw"
	_ "golang.org/x/image/tiff"
	_ "golang.org/x/image/webp"
)

// 明暗分析相关的三个参数：
const (
	// brightnessDecodeLimit 允许分析的最大像素数。超过就跳过分析（按中性归类），
	// 防止超大图把内存吃爆。
	brightnessDecodeLimit = 40_000_000
	// thumbnailMaxEdge 大图缩放后的最长边像素。先缩小再统计，省内存省时间。
	thumbnailMaxEdge = 256
	// thumbnailScaleThreshold 像素数超过这个值才缩放成缩略图再统计。
	thumbnailScaleThreshold = 250_000
)

// AnalyzeBrightness 打开一张图片，分析它偏暗还是偏亮。
//   - darkScore：暗色像素（亮度 < 80）占多少比例
//   - lightScore：亮色像素（亮度 > 175）占多少比例
//
// 超大图直接跳过（返回两个 0，调用方会按"中性"归类）；
// 普通大图先缩小成缩略图再统计，既快又不失真（整体明暗分布基本保留）。
func AnalyzeBrightness(filePath string) (darkScore float64, lightScore float64, err error) {
	f, err := os.Open(filePath)
	if err != nil {
		return 0, 0, err
	}
	defer f.Close()

	// 先只读配置看看多大，太大就不分析了
	cfg, _, err := image.DecodeConfig(f)
	if err != nil {
		return 0, 0, err
	}
	if cfg.Width <= 0 || cfg.Height <= 0 || cfg.Width*cfg.Height > brightnessDecodeLimit {
		return 0, 0, nil // 超大图，跳过
	}
	// 回到文件开头，重新完整解码
	if _, err := f.Seek(0, 0); err != nil {
		return 0, 0, err
	}
	img, _, err := image.Decode(f)
	if err != nil {
		return 0, 0, err
	}

	// 大图先缩小成缩略图
	if img.Bounds().Dx()*img.Bounds().Dy() > thumbnailScaleThreshold {
		img = scaleToThumbnail(img, thumbnailMaxEdge)
	}
	dark, light := analyzeBrightness(img)
	return dark, light, nil
}

// scaleToThumbnail 把大图等比缩放到最长边不超过 maxEdge 的缩略图（只占内存，不写盘）。
// 缩小时是像素平均，整体明暗分布基本保留，用来粗分亮暗足够了。
func scaleToThumbnail(img image.Image, maxEdge int) image.Image {
	srcBounds := img.Bounds()
	srcW, srcH := srcBounds.Dx(), srcBounds.Dy()
	if srcW <= 0 || srcH <= 0 {
		return img
	}
	ratio := float64(maxEdge) / float64(max(srcW, srcH))
	dstW, dstH := int(float64(srcW)*ratio), int(float64(srcH)*ratio)
	if dstW < 1 {
		dstW = 1
	}
	if dstH < 1 {
		dstH = 1
	}
	dst := image.NewRGBA(image.Rect(0, 0, dstW, dstH))
	draw.BiLinear.Scale(dst, dst.Bounds(), img, srcBounds, draw.Over, nil)
	return dst
}

// analyzeBrightness 数一数暗色/亮色像素各占多少比例。
// 用 ITU-R BT.601 的亮度公式算每个像素的"亮度值"。
// 完全透明的像素（alpha=0）不算数。
// 为了快，常见的内存格式（NRGBA/RGBA/Gray）直接翻底层像素数组，
// 不常见的才用慢一点的通用方式。
func analyzeBrightness(img image.Image) (darkScore float64, lightScore float64) {
	bounds := img.Bounds()
	totalPixels := bounds.Dx() * bounds.Dy()
	if totalPixels == 0 {
		return 0, 0
	}

	var darkCount, lightCount, sampleCount int64
	switch m := img.(type) {
	case *image.NRGBA:
		pix := m.Pix
		for y := bounds.Min.Y; y < bounds.Max.Y; y++ {
			offset := m.PixOffset(bounds.Min.X, y)
			for x := bounds.Min.X; x < bounds.Max.X; x++ {
				r, g, b, a := pix[offset], pix[offset+1], pix[offset+2], pix[offset+3]
				if a == 0 {
					offset += 4
					continue
				}
				countBrightness(float64(r), float64(g), float64(b), &darkCount, &lightCount)
				sampleCount++
				offset += 4
			}
		}
	case *image.RGBA:
		pix := m.Pix
		for y := bounds.Min.Y; y < bounds.Max.Y; y++ {
			offset := m.PixOffset(bounds.Min.X, y)
			for x := bounds.Min.X; x < bounds.Max.X; x++ {
				// RGBA 是"预乘 alpha"，完全不透明就用原值，
				// 半透明的按比例还原出真实颜色分量再算亮度
				a := pix[offset+3]
				if a == 0 {
					offset += 4
					continue
				}
				r, g, b := pix[offset], pix[offset+1], pix[offset+2]
				if a != 255 {
					r = uint8(int(r) * 255 / int(a))
					g = uint8(int(g) * 255 / int(a))
					b = uint8(int(b) * 255 / int(a))
				}
				countBrightness(float64(r), float64(g), float64(b), &darkCount, &lightCount)
				sampleCount++
				offset += 4
			}
		}
	case *image.Gray:
		// 灰度图：像素本身就是亮度值
		pix := m.Pix
		for y := bounds.Min.Y; y < bounds.Max.Y; y++ {
			offset := m.PixOffset(bounds.Min.X, y)
			for x := bounds.Min.X; x < bounds.Max.X; x++ {
				v := float64(pix[offset])
				if v < 80 {
					darkCount++
				} else if v > 175 {
					lightCount++
				}
				sampleCount++
				offset++
			}
		}
	case *image.YCbCr:
		// JPEG 解码出来是 YCbCr：Y 平面就是亮度，直接翻数组。
		// 这是最常见的照片格式，走这分支能避开 default 里逐像素 img.At().RGBA() 的慢路径。
		pix := m.Y
		for y := bounds.Min.Y; y < bounds.Max.Y; y++ {
			offset := y*m.YStride + bounds.Min.X
			for x := bounds.Min.X; x < bounds.Max.X; x++ {
				v := float64(pix[offset])
				if v < 80 {
					darkCount++
				} else if v > 175 {
					lightCount++
				}
				sampleCount++
				offset++
			}
		}
	default:
		// 其他格式：用通用的读法
		for y := bounds.Min.Y; y < bounds.Max.Y; y++ {
			for x := bounds.Min.X; x < bounds.Max.X; x++ {
				r, g, b, a := img.At(x, y).RGBA()
				if a == 0 {
					continue
				}
				countBrightness(float64(r>>8), float64(g>>8), float64(b>>8), &darkCount, &lightCount)
				sampleCount++
			}
		}
	}

	if sampleCount == 0 {
		return 0, 0
	}
	return float64(darkCount) / float64(sampleCount), float64(lightCount) / float64(sampleCount)
}

// countBrightness 按亮度公式（BT.601）判断一个像素偏暗还是偏亮，并计数。
func countBrightness(r, g, b float64, darkCount, lightCount *int64) {
	brightness := 0.299*r + 0.587*g + 0.114*b
	if brightness < 80 {
		*darkCount++
	} else if brightness > 175 {
		*lightCount++
	}
}