Rough draft of how to implement connected component labelling using disjoint sets

connected-components.go

package main

import (
	"fmt"
	"image"
	_ "image/png"
	"log"
	"os"
)

type Pos struct {
	x, y int
}

type DisjointSet struct {
	parents map[Pos]Pos
	size    map[Pos]int
	labels  map[Pos]int
}

func newDisjointSet() *DisjointSet {
	return &DisjointSet{
		parents: make(map[Pos]Pos),
		size:    make(map[Pos]int),
		labels:  make(map[Pos]int),
	}
}

func (s *DisjointSet) add(x Pos, l int) {
	// ignore already registered
	_, ok := s.size[x]

	if ok {
		return
	}

	s.labels[x] = l
	s.parents[x] = x
	s.size[x] = 1
}

func (s *DisjointSet) find(x Pos) Pos {
	if s.parents[x] == x {
		return x
	}

	return s.find(s.parents[x])
}

func (s *DisjointSet) union(s1 Pos, l1 int, s2 Pos, l2 int) {
	r1 := s.find(s1)
	r2 := s.find(s2)

	// already in the same set
	if r1 == r2 {
		return
	}

	fmt.Printf("union of %#v and %#v (%d %d)\n", s1, s2, l1, l2)
	fmt.Printf("cur root %#v and %#v (%d %d)\n", r1, r2, s.labels[r1], s.labels[r2])

	if s.size[r1] >= s.size[r2] {
		// make r1 parent of r2
		s.size[r1] = s.size[r1] + s.size[r2]
		s.parents[r2] = r1

		s.labels[r1] = minInt(s.labels[r1], l1, l2)

		fmt.Printf("\t %d %d %d %d\n", s.labels[r2], l1, l2, minInt(4, 4, 4))
		println("\t A root label ", s.labels[r1])
	} else {
		// make r2 parent of r1
		s.size[r2] = s.size[r1] + s.size[r2]
		s.parents[r1] = r2

		s.labels[r2] = minInt(s.labels[r2], l1, l2)

		println("\t B root label ", s.labels[r2])
	}
}

func (s *DisjointSet) sameComponent(s1, s2 Pos) bool {
	return s.find(s1) == s.find(s2)
}

func minInt(vals ...int) int {
	var tmp = 0

	for _, val := range vals {
		if tmp == 0 && val != 0 {
			tmp = val
		}

		if val == 0 { // 0 is just placeholder
			continue
		}

		if val < tmp {
			tmp = val
		}
	}

	return tmp
}

var labels [6][5]int // @FIXME change hardcoded labels size

func main() {
	reader, err := os.Open("/Users/max/Desktop/frob-tiny.png")

	if err != nil {
		log.Fatal(err)
	}

	defer reader.Close()

	m, _, err := image.Decode(reader)
	if err != nil {
		log.Fatal(err)
	}

	bounds := m.Bounds()
	set := newDisjointSet()

	var labelCount = 1

	for y := bounds.Min.Y; y < bounds.Max.Y; y++ {
		for x := bounds.Min.X; x < bounds.Max.X; x++ {
			r, _, _, _ := m.At(x, y).RGBA()
			color := int(r / 257)
			curPos := Pos{x, y}

			// get neighbors in same color
			neighbors := neighbors(m, x, y, bounds.Max.X, bounds.Max.Y, color)

			// get min label value from all labelled neighbors
			minNeighbor, minLabel := min(neighbors)

			// none of the neighbors has been labelled yet, label this cell
			if minNeighbor == nil {
				set.add(curPos, labelCount)

				labels[x][y] = labelCount

				// store equivalence relation between current cell and its neighbors
				// notice: the neighbors don't necessarily have the same label as the current cell,
				// but they do have the same _color_ as the current cell
				for neighPos, neighLabel := range neighbors {
					set.add(neighPos, neighLabel)
					set.union(curPos, labelCount, neighPos, neighLabel)
				}

				labelCount++

				continue
			}

			set.add(curPos, minLabel)

			// store equivalence relation between current cell and its neighbors
			// notice: the neighbors don't necessarily have the same label as the current cell,
			// but they do have the same _color_ as the current cell
			for neighPos, neighLabel := range neighbors {
				set.add(neighPos, neighLabel)
				set.union(curPos, minLabel, neighPos, neighLabel)
			}

			// @TODO define more invariant asserts?

			// find the parent of the set the min neighbor belongs to
			parentPos := set.find(*minNeighbor)
			set.union(curPos, minLabel, parentPos, minLabel)

			labels[x][y] = set.labels[parentPos]

			// sanity check
			r2, _, _, _ := m.At(parentPos.x, parentPos.y).RGBA()
			color2 := int(r2 / 257)

			if color != color2 {
				println("FAIL at", x, y, "parent is supposed to be", parentPos.x, parentPos.y)
				panic("wtf?")
			}
		}
	}

	// second pass, assign the cell the min label of the connected component it belongs to
	for y := bounds.Min.Y; y < bounds.Max.Y; y++ {
		for x := bounds.Min.X; x < bounds.Max.X; x++ {
			// @TODO assert parent pos is smaller?
			parentPos := set.find(Pos{x, y})
			labels[x][y] = set.labels[parentPos]
		}
	}

	// dump the results
	for y := bounds.Min.Y; y < bounds.Max.Y; y++ {
		for x := bounds.Min.X; x < bounds.Max.X; x++ {
			cur := labels[x][y]
			fmt.Printf("\t%d ", cur)
		}

		fmt.Println()
	}
}

var deltas = []struct {
	x, y int
}{
	// 4-connected
	{0, -1}, // North
	{0, 1},  // South
	{1, 0},  // East
	{-1, 0}, // West

	// 8-connected
	{-1, -1}, // North-West
	{1, -1},  // North-East
	{-1, 1},  // South-West
	{1, -1},  // South-East
}

// neighbors in same color
func neighbors(img image.Image, x, y, width, height, color int) map[Pos]int {
	tmp := make(map[Pos]int)

	for _, d := range deltas {
		dx := x + d.x
		dy := y + d.y

		if dx < 0 || dy < 0 || dx >= width || dy >= height {
			continue
		}

		r, _, _, _ := img.At(dx, dy).RGBA()
		curColor := int(r / 257)

		if curColor != color {
			continue
		}

		if labels[dx][dy] != 0 {
			tmp[Pos{dx, dy}] = labels[dx][dy]
		}
	}

	return tmp
}

// find the first one that is not 0
func min(neighbors map[Pos]int) (*Pos, int) {
	var min = 0
	var p *Pos

	for neigh, label := range neighbors {
		if min == 0 && label != 0 {
			min = label
			p = &neigh

			continue
		}

		if label == 0 {
			continue
		}

		if label < min {
			min = label
			p = &neigh
		}
	}

	return p, min
}