113 lines
2.1 KiB
Go
113 lines
2.1 KiB
Go
package pipeline
|
|
|
|
import (
|
|
"math"
|
|
"sort"
|
|
)
|
|
|
|
// ClusterRows runs 1D DBSCAN on Y coordinates (parity with sklearn DBSCAN).
|
|
func ClusterRows(boxes []indexedBox, eps float64, minSamples int) {
|
|
if len(boxes) == 0 {
|
|
return
|
|
}
|
|
ys := make([]float64, len(boxes))
|
|
for i, b := range boxes {
|
|
ys[i] = float64(b.Y)
|
|
}
|
|
labels := dbscan1D(ys, eps, minSamples)
|
|
for i := range boxes {
|
|
boxes[i].Cluster = labels[i]
|
|
}
|
|
}
|
|
|
|
func dbscan1D(points []float64, eps float64, minSamples int) []int {
|
|
n := len(points)
|
|
labels := make([]int, n)
|
|
for i := range labels {
|
|
labels[i] = -1
|
|
}
|
|
clusterID := 0
|
|
visited := make([]bool, n)
|
|
|
|
for i := 0; i < n; i++ {
|
|
if visited[i] {
|
|
continue
|
|
}
|
|
visited[i] = true
|
|
neighbors := regionQuery(points, i, eps)
|
|
if len(neighbors) < minSamples {
|
|
continue
|
|
}
|
|
labels[i] = clusterID
|
|
seed := append([]int(nil), neighbors...)
|
|
for j := 0; j < len(seed); j++ {
|
|
q := seed[j]
|
|
if !visited[q] {
|
|
visited[q] = true
|
|
neighborsQ := regionQuery(points, q, eps)
|
|
if len(neighborsQ) >= minSamples {
|
|
seed = appendUnique(seed, neighborsQ)
|
|
}
|
|
}
|
|
if labels[q] == -1 {
|
|
labels[q] = clusterID
|
|
}
|
|
}
|
|
clusterID++
|
|
}
|
|
return labels
|
|
}
|
|
|
|
func regionQuery(points []float64, idx int, eps float64) []int {
|
|
var neighbors []int
|
|
for i, p := range points {
|
|
if math.Abs(p-points[idx]) <= eps {
|
|
neighbors = append(neighbors, i)
|
|
}
|
|
}
|
|
return neighbors
|
|
}
|
|
|
|
func appendUnique(seed, add []int) []int {
|
|
inSeed := make(map[int]bool, len(seed))
|
|
for _, v := range seed {
|
|
inSeed[v] = true
|
|
}
|
|
for _, v := range add {
|
|
if !inSeed[v] {
|
|
seed = append(seed, v)
|
|
inSeed[v] = true
|
|
}
|
|
}
|
|
return seed
|
|
}
|
|
|
|
type indexedBox struct {
|
|
Text string
|
|
X int
|
|
Y int
|
|
W int
|
|
H int
|
|
Cluster int
|
|
}
|
|
|
|
func groupByCluster(boxes []indexedBox) map[int][]indexedBox {
|
|
groups := make(map[int][]indexedBox)
|
|
for _, b := range boxes {
|
|
if b.Cluster < 0 {
|
|
continue
|
|
}
|
|
groups[b.Cluster] = append(groups[b.Cluster], b)
|
|
}
|
|
return groups
|
|
}
|
|
|
|
func sortClusters(groups map[int][]indexedBox) []int {
|
|
keys := make([]int, 0, len(groups))
|
|
for k := range groups {
|
|
keys = append(keys, k)
|
|
}
|
|
sort.Ints(keys)
|
|
return keys
|
|
}
|