octree.go

package quantise

import (
	"image"
	"image/color"
)

type Strategy int

const (
	// Merge the colours representing the fewest pixels
	LEAST Strategy = iota

	// Merge the colours representing the greatest pixels
	MOST
)

type OctreeQuantiser struct {
	Size     int
	Depth    uint8
	Strategy Strategy
}

func (q OctreeQuantiser) Quantise(in image.Image) color.Palette {
	octree := &oct{isLeaf: false}
	bounds := in.Bounds()

	for y := bounds.Min.Y; y < bounds.Max.Y; y++ {
		for x := bounds.Min.X; x < bounds.Max.X; x++ {
			octree.insert(in.At(x, y), q.Size, q.Depth, q.Strategy)
		}
	}

	return octree.palette()
}

// returns the dth bit of n
func bit(n, d uint8) uint8 {
	if n&(1<<(7-d)) == 0 {
		return 0
	}
	return 1
}

// An oct represents a node or leaf in the oct-tree.
type oct struct {
	isLeaf bool

	// A node in an octree simply has eight children
	children [8]*oct

	// A leaf has a color, and count
	color *color.Color
	count uint64
}

func (tree *oct) justInsert(c *color.Color, r, g, b, depth, maxDepth uint8) {
	tree.count += 1

	if tree.isLeaf {
		return
	}

	index := bit(r, depth)<<2 | bit(g, depth)<<1 | bit(b, depth)

	if tree.children[index] == nil {
		if depth == maxDepth {
			tree.children[index] = &oct{isLeaf: true, color: c, count: 1}
			return
		}

		tree.children[index] = &oct{isLeaf: false, count: 1}
	}

	tree.children[index].justInsert(c, r, g, b, depth+1, maxDepth)
}

func (tree *oct) deepest() []*oct {
	nodes := []*oct{}
	last := []*oct{tree}

	for {
		for _, l := range last {
			for _, child := range l.children {
				if child != nil && !child.isLeaf {
					nodes = append(nodes, child)
				}
			}
		}

		if len(nodes) == 0 {
			return last
		}

		last = nodes
		nodes = []*oct{}
	}
}

func (tree *oct) leaves() []*oct {
	if tree.isLeaf {
		return []*oct{tree}
	}

	leaves := []*oct{}
	for _, child := range tree.children {
		if child != nil {
			leaves = append(leaves, child.leaves()...)
		}
	}

	return leaves
}

func (tree *oct) insert(c color.Color, size int, maxDepth uint8, strategy Strategy) {
	if len(tree.leaves()) <= size {
		r, g, b, _ := c.RGBA()
		tree.justInsert(&c, uint8(r), uint8(g), uint8(b), 0, maxDepth)

	} else {
		deepest := tree.deepest()
		toMerge := deepest[0]

		for _, node := range deepest {
			if strategy == LEAST {
				if node.count < toMerge.count {
					toMerge = node
				}
			} else {
				if node.count > toMerge.count {
					toMerge = node
				}
			}
		}

		toMerge.average()
		tree.insert(c, size, maxDepth, strategy)
	}
}

func (tree *oct) average2() (color.Color, uint64) {
	if tree == nil {
		return nil, 0
	}

	if tree.isLeaf {
		if tree.color == nil {
			return nil, 0
		}
		return *tree.color, tree.count
	}

	var rt, gt, bt, ct uint64

	for i := 0; i < 8; i++ {
		child := tree.children[i]

		avg, c := child.average2()
		if avg == nil {
			continue
		}

		r, g, b, _ := avg.RGBA()
		rt += uint64(r) * c / 255
		gt += uint64(g) * c / 255
		bt += uint64(b) * c / 255
		ct += c
	}

	if ct == 0 {
		return nil, 0
	}

	return color.RGBA{uint8(rt / ct), uint8(gt / ct), uint8(bt / ct), 255}, ct
}

func (tree *oct) average() {
	if tree.isLeaf {
		return
	}

	avg, ct := tree.average2()

	tree.children = [8]*oct{}
	tree.isLeaf = true
	tree.color = &avg
	tree.count = ct
}

func (tree *oct) palette() color.Palette {
	colors := []color.Color{}
	leaves := tree.leaves()

	for i := 0; i < len(leaves); i++ {
		colors = append(colors, *leaves[i].color)
	}

	return colors
}