Simplified area recognition

patches/improve-image-quality/improve-image-quality.go

@@ -8,10 +8,7 @@ import (
 	_ "image/jpeg"
 	"image/png"
 	"io/ioutil"
-	"math"
 	"os"
-
-	"github.com/animenotifier/arn"
 )
 
 func main() {
@@ -43,32 +40,58 @@ const max = float64(65535)
 
 // Pixel ...
 type Pixel struct {
-	X     int
+	X int
-	Y     int
+	Y int
-	Color arn.HSLColor
 }
 
 // Area ...
 type Area struct {
-	AverageColor color.Color
+	color.Color
-	Pixels       []Pixel
+	Pixels []Pixel
+	totalR uint64
+	totalG uint64
+	totalB uint64
+	totalA uint64
 }
 
 // Add ...
-func (area *Area) Add(x, y int, hsl arn.HSLColor) {
+func (area *Area) Add(x, y int, r, g, b, a uint32) {
 	area.Pixels = append(area.Pixels, Pixel{
-		X:     x,
+		X: x,
-		Y:     y,
+		Y: y,
-		Color: hsl,
 	})
+
+	area.totalR += uint64(r)
+	area.totalG += uint64(g)
+	area.totalB += uint64(b)
+	area.totalA += uint64(a)
+}
+
+// AverageColor ...
+func (area *Area) AverageColor() color.Color {
+	return color.RGBA64{
+		R: uint16(area.totalR / uint64(len(area.Pixels))),
+		G: uint16(area.totalG / uint64(len(area.Pixels))),
+		B: uint16(area.totalB / uint64(len(area.Pixels))),
+		A: uint16(area.totalA / uint64(len(area.Pixels))),
+	}
 }
 
 const (
+	tolerance           = uint32(3000)
 	hueTolerance        = 0.1
 	lightnessTolerance  = 0.1
 	saturationTolerance = 0.1
 )
 
+func diffAbs(a uint32, b uint32) uint32 {
+	if a > b {
+		return a - b
+	}
+
+	return b - a
+}
+
 // ImproveQuality returns the average color of an image in HSL format.
 func ImproveQuality(img image.Image) *image.NRGBA {
 	width := img.Bounds().Dx()
@@ -79,40 +102,17 @@ func ImproveQuality(img image.Image) *image.NRGBA {
 	for x := 0; x < width; x++ {
 		for y := 0; y < height; y++ {
 			color := img.At(x, y)
-			rUint, gUint, bUint, _ := color.RGBA()
+			r, g, b, a := color.RGBA()
-			r := float64(rUint) / max
-			g := float64(gUint) / max
-			b := float64(bUint) / max
-			h, s, l := arn.RGBToHSL(r, g, b)
 			areaIndex := -1
 
 			// Find similar area
 			for i := 0; i < len(hueAreas); i++ {
 				area := hueAreas[i]
+				avgR, avgG, avgB, _ := area.AverageColor().RGBA()
 
-				// Is the pixel close to any pixel in the area we're checking?
+				// Is the color similar?
-				for _, pixel := range area.Pixels {
+				if diffAbs(r, avgR) <= tolerance && diffAbs(g, avgG) <= tolerance && diffAbs(b, avgB) <= tolerance {
-					xDist := x - pixel.X
+					areaIndex = i
-					yDist := y - pixel.Y
-
-					if xDist < 0 {
-						xDist = -xDist
-					}
-
-					if yDist < 0 {
-						yDist = -yDist
-					}
-
-					if xDist <= 1 && yDist <= 1 {
-						// Is the color similar?
-						if math.Abs(h-pixel.Color.Hue) <= hueTolerance && math.Abs(s-pixel.Color.Saturation) <= saturationTolerance && math.Abs(l-pixel.Color.Lightness) <= lightnessTolerance {
-							areaIndex = i
-							break
-						}
-					}
-				}
-
-				if areaIndex != -1 {
 					break
 				}
 			}
@@ -123,11 +123,7 @@ func ImproveQuality(img image.Image) *image.NRGBA {
 				hueAreas = append(hueAreas, Area{})
 			}
 
-			hueAreas[areaIndex].Add(x, y, arn.HSLColor{
+			hueAreas[areaIndex].Add(x, y, r, g, b, a)
-				Hue:        h,
-				Saturation: s,
-				Lightness:  l,
-			})
 		}
 	}
 
@@ -135,32 +131,10 @@ func ImproveQuality(img image.Image) *image.NRGBA {
 
 	// Build image from areas
 	for _, area := range hueAreas {
-		totalR := uint64(0)
+		avgColor := area.AverageColor()
-		totalG := uint64(0)
-		totalB := uint64(0)
-
-		// Calculate area average color
-		for _, pixel := range area.Pixels {
-			col := img.At(pixel.X, pixel.Y)
-			r, g, b, _ := col.RGBA()
-			totalR += uint64(r)
-			totalG += uint64(g)
-			totalB += uint64(b)
-		}
-
-		averageR := float64(totalR/uint64(len(area.Pixels))) / max
-		averageG := float64(totalG/uint64(len(area.Pixels))) / max
-		averageB := float64(totalB/uint64(len(area.Pixels))) / max
-
-		area.AverageColor = color.RGBA{
-			R: uint8(averageR * 255),
-			G: uint8(averageG * 255),
-			B: uint8(averageB * 255),
-			A: 255,
-		}
 
 		for _, pixel := range area.Pixels {
-			clone.Set(pixel.X, pixel.Y, area.AverageColor)
+			clone.Set(pixel.X, pixel.Y, avgColor)
 		}
 	}