290 lines
5.7 KiB
Go
Raw Normal View History

2018-04-11 17:01:09 +00:00
package main
import (
"bytes"
"fmt"
"image"
"image/color"
_ "image/jpeg"
"image/png"
"io/ioutil"
"os"
)
func main() {
data, err := ioutil.ReadFile("input.jpg")
if err != nil {
panic(err)
}
2018-04-11 20:25:10 +00:00
img, _, err := image.Decode(bytes.NewReader(data))
2018-04-11 17:01:09 +00:00
if err != nil {
panic(err)
}
improved := ImproveQuality(img)
f, err := os.Create("output.png")
if err != nil {
panic(err)
}
defer f.Close()
png.Encode(f, improved)
}
const max = float64(65535)
// Pixel ...
type Pixel struct {
2018-04-11 17:41:23 +00:00
X int
Y int
2018-04-11 17:01:09 +00:00
}
// Area ...
type Area struct {
2018-04-11 17:41:23 +00:00
Pixels []Pixel
totalR uint64
totalG uint64
totalB uint64
totalA uint64
2018-04-11 17:01:09 +00:00
}
// Add ...
2018-04-11 17:41:23 +00:00
func (area *Area) Add(x, y int, r, g, b, a uint32) {
2018-04-11 17:01:09 +00:00
area.Pixels = append(area.Pixels, Pixel{
2018-04-11 17:41:23 +00:00
X: x,
Y: y,
2018-04-11 17:01:09 +00:00
})
2018-04-11 17:41:23 +00:00
area.totalR += uint64(r)
area.totalG += uint64(g)
area.totalB += uint64(b)
area.totalA += uint64(a)
}
// AverageColor ...
func (area *Area) AverageColor() color.Color {
2018-04-11 20:25:10 +00:00
if len(area.Pixels) == 0 {
return color.Transparent
}
2018-04-11 17:41:23 +00:00
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))),
}
2018-04-11 17:01:09 +00:00
}
const (
2018-04-11 20:59:04 +00:00
tolerance = uint32(3000)
2018-04-11 17:01:09 +00:00
)
2018-04-11 17:41:23 +00:00
func diffAbs(a uint32, b uint32) uint32 {
if a > b {
return a - b
}
return b - a
}
2018-04-11 17:01:09 +00:00
// ImproveQuality returns the average color of an image in HSL format.
func ImproveQuality(img image.Image) *image.NRGBA {
width := img.Bounds().Dx()
height := img.Bounds().Dy()
clone := image.NewNRGBA(image.Rect(0, 0, width, height))
2018-04-11 20:25:10 +00:00
areas := []*Area{}
2018-04-11 18:23:16 +00:00
areaIndexMap := make([]int, width*height)
2018-04-11 17:01:09 +00:00
for x := 0; x < width; x++ {
for y := 0; y < height; y++ {
color := img.At(x, y)
2018-04-11 17:41:23 +00:00
r, g, b, a := color.RGBA()
2018-04-11 17:01:09 +00:00
areaIndex := -1
// Find similar area
2018-04-11 18:23:16 +00:00
for i := 0; i < len(areas); i++ {
area := areas[i]
2018-04-11 17:41:23 +00:00
avgR, avgG, avgB, _ := area.AverageColor().RGBA()
2018-04-11 17:01:09 +00:00
2018-04-11 17:41:23 +00:00
// Is the color similar?
if diffAbs(r, avgR) <= tolerance && diffAbs(g, avgG) <= tolerance && diffAbs(b, avgB) <= tolerance {
areaIndex = i
2018-04-11 17:01:09 +00:00
break
}
}
// Insert new area
if areaIndex == -1 {
2018-04-11 18:23:16 +00:00
areaIndex = len(areas)
2018-04-11 20:25:10 +00:00
areas = append(areas, &Area{})
2018-04-11 17:01:09 +00:00
}
2018-04-11 18:23:16 +00:00
areaIndexMap[y*width+x] = areaIndex
areas[areaIndex].Add(x, y, r, g, b, a)
2018-04-11 17:01:09 +00:00
}
}
2018-04-11 18:23:16 +00:00
fmt.Println(len(areas), "areas")
// Reduce noise
noiseCount := 0
2018-04-11 20:25:10 +00:00
for areaIndex, area := range areas {
2018-04-11 20:59:04 +00:00
noisePixelIndices := []int{}
2018-04-11 20:25:10 +00:00
areaSurroundedBy := map[int]int{}
for i := 0; i < len(area.Pixels); i++ {
// If pixel is surrounded by 4 different areas, remove it
pixel := area.Pixels[i]
x := pixel.X
y := pixel.Y
left := areaIndex
right := areaIndex
top := areaIndex
bottom := areaIndex
if x > 0 {
left = areaIndexMap[y*width+(x-1)]
}
2018-04-11 18:23:16 +00:00
2018-04-11 20:25:10 +00:00
if x < width-1 {
right = areaIndexMap[y*width+(x+1)]
}
2018-04-11 18:23:16 +00:00
2018-04-11 20:25:10 +00:00
if y > 0 {
top = areaIndexMap[(y-1)*width+x]
}
2018-04-11 18:23:16 +00:00
2018-04-11 20:25:10 +00:00
if y < height-1 {
bottom = areaIndexMap[(y+1)*width+x]
}
2018-04-11 18:23:16 +00:00
2018-04-11 20:25:10 +00:00
differentNeighbors := 0
2018-04-11 18:23:16 +00:00
2018-04-11 20:25:10 +00:00
if left != areaIndex {
differentNeighbors++
}
2018-04-11 18:23:16 +00:00
2018-04-11 20:25:10 +00:00
if right != areaIndex {
differentNeighbors++
}
2018-04-11 18:23:16 +00:00
2018-04-11 20:25:10 +00:00
if top != areaIndex {
differentNeighbors++
}
2018-04-11 18:23:16 +00:00
2018-04-11 20:25:10 +00:00
if bottom != areaIndex {
differentNeighbors++
}
2018-04-11 18:23:16 +00:00
2018-04-11 20:25:10 +00:00
// Determine surrounding area
areaIndexScore := map[int]int{}
areaIndexScore[left]++
areaIndexScore[right]++
areaIndexScore[top]++
areaIndexScore[bottom]++
2018-04-11 18:23:16 +00:00
2018-04-11 20:25:10 +00:00
areaSurroundedBy[left]++
areaSurroundedBy[right]++
areaSurroundedBy[top]++
areaSurroundedBy[bottom]++
2018-04-11 18:23:16 +00:00
2018-04-11 20:25:10 +00:00
newAreaIndex := -1
bestScore := 0
for checkIndex, score := range areaIndexScore {
if score > bestScore {
bestScore = score
newAreaIndex = checkIndex
2018-04-11 18:23:16 +00:00
}
}
2018-04-11 20:25:10 +00:00
if differentNeighbors >= 3 && bestScore >= 3 {
noiseCount++
2018-04-11 20:59:04 +00:00
noisePixelIndices = append(noisePixelIndices, i)
2018-04-11 18:23:16 +00:00
2018-04-11 20:25:10 +00:00
// Add to surrounding area
r, g, b, a := img.At(x, y).RGBA()
areas[newAreaIndex].Add(x, y, r, g, b, a)
2018-04-11 20:59:04 +00:00
area.totalR -= uint64(r)
area.totalG -= uint64(g)
area.totalB -= uint64(b)
area.totalA -= uint64(a)
2018-04-11 18:23:16 +00:00
}
}
2018-04-11 20:25:10 +00:00
// Remove noise pixels
2018-04-11 20:59:04 +00:00
offset := 0
2018-04-11 20:25:10 +00:00
2018-04-11 20:59:04 +00:00
for _, removal := range noisePixelIndices {
index := removal - offset
area.Pixels = append(area.Pixels[:index], area.Pixels[index+1:]...)
offset++
}
2018-04-11 20:25:10 +00:00
// Determine surrounding area
surroundingAreaIndex := -1
bestScore := 0
for checkIndex, score := range areaSurroundedBy {
2018-04-11 20:59:04 +00:00
if score > bestScore && checkIndex != areaIndex {
2018-04-11 20:25:10 +00:00
bestScore = score
surroundingAreaIndex = checkIndex
}
}
2018-04-11 20:59:04 +00:00
surroundingArea := areas[surroundingAreaIndex]
2018-04-11 20:25:10 +00:00
2018-04-11 20:59:04 +00:00
if areaIndex != surroundingAreaIndex && len(surroundingArea.Pixels) > len(area.Pixels)*2 {
// const surroundTolerance = 5000
// r1, g1, b1, a1 := area.AverageColor().RGBA()
// r2, g2, b2, a2 := surroundingArea.AverageColor().RGBA()
// if diffAbs(r1, r2) < surroundTolerance && diffAbs(g1, g2) < surroundTolerance && diffAbs(b1, b2) < surroundTolerance && diffAbs(a1, a2) < surroundTolerance {
// // fmt.Println(areaIndex, "surrounded by", surroundingAreaIndex, "|", len(area.Pixels), len(surroundingArea.Pixels))
2018-04-11 20:25:10 +00:00
2018-04-11 20:59:04 +00:00
// // Add pixels to surrounding area
// for _, pixel := range area.Pixels {
// r, g, b, a := img.At(pixel.X, pixel.Y).RGBA()
// surroundingArea.Add(pixel.X, pixel.Y, r, g, b, a)
// }
// // Remove this area
// area.Pixels = nil
// area.totalR = 0
// area.totalG = 0
// area.totalB = 0
// area.totalA = 0
// }
2018-04-11 20:25:10 +00:00
}
2018-04-11 18:23:16 +00:00
}
fmt.Println(noiseCount, "noise pixels")
2018-04-11 17:01:09 +00:00
2018-04-11 20:59:04 +00:00
pixelCount := 0
2018-04-11 20:25:10 +00:00
for _, area := range areas {
pixelCount += len(area.Pixels)
}
fmt.Println(pixelCount, "pixels", width*height)
2018-04-11 17:01:09 +00:00
// Build image from areas
2018-04-11 18:23:16 +00:00
for _, area := range areas {
2018-04-11 17:41:23 +00:00
avgColor := area.AverageColor()
2018-04-11 17:01:09 +00:00
for _, pixel := range area.Pixels {
2018-04-11 17:41:23 +00:00
clone.Set(pixel.X, pixel.Y, avgColor)
2018-04-11 17:01:09 +00:00
}
}
return clone
}