251 lines
6.6 KiB
Go
251 lines
6.6 KiB
Go
package imageproc
|
|
|
|
import (
|
|
"image"
|
|
"math"
|
|
"sort"
|
|
|
|
"materialanalyzer/internal/model"
|
|
)
|
|
|
|
type rgbSample struct {
|
|
r, g, b float64
|
|
lum float64
|
|
sat float64
|
|
}
|
|
|
|
type colorCluster struct {
|
|
center rgbSample
|
|
count int
|
|
}
|
|
|
|
func ExtractVisual(img image.Image) (model.VisualFeatures, model.ColorHistogram) {
|
|
samples := collectSamples(img, 65000)
|
|
hist := buildHSVHistogram(samples, 16, 4, 4)
|
|
if len(samples) == 0 {
|
|
return model.VisualFeatures{HistogramBins: 256, HistogramSpace: "hsv"}, hist
|
|
}
|
|
|
|
var sumR, sumG, sumB, sumLum, sumLumSq, sumSat float64
|
|
for _, s := range samples {
|
|
sumR += s.r
|
|
sumG += s.g
|
|
sumB += s.b
|
|
sumLum += s.lum
|
|
sumLumSq += s.lum * s.lum
|
|
sumSat += s.sat
|
|
}
|
|
n := float64(len(samples))
|
|
avg := model.RGB{R: int(math.Round(sumR / n)), G: int(math.Round(sumG / n)), B: int(math.Round(sumB / n))}
|
|
avgHSV := RGBToHSV(avg)
|
|
avgLAB := RGBToLAB(avg)
|
|
brightness := sumLum / n / 255.0
|
|
contrast := math.Sqrt(math.Max(0, sumLumSq/n-(sumLum/n)*(sumLum/n))) / 255.0
|
|
saturation := sumSat / n
|
|
|
|
clusters := kmeansRGB(samples, 5, 12)
|
|
colorClusters := make([]model.ColorCluster, 0, len(clusters))
|
|
for _, c := range clusters {
|
|
if c.count == 0 {
|
|
continue
|
|
}
|
|
rgb := model.RGB{
|
|
R: int(math.Round(c.center.r)),
|
|
G: int(math.Round(c.center.g)),
|
|
B: int(math.Round(c.center.b)),
|
|
}
|
|
colorClusters = append(colorClusters, model.ColorCluster{
|
|
RGB: rgb, LAB: RGBToLAB(rgb), HSV: RGBToHSV(rgb),
|
|
Percentage: round4(float64(c.count) / n),
|
|
})
|
|
}
|
|
sort.Slice(colorClusters, func(i, j int) bool { return colorClusters[i].Percentage > colorClusters[j].Percentage })
|
|
|
|
dominant := avg
|
|
secondary := avg
|
|
if len(colorClusters) > 0 {
|
|
dominant = colorClusters[0].RGB
|
|
}
|
|
if len(colorClusters) > 1 {
|
|
secondary = colorClusters[1].RGB
|
|
}
|
|
|
|
return model.VisualFeatures{
|
|
DominantRGB: dominant, SecondaryRGB: secondary,
|
|
DominantLAB: RGBToLAB(dominant), DominantHSV: RGBToHSV(dominant),
|
|
AverageRGB: avg, AverageLAB: avgLAB, AverageHSV: avgHSV,
|
|
Brightness: round4(brightness), Contrast: round4(contrast), Saturation: round4(saturation),
|
|
ColorClusters: colorClusters,
|
|
HistogramBins: 256, HistogramSpace: "hsv",
|
|
}, hist
|
|
}
|
|
|
|
func collectSamples(img image.Image, maxSamples int) []rgbSample {
|
|
b := img.Bounds()
|
|
total := b.Dx() * b.Dy()
|
|
if total <= 0 {
|
|
return nil
|
|
}
|
|
step := 1
|
|
if total > maxSamples {
|
|
step = int(math.Ceil(math.Sqrt(float64(total) / float64(maxSamples))))
|
|
}
|
|
samples := make([]rgbSample, 0, min(total/(step*step)+1, maxSamples+1024))
|
|
for y := b.Min.Y; y < b.Max.Y; y += step {
|
|
for x := b.Min.X; x < b.Max.X; x += step {
|
|
r, g, bb := rgba8(img.At(x, y))
|
|
hsv := rgbToHSVFloat(float64(r), float64(g), float64(bb))
|
|
lum := 0.2126*float64(r) + 0.7152*float64(g) + 0.0722*float64(bb)
|
|
samples = append(samples, rgbSample{r: float64(r), g: float64(g), b: float64(bb), lum: lum, sat: hsv.S})
|
|
}
|
|
}
|
|
return samples
|
|
}
|
|
|
|
func buildHSVHistogram(samples []rgbSample, hueBins, saturationBins, valueBins int) model.ColorHistogram {
|
|
values := make([]float64, hueBins*saturationBins*valueBins)
|
|
if len(samples) == 0 {
|
|
return model.ColorHistogram{Space: "hsv", Bins: []int{hueBins, saturationBins, valueBins}, Values: values}
|
|
}
|
|
for _, s := range samples {
|
|
hsv := rgbToHSVFloat(s.r, s.g, s.b)
|
|
hi := min(hueBins-1, int(hsv.H/360.0*float64(hueBins)))
|
|
si := min(saturationBins-1, int(hsv.S*float64(saturationBins)))
|
|
vi := min(valueBins-1, int(hsv.V*float64(valueBins)))
|
|
values[hi*saturationBins*valueBins+si*valueBins+vi]++
|
|
}
|
|
denom := float64(len(samples))
|
|
for i := range values {
|
|
values[i] = round6(values[i] / denom)
|
|
}
|
|
return model.ColorHistogram{Space: "hsv", Bins: []int{hueBins, saturationBins, valueBins}, Values: values}
|
|
}
|
|
|
|
func kmeansRGB(samples []rgbSample, k, iterations int) []colorCluster {
|
|
if len(samples) == 0 || k <= 0 {
|
|
return nil
|
|
}
|
|
sorted := append([]rgbSample(nil), samples...)
|
|
sort.Slice(sorted, func(i, j int) bool { return sorted[i].lum < sorted[j].lum })
|
|
if k > len(sorted) {
|
|
k = len(sorted)
|
|
}
|
|
centers := make([]rgbSample, k)
|
|
for i := 0; i < k; i++ {
|
|
idx := int((float64(i) + 0.5) / float64(k) * float64(len(sorted)-1))
|
|
centers[i] = sorted[idx]
|
|
}
|
|
|
|
assignments := make([]int, len(samples))
|
|
for iter := 0; iter < iterations; iter++ {
|
|
sums := make([]rgbSample, k)
|
|
counts := make([]int, k)
|
|
for i, s := range samples {
|
|
bestIdx := 0
|
|
bestDist := math.MaxFloat64
|
|
for c, center := range centers {
|
|
dr, dg, db := s.r-center.r, s.g-center.g, s.b-center.b
|
|
dist := dr*dr + dg*dg + db*db
|
|
if dist < bestDist {
|
|
bestDist = dist
|
|
bestIdx = c
|
|
}
|
|
}
|
|
assignments[i] = bestIdx
|
|
sums[bestIdx].r += s.r
|
|
sums[bestIdx].g += s.g
|
|
sums[bestIdx].b += s.b
|
|
sums[bestIdx].lum += s.lum
|
|
sums[bestIdx].sat += s.sat
|
|
counts[bestIdx]++
|
|
}
|
|
for i := range centers {
|
|
if counts[i] == 0 {
|
|
continue
|
|
}
|
|
denom := float64(counts[i])
|
|
centers[i] = rgbSample{
|
|
r: sums[i].r / denom, g: sums[i].g / denom, b: sums[i].b / denom,
|
|
lum: sums[i].lum / denom, sat: sums[i].sat / denom,
|
|
}
|
|
}
|
|
}
|
|
|
|
clusters := make([]colorCluster, k)
|
|
for i := range clusters {
|
|
clusters[i].center = centers[i]
|
|
}
|
|
for _, idx := range assignments {
|
|
clusters[idx].count++
|
|
}
|
|
sort.Slice(clusters, func(i, j int) bool { return clusters[i].count > clusters[j].count })
|
|
return clusters
|
|
}
|
|
|
|
func RGBToHSV(rgb model.RGB) model.HSV {
|
|
return rgbToHSVFloat(float64(rgb.R), float64(rgb.G), float64(rgb.B))
|
|
}
|
|
|
|
func rgbToHSVFloat(r, g, b float64) model.HSV {
|
|
r /= 255
|
|
g /= 255
|
|
b /= 255
|
|
maxV := math.Max(r, math.Max(g, b))
|
|
minV := math.Min(r, math.Min(g, b))
|
|
d := maxV - minV
|
|
h := 0.0
|
|
if d != 0 {
|
|
switch maxV {
|
|
case r:
|
|
h = math.Mod((g-b)/d, 6)
|
|
case g:
|
|
h = (b-r)/d + 2
|
|
default:
|
|
h = (r-g)/d + 4
|
|
}
|
|
h *= 60
|
|
if h < 0 {
|
|
h += 360
|
|
}
|
|
}
|
|
s := 0.0
|
|
if maxV != 0 {
|
|
s = d / maxV
|
|
}
|
|
return model.HSV{H: round4(h), S: round4(s), V: round4(maxV)}
|
|
}
|
|
|
|
func RGBToLAB(rgb model.RGB) model.LAB {
|
|
r := pivotRGB(float64(rgb.R) / 255.0)
|
|
g := pivotRGB(float64(rgb.G) / 255.0)
|
|
b := pivotRGB(float64(rgb.B) / 255.0)
|
|
|
|
x := (r*0.4124 + g*0.3576 + b*0.1805) / 0.95047
|
|
y := (r*0.2126 + g*0.7152 + b*0.0722) / 1.00000
|
|
z := (r*0.0193 + g*0.1192 + b*0.9505) / 1.08883
|
|
|
|
fx, fy, fz := pivotXYZ(x), pivotXYZ(y), pivotXYZ(z)
|
|
return model.LAB{
|
|
L: round4(116*fy - 16),
|
|
A: round4(500 * (fx - fy)),
|
|
B: round4(200 * (fy - fz)),
|
|
}
|
|
}
|
|
|
|
func pivotRGB(v float64) float64 {
|
|
if v > 0.04045 {
|
|
return math.Pow((v+0.055)/1.055, 2.4)
|
|
}
|
|
return v / 12.92
|
|
}
|
|
|
|
func pivotXYZ(v float64) float64 {
|
|
if v > 0.008856 {
|
|
return math.Cbrt(v)
|
|
}
|
|
return 7.787*v + 16.0/116.0
|
|
}
|
|
|
|
func round4(v float64) float64 { return math.Round(v*10000) / 10000 }
|
|
func round6(v float64) float64 { return math.Round(v*1000000) / 1000000 }
|