Simplified area recognition

This commit is contained in:
Eduard Urbach 2018-04-11 19:41:23 +02:00
parent bd6391d885
commit 5d4c0fdce2

View File

@ -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
Y int
Color arn.HSLColor
X int
Y int
}
// Area ...
type Area struct {
AverageColor color.Color
Pixels []Pixel
color.Color
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,
Y: y,
Color: hsl,
X: x,
Y: y,
})
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 := float64(rUint) / max
g := float64(gUint) / max
b := float64(bUint) / max
h, s, l := arn.RGBToHSL(r, g, b)
r, g, b, a := color.RGBA()
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?
for _, pixel := range area.Pixels {
xDist := x - pixel.X
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 {
// Is the color similar?
if diffAbs(r, avgR) <= tolerance && diffAbs(g, avgG) <= tolerance && diffAbs(b, avgB) <= tolerance {
areaIndex = i
break
}
}
@ -123,11 +123,7 @@ func ImproveQuality(img image.Image) *image.NRGBA {
hueAreas = append(hueAreas, Area{})
}
hueAreas[areaIndex].Add(x, y, arn.HSLColor{
Hue: h,
Saturation: s,
Lightness: l,
})
hueAreas[areaIndex].Add(x, y, r, g, b, a)
}
}
@ -135,32 +131,10 @@ func ImproveQuality(img image.Image) *image.NRGBA {
// Build image from areas
for _, area := range hueAreas {
totalR := uint64(0)
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,
}
avgColor := area.AverageColor()
for _, pixel := range area.Pixels {
clone.Set(pixel.X, pixel.Y, area.AverageColor)
clone.Set(pixel.X, pixel.Y, avgColor)
}
}