mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-08-12 20:03:40 +08:00
fix(pdf/layout): replace KMeans-only column detection with gap + KMeans hybrid (#18023)
This commit is contained in:
Jack
474
internal/deepdoc/parser/pdf/layout/combined_column.go
Normal file
474
internal/deepdoc/parser/pdf/layout/combined_column.go
Normal file
@@ -0,0 +1,474 @@
|
||||
package layout
|
||||
|
||||
import (
|
||||
"math"
|
||||
"math/rand"
|
||||
"sort"
|
||||
|
||||
pdf "ragflow/internal/deepdoc/parser/pdf/type"
|
||||
util "ragflow/internal/deepdoc/parser/pdf/util"
|
||||
)
|
||||
|
||||
// AssignColumn groups boxes into columns using the hybrid gap + KMeans
|
||||
// strategy that beats gap-only column detection on real documents.
|
||||
//
|
||||
// Decision per page (mirrors tool-py/diagnose_combined.py):
|
||||
// 1. Geometric gap (whitespace gutter voting) finds candidate column
|
||||
// separators. But "gap >= 2" is NOT blindly trusted:
|
||||
// - If the resulting columns are NARROW (max column width <
|
||||
// tableMaxColFrac of the page), they are table cells, not text columns:
|
||||
// the page is a single reading block -> return 1 directly (and do NOT
|
||||
// fall through to the balance gate, which would re-split the table's
|
||||
// bimodal x0 into 2).
|
||||
// - If gap == 2, the separator is unreliable (it is often a fake gutter
|
||||
// from indentation/line-width variation, not a real column). Defer to
|
||||
// the balance gate below.
|
||||
// - If gap >= 3 with WIDE columns, it is a real multi-column layout:
|
||||
// trust it and partition by KMeans(g).
|
||||
// 2. When gap reports 1 (single column OR a double column whose gutter is
|
||||
// bridged by full-width front matter), or gap == 2 was deferred, a forced
|
||||
// k=2 KMeans on the BODY x0 decides whether the lines form TWO clusters
|
||||
// each holding >= minModeFrac of body lines, separated by >=
|
||||
// minSepFrac*width. A balanced split is a real second column; an
|
||||
// unbalanced split (the usual KMeans false-split on a single page) is
|
||||
// dropped -> stays 1.
|
||||
//
|
||||
// Net effect: tables and fake gutters no longer over-split, while the
|
||||
// double-column pages that gap alone misses are recovered by the balance gate.
|
||||
func AssignColumn(boxes []pdf.TextBox) []pdf.TextBox {
|
||||
if len(boxes) == 0 {
|
||||
return boxes
|
||||
}
|
||||
pageGroups, sortedPages := groupBoxesByPage(boxes)
|
||||
result := make([]pdf.TextBox, len(boxes))
|
||||
copy(result, boxes)
|
||||
for _, pg := range sortedPages {
|
||||
indices := pageGroups[pg]
|
||||
k, cents := detectColumnCount(boxes, indices)
|
||||
assignColIDs(boxes, result, indices, k, cents)
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// tableMaxColFrac: a column narrower than this fraction of the page width is
|
||||
// treated as a table cell, not a text column. Above this, the columns are
|
||||
// wide enough to be real reading columns.
|
||||
const tableMaxColFrac = 0.22
|
||||
|
||||
// maxColumnCount caps how many columns the gap detector may report. Gap
|
||||
// voting can over-split a single page into many spurious gutters (e.g.
|
||||
// first-line indentation), so we bound the count to the old detector's best-k
|
||||
// cap of min(4, n). This prevents catastrophic splits (a single page reported
|
||||
// as 7+ columns) that the old code could never produce.
|
||||
const maxColumnCount = 4
|
||||
|
||||
// minColLineFrac: a column holding fewer than this fraction of the page's
|
||||
// lines (or zero lines) is not a real reading column — it is a spurious
|
||||
// gutter sliver (an indented block, a stray caption, an empty kmeans
|
||||
// centroid). Drop it so the detector does not over-split.
|
||||
//
|
||||
// The threshold is set with margin below the smallest genuine column ratio
|
||||
// observed on the 70-page labeled corpus: the sparsest real double's minority
|
||||
// column is ~17.7% of lines, and the only real triple's columns are each
|
||||
// >=22%. 12% prunes genuine outliers (e.g. a 4-line footnote, 7.3%) without
|
||||
// touching those.
|
||||
const minColLineFrac = 0.12
|
||||
|
||||
// detectColumnCount returns (columnCount, centroids) for one page.
|
||||
// columnCount is 1, 2, or up to maxColumnCount; centroids are the k cluster
|
||||
// means in x0 space (snapshot of the gate decision) and are reused for ColID
|
||||
// assignment.
|
||||
func detectColumnCount(boxes []pdf.TextBox, indices []int) (int, []float64) {
|
||||
lines := make([]pdf.TextBox, len(indices))
|
||||
for i, idx := range indices {
|
||||
lines[i] = boxes[idx]
|
||||
}
|
||||
g := gapColumnCount(lines, 0.04, 0.15, 2.0)
|
||||
if g >= 2 {
|
||||
_, width := pageExtent(lines)
|
||||
if width > 0 {
|
||||
widths := gapColumnWidths(lines)
|
||||
maxw := 0.0
|
||||
for _, w := range widths {
|
||||
if w > maxw {
|
||||
maxw = w
|
||||
}
|
||||
}
|
||||
if maxw < tableMaxColFrac*width {
|
||||
// Narrow columns => table cells, not text columns. The page
|
||||
// is one reading block; return 1 and skip the balance gate
|
||||
// (which would otherwise re-split the table's x0).
|
||||
return 1, nil
|
||||
}
|
||||
}
|
||||
if g > 2 {
|
||||
// gap >= 3 with wide columns: a real multi-column layout.
|
||||
// Cap the count (maxColumnCount) so spurious gutters cannot
|
||||
// split a single page into many columns, then prune empty or
|
||||
// too-sparse columns so an indentation-created sliver does not
|
||||
// survive as a spurious column.
|
||||
k := g
|
||||
if k > maxColumnCount {
|
||||
k = maxColumnCount
|
||||
}
|
||||
if k > len(lines) {
|
||||
k = len(lines)
|
||||
}
|
||||
_, w := pageExtent(lines)
|
||||
cents := kmeansCentroids(lines, k, w)
|
||||
if pk, pc, ok := pruneColumns(lines, cents); ok {
|
||||
return pk, pc
|
||||
}
|
||||
return 1, nil
|
||||
}
|
||||
// g == 2: unreliable (fake gutter or real 2-col) -> defer to balance.
|
||||
}
|
||||
if ok, cents, body := balancedBodyK2(lines, 0.30, 0.10); ok {
|
||||
// prune on the SAME body the gate clustered, not all lines: full-width
|
||||
// titles/abstracts were deliberately excluded from the balance check
|
||||
// and must not be re-counted here (they would inflate one column and
|
||||
// let prune wrongly collapse a real two-column page to one).
|
||||
if pk, pc, ok2 := pruneColumns(body, cents); ok2 {
|
||||
return pk, pc
|
||||
}
|
||||
return 1, nil
|
||||
}
|
||||
return 1, nil
|
||||
}
|
||||
|
||||
// pruneColumns drops empty (0-line) or too-sparse (< minColLineFrac) columns
|
||||
// from a k-centroid partition and returns the surviving (k', cents'). A column
|
||||
// is "real" only if it captures enough of the page's lines. If fewer than 2
|
||||
// real columns survive, ok is false and the caller should treat the page as a
|
||||
// single column.
|
||||
func pruneColumns(lines []pdf.TextBox, cents []float64) (int, []float64, bool) {
|
||||
n := len(lines)
|
||||
if n == 0 || len(cents) < 2 {
|
||||
return len(cents), cents, len(cents) >= 2
|
||||
}
|
||||
counts := make([]int, len(cents))
|
||||
for _, b := range lines {
|
||||
best, bestD := 0, math.Abs(b.X0-cents[0])
|
||||
for c := 1; c < len(cents); c++ {
|
||||
if d := math.Abs(b.X0 - cents[c]); d < bestD {
|
||||
bestD, best = d, c
|
||||
}
|
||||
}
|
||||
counts[best]++
|
||||
}
|
||||
keep := make([]int, 0, len(cents))
|
||||
for c := range cents {
|
||||
if counts[c] > 0 && float64(counts[c]) >= minColLineFrac*float64(n) {
|
||||
keep = append(keep, c)
|
||||
}
|
||||
}
|
||||
if len(keep) < 2 {
|
||||
return len(keep), nil, false
|
||||
}
|
||||
newCents := make([]float64, len(keep))
|
||||
for i, c := range keep {
|
||||
newCents[i] = cents[c]
|
||||
}
|
||||
return len(keep), newCents, true
|
||||
}
|
||||
|
||||
// gapColumnWidths returns the width (in page units) of each column found by
|
||||
// the same gutter voting as gapColumnCount. Used to tell real wide text
|
||||
// columns apart from narrow table-cell columns.
|
||||
func gapColumnWidths(lines []pdf.TextBox) []float64 {
|
||||
n := len(lines)
|
||||
if n == 0 {
|
||||
return nil
|
||||
}
|
||||
minX0, width := pageExtent(lines)
|
||||
if width <= 0 {
|
||||
return nil
|
||||
}
|
||||
binPt := 2.0
|
||||
nb := int(width/binPt) + 1
|
||||
cov := make([]int, nb)
|
||||
for _, b := range lines {
|
||||
i0 := clampInt(int((b.X0-minX0)/binPt), 0, nb-1)
|
||||
i1 := clampInt(int((b.X1-minX0)/binPt), 0, nb-1)
|
||||
for i := i0; i <= i1; i++ {
|
||||
cov[i]++
|
||||
}
|
||||
}
|
||||
thr := 0.15 * float64(n)
|
||||
var widths []float64
|
||||
i := 0
|
||||
for i < nb {
|
||||
if float64(cov[i]) < thr {
|
||||
i++
|
||||
continue
|
||||
}
|
||||
j := i
|
||||
for j < nb && float64(cov[j]) >= thr {
|
||||
j++
|
||||
}
|
||||
widths = append(widths, float64(j-i)*binPt)
|
||||
i = j
|
||||
}
|
||||
return widths
|
||||
}
|
||||
|
||||
func clampInt(v, lo, hi int) int {
|
||||
if v < lo {
|
||||
return lo
|
||||
}
|
||||
if v > hi {
|
||||
return hi
|
||||
}
|
||||
return v
|
||||
}
|
||||
|
||||
// gapColumnCount mirrors column_detectors.gap_column_counts: rasterize the
|
||||
// [minX0, maxX1] text region into x-bins, count how many lines cover each bin,
|
||||
// and treat a covered-fraction-below-crossTol run wider than gapMinFrac*width
|
||||
// as a column-separating gutter.
|
||||
func gapColumnCount(lines []pdf.TextBox, gapMinFrac, crossTol, binPt float64) int {
|
||||
n := len(lines)
|
||||
if n == 0 {
|
||||
return 1
|
||||
}
|
||||
minX0, width := pageExtent(lines)
|
||||
if width <= 0 {
|
||||
return 1
|
||||
}
|
||||
minGap := gapMinFrac * width
|
||||
nb := int(width/binPt) + 1
|
||||
cov := make([]int, nb)
|
||||
for _, b := range lines {
|
||||
i0 := int((b.X0 - minX0) / binPt)
|
||||
if i0 < 0 {
|
||||
i0 = 0
|
||||
}
|
||||
i1 := int((b.X1 - minX0) / binPt)
|
||||
if i1 > nb-1 {
|
||||
i1 = nb - 1
|
||||
}
|
||||
for i := i0; i <= i1; i++ {
|
||||
cov[i]++
|
||||
}
|
||||
}
|
||||
thr := crossTol * float64(n)
|
||||
cols := 1
|
||||
run := 0.0
|
||||
for _, c := range cov {
|
||||
if float64(c) < thr {
|
||||
run += binPt
|
||||
} else {
|
||||
if run >= minGap {
|
||||
cols++
|
||||
}
|
||||
run = 0
|
||||
}
|
||||
}
|
||||
if run >= minGap {
|
||||
cols++
|
||||
}
|
||||
return cols
|
||||
}
|
||||
|
||||
// balancedBodyK2 runs a forced k=2 KMeans on the BODY x0 (full-width front
|
||||
// matter excluded) and reports whether the split is a real two-column: two
|
||||
// clusters each holding >= minModeFrac of body lines, separated by >=
|
||||
// minSepFrac*width. Returns the 2 cluster centroids on success, plus the body
|
||||
// slice it clustered on so the caller's prune step counts the SAME line set
|
||||
// (otherwise full-width lines re-inflated into one column would let prune
|
||||
// collapse a balanced two-column page back to one).
|
||||
func balancedBodyK2(lines []pdf.TextBox, minModeFrac, minSepFrac float64) (bool, []float64, []pdf.TextBox) {
|
||||
minX0, width := pageExtent(lines)
|
||||
if width <= 0 {
|
||||
return false, nil, nil
|
||||
}
|
||||
body := dropFullWidth(lines, width)
|
||||
if len(body) < 4 {
|
||||
return false, nil, nil
|
||||
}
|
||||
x0s := make([]float64, len(body))
|
||||
for i, b := range body {
|
||||
x0s[i] = b.X0
|
||||
}
|
||||
indentTol := width * 0.12
|
||||
sx := snapX0s(x0s, minX0, indentTol)
|
||||
labels, cents := kmeansK2PlusPlus(sx, 42)
|
||||
if len(uniqueInts(labels)) < 2 {
|
||||
return false, nil, nil
|
||||
}
|
||||
counts := make(map[int]int, 2)
|
||||
for _, l := range labels {
|
||||
counts[l]++
|
||||
}
|
||||
minCount := math.MaxInt32
|
||||
for _, c := range counts {
|
||||
if c < minCount {
|
||||
minCount = c
|
||||
}
|
||||
}
|
||||
if float64(minCount) < minModeFrac*float64(len(body)) {
|
||||
return false, nil, nil
|
||||
}
|
||||
if math.Abs(cents[0]-cents[1]) < minSepFrac*width {
|
||||
return false, nil, nil
|
||||
}
|
||||
return true, cents, body
|
||||
}
|
||||
|
||||
// dropFullWidth removes lines whose width spans >=90% of the page text width
|
||||
// (titles / abstracts / headings that legitimately bridge a gutter).
|
||||
func dropFullWidth(lines []pdf.TextBox, width float64) []pdf.TextBox {
|
||||
fwThr := 0.9 * width
|
||||
out := make([]pdf.TextBox, 0, len(lines))
|
||||
for _, b := range lines {
|
||||
if b.X1-b.X0 < fwThr {
|
||||
out = append(out, b)
|
||||
}
|
||||
}
|
||||
if len(out) == 0 {
|
||||
// Every line is full-width: there is no narrow body to form a second
|
||||
// column. Return nil (not the original lines) so the caller's
|
||||
// len(body) < 4 guard treats the page as a single column instead of
|
||||
// pushing the whole page through the balance gate, which could
|
||||
// mis-split a full-width single column whose x0 happens to be bimodal.
|
||||
return nil
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// pageExtent returns minX0 (leftmost x0) and the text width (maxX1 - minX0).
|
||||
func pageExtent(lines []pdf.TextBox) (minX0, width float64) {
|
||||
minX0 = math.MaxFloat64
|
||||
maxX1 := 0.0
|
||||
for _, b := range lines {
|
||||
if b.X0 < minX0 {
|
||||
minX0 = b.X0
|
||||
}
|
||||
if b.X1 > maxX1 {
|
||||
maxX1 = b.X1
|
||||
}
|
||||
}
|
||||
return minX0, maxX1 - minX0
|
||||
}
|
||||
|
||||
// snapX0s pulls x0 values within indentTol of minX0 back to minX0, so slightly
|
||||
// indented lines still cluster with the left edge (mirrors _assign_column).
|
||||
func snapX0s(x0s []float64, minX0, indentTol float64) []float64 {
|
||||
out := make([]float64, len(x0s))
|
||||
for i, v := range x0s {
|
||||
if math.Abs(v-minX0) < indentTol {
|
||||
out[i] = minX0
|
||||
} else {
|
||||
out[i] = v
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// kmeansK2PlusPlus is a density-aware k=2 clustering (k-means++ init, single
|
||||
// Lloyd pass). Unlike util.KMeans1D (even-spaced init, a range partition), the
|
||||
// first center is a random data point and the second is the farthest point, so
|
||||
// it respects natural x0 density — required for the balance check to reject a
|
||||
// single column whose x0 merely has a wide range. Deterministic via seed.
|
||||
func kmeansK2PlusPlus(x0s []float64, seed int64) ([]int, []float64) {
|
||||
n := len(x0s)
|
||||
labels := make([]int, n)
|
||||
if n == 0 {
|
||||
return labels, nil
|
||||
}
|
||||
rng := rand.New(rand.NewSource(seed))
|
||||
first := rng.Intn(n)
|
||||
c0 := x0s[first]
|
||||
bestJ, bestD := 0, -1.0
|
||||
for j, v := range x0s {
|
||||
d := (v - c0) * (v - c0)
|
||||
if d > bestD {
|
||||
bestD, bestJ = d, j
|
||||
}
|
||||
}
|
||||
c1 := x0s[bestJ]
|
||||
cents := []float64{c0, c1}
|
||||
for iter := 0; iter < 100; iter++ {
|
||||
changed := false
|
||||
for i, v := range x0s {
|
||||
bestC := 0
|
||||
if math.Abs(v-c1) < math.Abs(v-c0) {
|
||||
bestC = 1
|
||||
}
|
||||
if labels[i] != bestC {
|
||||
changed = true
|
||||
labels[i] = bestC
|
||||
}
|
||||
}
|
||||
if !changed {
|
||||
break
|
||||
}
|
||||
sum := [2]float64{}
|
||||
cnt := [2]int{}
|
||||
for i, v := range x0s {
|
||||
sum[labels[i]] += v
|
||||
cnt[labels[i]]++
|
||||
}
|
||||
for c := 0; c < 2; c++ {
|
||||
if cnt[c] > 0 {
|
||||
cents[c] = sum[c] / float64(cnt[c])
|
||||
}
|
||||
}
|
||||
}
|
||||
return labels, cents
|
||||
}
|
||||
|
||||
// kmeansCentroids returns the k cluster centroids from util.KMeans1D on the
|
||||
// snapped x0s of all lines; used to partition a page when gap reports >=2.
|
||||
func kmeansCentroids(lines []pdf.TextBox, k int, width float64) []float64 {
|
||||
minX0, _ := pageExtent(lines)
|
||||
x0s := make([]float64, len(lines))
|
||||
for i, b := range lines {
|
||||
x0s[i] = b.X0
|
||||
}
|
||||
sx := snapX0s(x0s, minX0, width*0.12)
|
||||
_, cents := util.KMeans1D(sx, k)
|
||||
return cents
|
||||
}
|
||||
|
||||
// assignColIDs sets ColID for a page's boxes by nearest centroid, remapped so
|
||||
// the leftmost centroid becomes column 0.
|
||||
func assignColIDs(boxes, result []pdf.TextBox, indices []int, k int, cents []float64) {
|
||||
if k <= 1 || len(cents) == 0 {
|
||||
for _, idx := range indices {
|
||||
result[idx].ColID = 0
|
||||
}
|
||||
return
|
||||
}
|
||||
order := make([]int, len(cents))
|
||||
idxByVal := make([]int, len(cents))
|
||||
for i := range cents {
|
||||
idxByVal[i] = i
|
||||
}
|
||||
sort.Slice(idxByVal, func(a, b int) bool { return cents[idxByVal[a]] < cents[idxByVal[b]] })
|
||||
for newL, oldL := range idxByVal {
|
||||
order[oldL] = newL
|
||||
}
|
||||
for _, idx := range indices {
|
||||
x := boxes[idx].X0
|
||||
best, bestD := 0, math.Abs(x-cents[0])
|
||||
for c := 1; c < len(cents); c++ {
|
||||
if d := math.Abs(x - cents[c]); d < bestD {
|
||||
bestD, best = d, c
|
||||
}
|
||||
}
|
||||
result[idx].ColID = order[best]
|
||||
}
|
||||
}
|
||||
|
||||
func uniqueInts(xs []int) []int {
|
||||
seen := make(map[int]struct{}, len(xs))
|
||||
for _, x := range xs {
|
||||
seen[x] = struct{}{}
|
||||
}
|
||||
out := make([]int, 0, len(seen))
|
||||
for x := range seen {
|
||||
out = append(out, x)
|
||||
}
|
||||
return out
|
||||
}
|
||||
@@ -0,0 +1,109 @@
|
||||
package layout
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
pdf "ragflow/internal/deepdoc/parser/pdf/type"
|
||||
)
|
||||
|
||||
// Synthetic fixtures regression test.
|
||||
//
|
||||
// This file is the CI-runnable guardrail for the column detector. Unlike
|
||||
// TestAssignColumnCombined_Labeled (which needs the gitignored 70-page label
|
||||
// sheet + 343MB corpus and is therefore skipped in CI), these fixtures are
|
||||
// hand-built []pdf.TextBox geometries committed to the repo, so they run on
|
||||
// every `go test ./...` with no external data.
|
||||
//
|
||||
// Green assertions pin behavior that must NOT regress. The title-bridged
|
||||
// double is a skipped TODO that captures the #18079 acceptance target: a 2D
|
||||
// spatial column detector should split it into 2, but the current x0-based
|
||||
// balance gate rejects it (sparse minority + x0 overlap).
|
||||
|
||||
// columnCount runs AssignColumn on one page of boxes and returns the number of
|
||||
// distinct columns (max ColID + 1).
|
||||
func columnCount(boxes []pdf.TextBox) int {
|
||||
in := make([]pdf.TextBox, len(boxes))
|
||||
copy(in, boxes)
|
||||
for i := range in {
|
||||
in[i].PageNumber = 0
|
||||
}
|
||||
res := AssignColumn(in)
|
||||
k := 1
|
||||
for _, b := range res {
|
||||
if b.ColID+1 > k {
|
||||
k = b.ColID + 1
|
||||
}
|
||||
}
|
||||
return k
|
||||
}
|
||||
|
||||
// stackedColumn builds n text boxes in a vertical column at [x0,x1], starting
|
||||
// at top0 with line spacing dy.
|
||||
func stackedColumn(x0, x1, n int, top0, dy float64) []pdf.TextBox {
|
||||
boxes := make([]pdf.TextBox, n)
|
||||
for i := 0; i < n; i++ {
|
||||
top := top0 + float64(i)*dy
|
||||
boxes[i] = pdf.TextBox{X0: float64(x0), X1: float64(x1), Top: top, Bottom: top + 8}
|
||||
}
|
||||
return boxes
|
||||
}
|
||||
|
||||
func concat(dst, src []pdf.TextBox) []pdf.TextBox { return append(dst, src...) }
|
||||
|
||||
// A single reading column must stay single (no over-split).
|
||||
func TestSyntheticSingleColumn(t *testing.T) {
|
||||
boxes := stackedColumn(50, 240, 10, 10, 10)
|
||||
if got := columnCount(boxes); got != 1 {
|
||||
t.Errorf("single column: got %d, want 1", got)
|
||||
}
|
||||
}
|
||||
|
||||
// A balanced two-column page must be recovered by the balance gate.
|
||||
func TestSyntheticBalancedDouble(t *testing.T) {
|
||||
boxes := concat(stackedColumn(50, 240, 10, 10, 10), stackedColumn(270, 460, 10, 10, 10))
|
||||
if got := columnCount(boxes); got != 2 {
|
||||
t.Errorf("balanced double: got %d, want 2", got)
|
||||
}
|
||||
}
|
||||
|
||||
// Narrow side-by-side columns are a table, not text columns: returned as 1.
|
||||
func TestSyntheticTableNarrowColumns(t *testing.T) {
|
||||
var boxes []pdf.TextBox
|
||||
for _, c := range [][2]int{{50, 130}, {200, 280}, {350, 430}} {
|
||||
boxes = concat(boxes, stackedColumn(c[0], c[1], 5, 10, 10))
|
||||
}
|
||||
if got := columnCount(boxes); got != 1 {
|
||||
t.Errorf("narrow table: got %d, want 1", got)
|
||||
}
|
||||
}
|
||||
|
||||
// A real left column plus a sparse right column: the balance gate must reject
|
||||
// the sparse column (minority < 30%), keeping the page single. Guards the
|
||||
// sparse-column prune (minColLineFrac) against regression.
|
||||
func TestSyntheticSparseSecondColumn(t *testing.T) {
|
||||
boxes := concat(stackedColumn(50, 240, 30, 10, 10), stackedColumn(300, 460, 3, 10, 20))
|
||||
if got := columnCount(boxes); got != 1 {
|
||||
t.Errorf("sparse second column: got %d, want 1", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TODO #18079: a 2D spatial column detector should split this into 2.
|
||||
//
|
||||
// The page has a full-width title block at the top that bridges the gutter,
|
||||
// and below it two side-by-side columns where the right column is sparse and
|
||||
// its x0 overlaps the left column's x0 range. The x0-based balance gate
|
||||
// therefore rejects it (minority < 30%, x0 overlap) and the page is reported
|
||||
// as single. Skipped until #18079 lands; then unskip and assert got == 2.
|
||||
func TestSyntheticTitleBridgedDouble(t *testing.T) {
|
||||
t.Skip("TODO #18079: title-bridged double should be detected as 2 columns")
|
||||
boxes := concat(
|
||||
concat(
|
||||
stackedColumn(50, 440, 3, 10, 10), // full-width title (bridges the gutter)
|
||||
stackedColumn(50, 240, 30, 40, 10), // left column
|
||||
),
|
||||
stackedColumn(200, 440, 4, 40, 20), // right column: sparse, x0 overlaps left
|
||||
)
|
||||
if got := columnCount(boxes); got != 2 {
|
||||
t.Errorf("title-bridged double: got %d, want 2", got)
|
||||
}
|
||||
}
|
||||
283
internal/deepdoc/parser/pdf/layout/combined_column_test.go
Normal file
283
internal/deepdoc/parser/pdf/layout/combined_column_test.go
Normal file
@@ -0,0 +1,283 @@
|
||||
package layout
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"math"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"testing"
|
||||
|
||||
pdf "ragflow/internal/deepdoc/parser/pdf/type"
|
||||
)
|
||||
|
||||
// ---- charspy JSON binding + line reconstruction (mirrors tool-py
|
||||
// extract_column_divergence._reconstruct_lines so the Go detector is scored on
|
||||
// the SAME line boxes the Python reference used) ----
|
||||
|
||||
type charBox struct {
|
||||
Text string `json:"text"`
|
||||
X0 float64 `json:"x0"`
|
||||
X1 float64 `json:"x1"`
|
||||
Top float64 `json:"top"`
|
||||
Bottom float64 `json:"bottom"`
|
||||
Size float64 `json:"size"`
|
||||
}
|
||||
|
||||
type charPage struct {
|
||||
Pages [][]charBox `json:"pages"`
|
||||
}
|
||||
|
||||
func reconstructLines(chars []charBox) []pdf.TextBox {
|
||||
if len(chars) == 0 {
|
||||
return nil
|
||||
}
|
||||
sizes := make([]float64, len(chars))
|
||||
for i, c := range chars {
|
||||
sizes[i] = c.Size
|
||||
}
|
||||
sort.Float64s(sizes)
|
||||
medSize := sizes[len(sizes)/2]
|
||||
vTol := math.Max(medSize*0.8, 4.0)
|
||||
hGap := math.Max(medSize*3.0, 15.0)
|
||||
|
||||
rows := map[int][]charBox{}
|
||||
for _, c := range chars {
|
||||
key := int(math.Round(c.Top / vTol))
|
||||
rows[key] = append(rows[key], c)
|
||||
}
|
||||
keys := make([]int, 0, len(rows))
|
||||
for k := range rows {
|
||||
keys = append(keys, k)
|
||||
}
|
||||
sort.Ints(keys)
|
||||
|
||||
var lines []pdf.TextBox
|
||||
for _, key := range keys {
|
||||
row := rows[key]
|
||||
sort.Slice(row, func(a, b int) bool { return row[a].X0 < row[b].X0 })
|
||||
var cur *pdf.TextBox
|
||||
for _, c := range row {
|
||||
if cur == nil {
|
||||
cur = &pdf.TextBox{X0: c.X0, X1: c.X1, Top: c.Top, Bottom: c.Bottom, Text: c.Text}
|
||||
} else if c.X0-cur.X1 > hGap {
|
||||
lines = append(lines, *cur)
|
||||
cur = &pdf.TextBox{X0: c.X0, X1: c.X1, Top: c.Top, Bottom: c.Bottom, Text: c.Text}
|
||||
} else if c.X1 > cur.X1 {
|
||||
cur.X1 = c.X1
|
||||
}
|
||||
}
|
||||
if cur != nil {
|
||||
lines = append(lines, *cur)
|
||||
}
|
||||
}
|
||||
return lines
|
||||
}
|
||||
|
||||
// labeledPage mirrors the first 18 (locked) entries of
|
||||
// tool-py/column_labeling_sheet.json — human-confirmed column truth.
|
||||
type labeledPage struct {
|
||||
pdf string
|
||||
page int
|
||||
truth int
|
||||
}
|
||||
|
||||
// sheetEntry mirrors the JSON shape of tool-py/column_labeling_sheet.json.
|
||||
type sheetEntry struct {
|
||||
PDF string `json:"pdf"`
|
||||
Page int `json:"page"`
|
||||
TruthK *int `json:"truth_k"`
|
||||
}
|
||||
|
||||
// loadLabeledPages reads every page with a non-null truth_k from the shared
|
||||
// label sheet (single source of truth, also used by the Python reference).
|
||||
// Pages still marked null (need the rendered PDF to decide) are skipped.
|
||||
func loadLabeledPages(t *testing.T) []labeledPage {
|
||||
sheetPath := "../tool-py/column_labeling_sheet.json"
|
||||
raw, err := os.ReadFile(sheetPath)
|
||||
if err != nil {
|
||||
t.Skipf("label sheet not found at %s: %v", sheetPath, err)
|
||||
}
|
||||
var entries []sheetEntry
|
||||
if err := json.Unmarshal(raw, &entries); err != nil {
|
||||
t.Fatalf("unmarshal sheet: %v", err)
|
||||
}
|
||||
var out []labeledPage
|
||||
for _, e := range entries {
|
||||
if e.TruthK == nil {
|
||||
continue
|
||||
}
|
||||
out = append(out, labeledPage{pdf: e.PDF, page: e.Page, truth: *e.TruthK})
|
||||
}
|
||||
if len(out) == 0 {
|
||||
t.Skip("no labeled pages in sheet")
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// TestAssignColumnCombined_Labeled scores the gap+KMeans hybrid against the
|
||||
// LIVE gap-only baseline on every labeled page in the shared sheet (the sheet
|
||||
// is the single source of truth, also used by the Python reference). Hard
|
||||
// pages were labeled from the page coverage profile + strip map in
|
||||
// tool-py/analyze_hardcases.py, tagged H (high) / L (low) confidence in
|
||||
// truth_note. The 18-page labeled sample that motivated this detector reported
|
||||
// KMeans 11.1% / gap 61.1% accuracy, but that sample was hand-picked and must
|
||||
// not be used as a target. The per-page / single / double / ACC lines below
|
||||
// are the real signal, and the only asserted contract is "combined must not
|
||||
// regress gap by >5pt".
|
||||
func TestAssignColumnCombined_Labeled(t *testing.T) {
|
||||
charspyDir := "../testdata/charspy"
|
||||
if _, err := os.Stat(charspyDir); err != nil {
|
||||
t.Skipf("charspy corpus not found at %s (run from package dir): %v", charspyDir, err)
|
||||
}
|
||||
cases := loadLabeledPages(t)
|
||||
var correct, fs, miss, sCorrect, sTotal, dCorrect, dTotal int
|
||||
var gCorrect, gFs, gMiss, gsCorrect, gsTotal, gdCorrect, gdTotal int
|
||||
for _, c := range cases {
|
||||
raw, err := os.ReadFile(filepath.Join(charspyDir, c.pdf))
|
||||
if err != nil {
|
||||
t.Fatalf("read %s: %v", c.pdf, err)
|
||||
}
|
||||
var cp charPage
|
||||
if err := json.Unmarshal(raw, &cp); err != nil {
|
||||
t.Fatalf("unmarshal %s: %v", c.pdf, err)
|
||||
}
|
||||
lines := reconstructLines(cp.Pages[c.page])
|
||||
boxes := make([]pdf.TextBox, len(lines))
|
||||
for i, l := range lines {
|
||||
boxes[i] = l
|
||||
boxes[i].PageNumber = 0
|
||||
}
|
||||
res := AssignColumn(boxes)
|
||||
k := 1
|
||||
for _, b := range res {
|
||||
if b.ColID+1 > k {
|
||||
k = b.ColID + 1
|
||||
}
|
||||
}
|
||||
// Live gap-only baseline on the SAME lines (combined wraps gap, so
|
||||
// this is the honest reference — not the old 18-page Python 61.1%).
|
||||
g := gapColumnCount(lines, 0.04, 0.15, 2.0)
|
||||
record := func(got, truth int, corr, fsC, missC, sCorr, sTot, dCorr, dTot *int) {
|
||||
if got == truth {
|
||||
*corr++
|
||||
} else if got > truth {
|
||||
*fsC++
|
||||
} else {
|
||||
*missC++
|
||||
}
|
||||
if truth == 1 {
|
||||
*sTot++
|
||||
if got == truth {
|
||||
*sCorr++
|
||||
}
|
||||
} else {
|
||||
*dTot++
|
||||
if got == truth {
|
||||
*dCorr++
|
||||
}
|
||||
}
|
||||
}
|
||||
record(k, c.truth, &correct, &fs, &miss, &sCorrect, &sTotal, &dCorrect, &dTotal)
|
||||
record(g, c.truth, &gCorrect, &gFs, &gMiss, &gsCorrect, &gsTotal, &gdCorrect, &gdTotal)
|
||||
tag := "OK"
|
||||
if k != c.truth {
|
||||
tag = "WRONG"
|
||||
}
|
||||
t.Logf("[%s] %s p%d truth=%d got=%d (gap=%d) comb single=%d/%d double=%d/%d gap single=%d/%d double=%d/%d",
|
||||
tag, c.pdf, c.page, c.truth, k, g, sCorrect, sTotal, dCorrect, dTotal, gsCorrect, gsTotal, gdCorrect, gdTotal)
|
||||
}
|
||||
n := len(cases)
|
||||
acc := 100.0 * float64(correct) / float64(n)
|
||||
gapAcc := 100.0 * float64(gCorrect) / float64(n)
|
||||
t.Logf("Go gap+KMeans combined: ACC=%.1f%% false-split=%d miss=%d single=%d/%d double=%d/%d (n=%d)",
|
||||
acc, fs, miss, sCorrect, sTotal, dCorrect, dTotal, n)
|
||||
t.Logf("Go gap-only baseline: ACC=%.1f%% false-split=%d miss=%d single=%d/%d double=%d/%d (n=%d)",
|
||||
gapAcc, gFs, gMiss, gsCorrect, gsTotal, gdCorrect, gdTotal, n)
|
||||
// The hybrid's contract: gap is the safe base, KMeans only ADDS recoveries
|
||||
// (double columns gap misses). So combined must not regress gap by more
|
||||
// than a small tolerance. The 75% / 88.9% figures were tuned targets on
|
||||
// the comfortable 18 pages and must not be asserted (circular). Watch the
|
||||
// two ACC lines above as harder pages are labeled and added.
|
||||
if acc < gapAcc-5.0 {
|
||||
t.Errorf("combined ACC=%.1f%% regresses gap-only baseline %.1f%% by >5pt (KMeans addition is harmful)", acc, gapAcc)
|
||||
}
|
||||
}
|
||||
|
||||
// TestKmeansK2PlusPlus_Smoke sanity-checks the density-aware k=2 gate on a
|
||||
// clean bimodal vs single-mode input.
|
||||
func TestKmeansK2PlusPlus_Smoke(t *testing.T) {
|
||||
// Two well-separated modes, balanced -> two real clusters.
|
||||
bimodal := []float64{10, 11, 12, 13, 90, 91, 92, 93}
|
||||
labels, cents := kmeansK2PlusPlus(bimodal, 42)
|
||||
if len(uniqueInts(labels)) != 2 {
|
||||
t.Errorf("bimodal: expected 2 clusters, got %d", len(uniqueInts(labels)))
|
||||
}
|
||||
// kmeansK2PlusPlus does not guarantee centroid ordering, so assert the
|
||||
// property it DOES guarantee: a bimodal split yields two well-separated
|
||||
// centroids. (assignColIDs re-sorts them for ColID assignment.)
|
||||
if math.Abs(cents[0]-cents[1]) < 50 {
|
||||
t.Errorf("bimodal: centroids should be well separated, got %v", cents)
|
||||
}
|
||||
// Single tight mode -> still 2 labels after Lloyd, but minority tiny.
|
||||
tight := []float64{10, 10, 10, 10, 11, 10, 10, 10}
|
||||
labels2, _ := kmeansK2PlusPlus(tight, 42)
|
||||
counts := map[int]int{}
|
||||
for _, l := range labels2 {
|
||||
counts[l]++
|
||||
}
|
||||
minC := math.MaxInt32
|
||||
for _, c := range counts {
|
||||
if c < minC {
|
||||
minC = c
|
||||
}
|
||||
}
|
||||
if float64(minC) >= 0.30*float64(len(tight)) {
|
||||
t.Errorf("tight single-mode: minority %.0f%% should be <30%%", 100*float64(minC)/float64(len(tight)))
|
||||
}
|
||||
}
|
||||
|
||||
// TestAssignColumn_FullWidthBridgeKeepsBalancedTwoColumn is a regression test
|
||||
// for the balance-gate / prune asymmetry: the balance gate clusters the BODY
|
||||
// (full-width front matter excluded), but pruneColumns must count the SAME
|
||||
// body — not all lines. A real but minority-right two-column block dominated
|
||||
// by full-width title/abstract lines must stay two columns; otherwise the
|
||||
// full-width lines inflate the left column and prune wrongly collapses the
|
||||
// page to one.
|
||||
func TestAssignColumn_FullWidthBridgeKeepsBalancedTwoColumn(t *testing.T) {
|
||||
const (
|
||||
minX0 = 100.0
|
||||
maxX1 = 700.0
|
||||
)
|
||||
var boxes []pdf.TextBox
|
||||
add := func(x0, x1 float64) {
|
||||
top := float64(len(boxes)) * 12
|
||||
boxes = append(boxes, pdf.TextBox{
|
||||
PageNumber: 0, X0: x0, X1: x1, Top: top, Bottom: top + 10, Text: "b",
|
||||
})
|
||||
}
|
||||
// Body: a genuine two-column block with a real gutter. The right column is
|
||||
// the minority (6 of 20 body lines, i.e. 30% — exactly the gate floor).
|
||||
for i := 0; i < 14; i++ {
|
||||
add(minX0, 350) // left column
|
||||
}
|
||||
for i := 0; i < 6; i++ {
|
||||
add(450, maxX1) // right column
|
||||
}
|
||||
// 60 full-width front-matter lines (title + abstract) bridge the gutter so
|
||||
// gap reports a single column. They must not be re-counted by prune.
|
||||
for i := 0; i < 60; i++ {
|
||||
add(minX0, maxX1) // full width
|
||||
}
|
||||
|
||||
res := AssignColumn(boxes)
|
||||
k := 1
|
||||
for _, b := range res {
|
||||
if b.ColID+1 > k {
|
||||
k = b.ColID + 1
|
||||
}
|
||||
}
|
||||
if k != 2 {
|
||||
t.Fatalf("balanced two-column page with full-width front matter collapsed to %d column(s); want 2", k)
|
||||
}
|
||||
}
|
||||
@@ -14,148 +14,8 @@ import (
|
||||
)
|
||||
|
||||
// ---- Column assignment ----
|
||||
|
||||
// AssignColumn groups boxes into columns on each page by KMeans x0 clustering
|
||||
// with silhouette score selection, matching Python's _assign_column().
|
||||
//
|
||||
// Python: pdf_parser.py:739 _assign_column()
|
||||
func AssignColumn(boxes []pdf.TextBox) []pdf.TextBox {
|
||||
if len(boxes) == 0 {
|
||||
return boxes
|
||||
}
|
||||
|
||||
pageGroups, sortedPages := groupBoxesByPage(boxes)
|
||||
|
||||
result := make([]pdf.TextBox, len(boxes))
|
||||
copy(result, boxes)
|
||||
|
||||
// Step A: per-page best k using silhouette score.
|
||||
pageCols := make(map[int]int)
|
||||
for _, pg := range sortedPages {
|
||||
indices := pageGroups[pg]
|
||||
determineBestKForPage(boxes, result, indices, pg, pageCols)
|
||||
}
|
||||
|
||||
// Step B: assign col_id per page using per-page best k.
|
||||
// Labels are remapped by centroid x-order: leftmost column → 0.
|
||||
for _, pg := range sortedPages {
|
||||
indices := pageGroups[pg]
|
||||
assignColIDsForPage(boxes, result, indices, pg, pageCols)
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
// determineBestKForPage finds the best number of clusters (k) for a page using silhouette score
|
||||
func determineBestKForPage(boxes, result []pdf.TextBox, indices []int, pg int, pageCols map[int]int) {
|
||||
n := len(indices)
|
||||
if n < 2 {
|
||||
pageCols[pg] = 1
|
||||
for _, idx := range indices {
|
||||
result[idx].ColID = 0
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
x0s, minX0, maxX1 := extractX0Values(boxes, indices)
|
||||
pageWidth := maxX1 - minX0
|
||||
indentTol := pageWidth * 0.12
|
||||
applyIndentTolerance(x0s, minX0, indentTol)
|
||||
|
||||
bestK, _ := findBestK(x0s, n)
|
||||
pageCols[pg] = bestK
|
||||
}
|
||||
|
||||
// extractX0Values extracts x0 coordinates from boxes on a page and finds minX0 and maxX1
|
||||
func extractX0Values(boxes []pdf.TextBox, indices []int) (x0s []float64, minX0 float64, maxX1 float64) {
|
||||
n := len(indices)
|
||||
x0s = make([]float64, n)
|
||||
minX0 = math.MaxFloat64
|
||||
maxX1 = 0.0
|
||||
for i, idx := range indices {
|
||||
x0s[i] = boxes[idx].X0
|
||||
if x0s[i] < minX0 {
|
||||
minX0 = x0s[i]
|
||||
}
|
||||
if boxes[idx].X1 > maxX1 {
|
||||
maxX1 = boxes[idx].X1
|
||||
}
|
||||
}
|
||||
return x0s, minX0, maxX1
|
||||
}
|
||||
|
||||
// applyIndentTolerance adjusts x0 values that are close to minX0 to improve clustering
|
||||
func applyIndentTolerance(x0s []float64, minX0, indentTol float64) {
|
||||
for i := range x0s {
|
||||
if math.Abs(x0s[i]-minX0) < indentTol {
|
||||
x0s[i] = minX0
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// findBestK tries k from 1 to min(4, n) and returns the k with the best silhouette score
|
||||
func findBestK(x0s []float64, n int) (bestK int, bestScore float64) {
|
||||
maxTry := min(4, n)
|
||||
if maxTry < 2 {
|
||||
maxTry = 1
|
||||
}
|
||||
bestK, bestScore = 1, -1.0
|
||||
|
||||
for k := 1; k <= maxTry; k++ {
|
||||
labels, _ := util.KMeans1D(x0s, k)
|
||||
var score float64
|
||||
if k > 1 {
|
||||
score = util.Silhouette1D(x0s, labels)
|
||||
}
|
||||
// score = 0 for k=1; score = -1 if silhouette undefined.
|
||||
if score > bestScore {
|
||||
bestScore = score
|
||||
bestK = k
|
||||
}
|
||||
}
|
||||
return bestK, bestScore
|
||||
}
|
||||
|
||||
// assignColIDsForPage assigns column IDs to boxes on a page using the best k
|
||||
func assignColIDsForPage(boxes, result []pdf.TextBox, indices []int, pg int, pageCols map[int]int) {
|
||||
if len(indices) == 0 {
|
||||
return
|
||||
}
|
||||
k := pageCols[pg]
|
||||
if len(indices) < k {
|
||||
k = 1
|
||||
}
|
||||
|
||||
x0s := make([]float64, len(indices))
|
||||
for i, idx := range indices {
|
||||
x0s[i] = boxes[idx].X0
|
||||
}
|
||||
|
||||
labels, centroids := util.KMeans1D(x0s, k)
|
||||
remap := remapLabelsByCentroidOrder(centroids)
|
||||
|
||||
for i, idx := range indices {
|
||||
result[idx].ColID = remap[labels[i]]
|
||||
}
|
||||
}
|
||||
|
||||
// remapLabelsByCentroidOrder remaps cluster labels so leftmost column = 0
|
||||
func remapLabelsByCentroidOrder(centroids []float64) map[int]int {
|
||||
type clPair struct {
|
||||
center float64
|
||||
label int
|
||||
}
|
||||
var pairs []clPair
|
||||
for lbl, c := range centroids {
|
||||
pairs = append(pairs, clPair{c, lbl})
|
||||
}
|
||||
sort.Slice(pairs, func(i, j int) bool { return pairs[i].center < pairs[j].center })
|
||||
remap := make(map[int]int, len(centroids))
|
||||
for newL, p := range pairs {
|
||||
remap[p.label] = newL
|
||||
}
|
||||
return remap
|
||||
}
|
||||
// AssignColumn is implemented in combined_column.go (gap + KMeans hybrid).
|
||||
|
||||
// ---- Text merge (horizontal) ----
|
||||
|
||||
|
||||
@@ -22,14 +22,15 @@ func newTestTextBox(page int, x0, x1, top, bottom float64, text string) pdf.Text
|
||||
|
||||
func TestAssignColumn(t *testing.T) {
|
||||
boxes := []pdf.TextBox{
|
||||
{PageNumber: 0, X0: 50, Text: "col0-left"},
|
||||
{PageNumber: 0, X0: 55, Text: "col0-mid"},
|
||||
{PageNumber: 0, X0: 400, Text: "col1"},
|
||||
{PageNumber: 1, X0: 50, Text: "pg1-col0"},
|
||||
{PageNumber: 0, X0: 50, X1: 250, Text: "col0-left"},
|
||||
{PageNumber: 0, X0: 55, X1: 250, Text: "col0-mid"},
|
||||
{PageNumber: 0, X0: 400, X1: 600, Text: "col1"},
|
||||
{PageNumber: 0, X0: 410, X1: 610, Text: "col1-b"},
|
||||
{PageNumber: 1, X0: 50, X1: 250, Text: "pg1-col0"},
|
||||
}
|
||||
result := AssignColumn(boxes)
|
||||
if len(result) != 4 {
|
||||
t.Fatal("expected 4 boxes")
|
||||
if len(result) != 5 {
|
||||
t.Fatal("expected 5 boxes")
|
||||
}
|
||||
if result[0].ColID != result[1].ColID {
|
||||
t.Error("boxes 0 and 1 (close x0) should be same column")
|
||||
@@ -922,183 +923,3 @@ func TestProcessPageBoxes_NoMerge(t *testing.T) {
|
||||
t.Errorf("expected 2 boxes, got %d", len(result))
|
||||
}
|
||||
}
|
||||
|
||||
// ── Column-assignment helper tests ──────────────────────────────────
|
||||
|
||||
func TestExtractX0Values(t *testing.T) {
|
||||
boxes := []pdf.TextBox{
|
||||
{PageNumber: 0, X0: 50, X1: 200},
|
||||
{PageNumber: 0, X0: 30, X1: 100},
|
||||
{PageNumber: 0, X0: 80, X1: 300},
|
||||
}
|
||||
x0s, minX0, maxX1 := extractX0Values(boxes, []int{0, 1, 2})
|
||||
if len(x0s) != 3 {
|
||||
t.Fatalf("expected 3 x0s, got %d", len(x0s))
|
||||
}
|
||||
if x0s[0] != 50 || x0s[1] != 30 || x0s[2] != 80 {
|
||||
t.Errorf("x0s mismatch: %v", x0s)
|
||||
}
|
||||
if minX0 != 30 {
|
||||
t.Errorf("minX0 = %v, want 30", minX0)
|
||||
}
|
||||
if maxX1 != 300 {
|
||||
t.Errorf("maxX1 = %v, want 300", maxX1)
|
||||
}
|
||||
}
|
||||
|
||||
func TestApplyIndentTolerance(t *testing.T) {
|
||||
values := []float64{100, 105, 200, 210}
|
||||
applyIndentTolerance(values, 100, 10)
|
||||
if values[0] != 100 || values[1] != 100 {
|
||||
t.Errorf("close x0s should be adjusted to minX0: %v", values)
|
||||
}
|
||||
if values[2] != 200 || values[3] != 210 {
|
||||
t.Errorf("distant x0s should remain unchanged: %v", values)
|
||||
}
|
||||
}
|
||||
|
||||
func TestApplyIndentTolerance_Zero(t *testing.T) {
|
||||
values := []float64{100, 101, 200}
|
||||
applyIndentTolerance(values, 100, 0)
|
||||
if values[1] != 101 {
|
||||
t.Errorf("zero tolerance: x0s should be unchanged, got %v", values)
|
||||
}
|
||||
}
|
||||
|
||||
func TestApplyIndentTolerance_Negative(t *testing.T) {
|
||||
values := []float64{-100, -95, 0, 50}
|
||||
applyIndentTolerance(values, -100, 10)
|
||||
if values[0] != -100 || values[1] != -100 {
|
||||
t.Errorf("negative x0s close to minX0 should be adjusted: %v", values)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFindBestK_SingleCluster(t *testing.T) {
|
||||
// Note: KMeans1D uses random initialization, so non-identical values
|
||||
// may occasionally produce k>1. This test verifies the function runs
|
||||
// without error and returns k>=1 (not a correctness check).
|
||||
x0s := []float64{100, 99, 101}
|
||||
bestK, _ := findBestK(x0s, len(x0s))
|
||||
if bestK < 1 {
|
||||
t.Errorf("expected bestK>=1, got %d", bestK)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFindBestK_TwoColumns(t *testing.T) {
|
||||
x0s := []float64{50, 55, 60, 200, 210, 220}
|
||||
bestK, _ := findBestK(x0s, len(x0s))
|
||||
if bestK != 2 {
|
||||
t.Errorf("two columns: expected bestK=2, got %d", bestK)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFindBestK_OneValue(t *testing.T) {
|
||||
x0s := []float64{100}
|
||||
bestK, _ := findBestK(x0s, len(x0s))
|
||||
if bestK != 1 {
|
||||
t.Errorf("single value: expected bestK=1, got %d", bestK)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFindBestK_Identical(t *testing.T) {
|
||||
x0s := []float64{100, 100, 100, 100, 100}
|
||||
bestK, _ := findBestK(x0s, len(x0s))
|
||||
if bestK != 1 {
|
||||
t.Errorf("identical values: expected bestK=1, got %d", bestK)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRemapLabelsByCentroidOrder_Ordered(t *testing.T) {
|
||||
centroids := []float64{50, 200, 400}
|
||||
remap := remapLabelsByCentroidOrder(centroids)
|
||||
if remap[0] != 0 || remap[1] != 1 || remap[2] != 2 {
|
||||
t.Errorf("ordered centroids: expected 0->0,1->1,2->2, got %v", remap)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRemapLabelsByCentroidOrder_Unordered(t *testing.T) {
|
||||
centroids := []float64{200, 50, 400}
|
||||
remap := remapLabelsByCentroidOrder(centroids)
|
||||
if remap[0] != 1 || remap[1] != 0 || remap[2] != 2 {
|
||||
t.Errorf("unordered centroids: expected {0:1,1:0,2:2}, got %v", remap)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRemapLabelsByCentroidOrder_Nil(t *testing.T) {
|
||||
remap := remapLabelsByCentroidOrder(nil)
|
||||
if len(remap) != 0 {
|
||||
t.Errorf("nil centroids: expected empty map, got %v", remap)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDetermineBestKForPage_SingleBox(t *testing.T) {
|
||||
boxes := []pdf.TextBox{{PageNumber: 0, X0: 100, X1: 200}}
|
||||
result := make([]pdf.TextBox, len(boxes))
|
||||
copy(result, boxes)
|
||||
pageCols := make(map[int]int)
|
||||
determineBestKForPage(boxes, result, []int{0}, 0, pageCols)
|
||||
if pageCols[0] != 1 {
|
||||
t.Errorf("single box: expected pageCols[0]=1, got %d", pageCols[0])
|
||||
}
|
||||
if result[0].ColID != 0 {
|
||||
t.Errorf("single box: expected ColID=0, got %d", result[0].ColID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDetermineBestKForPage_TwoColumns(t *testing.T) {
|
||||
boxes := []pdf.TextBox{
|
||||
{PageNumber: 0, X0: 50, X1: 100},
|
||||
{PageNumber: 0, X0: 55, X1: 100},
|
||||
{PageNumber: 0, X0: 300, X1: 400},
|
||||
{PageNumber: 0, X0: 310, X1: 400},
|
||||
}
|
||||
result := make([]pdf.TextBox, len(boxes))
|
||||
copy(result, boxes)
|
||||
pageCols := make(map[int]int)
|
||||
determineBestKForPage(boxes, result, []int{0, 1, 2, 3}, 0, pageCols)
|
||||
if pageCols[0] != 2 {
|
||||
t.Errorf("two distinct columns: expected pageCols[0]=2, got %d", pageCols[0])
|
||||
}
|
||||
}
|
||||
|
||||
func TestAssignmentHelpers_IndentTolerance(t *testing.T) {
|
||||
boxes := []pdf.TextBox{
|
||||
{PageNumber: 0, X0: 50, X1: 150, Top: 10, Bottom: 30},
|
||||
{PageNumber: 0, X0: 205, X1: 350, Top: 10, Bottom: 30},
|
||||
{PageNumber: 0, X0: 58, X1: 150, Top: 40, Bottom: 60},
|
||||
}
|
||||
result := make([]pdf.TextBox, len(boxes))
|
||||
copy(result, boxes)
|
||||
pageCols := make(map[int]int)
|
||||
determineBestKForPage(boxes, result, []int{0, 1, 2}, 0, pageCols)
|
||||
if pageCols[0] != 2 {
|
||||
t.Errorf("expected 2 columns after indent tolerance, got %d", pageCols[0])
|
||||
}
|
||||
}
|
||||
|
||||
func TestAssignColIDsForPage_Normal(t *testing.T) {
|
||||
boxes := []pdf.TextBox{
|
||||
{PageNumber: 0, X0: 50, X1: 100},
|
||||
{PageNumber: 0, X0: 200, X1: 300},
|
||||
}
|
||||
result := make([]pdf.TextBox, len(boxes))
|
||||
copy(result, boxes)
|
||||
pageCols := map[int]int{0: 2}
|
||||
assignColIDsForPage(boxes, result, []int{0, 1}, 0, pageCols)
|
||||
if result[0].ColID != 0 || result[1].ColID != 1 {
|
||||
t.Errorf("expected ColIDs 0,1 but got %d,%d", result[0].ColID, result[1].ColID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAssignColIDsForPage_KTooLarge(t *testing.T) {
|
||||
boxes := []pdf.TextBox{
|
||||
{PageNumber: 0, X0: 100, X1: 200},
|
||||
}
|
||||
result := make([]pdf.TextBox, len(boxes))
|
||||
copy(result, boxes)
|
||||
pageCols := map[int]int{0: 3}
|
||||
assignColIDsForPage(boxes, result, []int{0}, 0, pageCols)
|
||||
if result[0].ColID != 0 {
|
||||
t.Errorf("expected ColID=0 (k clamped to 1), got %d", result[0].ColID)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2,7 +2,6 @@ package util
|
||||
|
||||
import (
|
||||
"math"
|
||||
"sort"
|
||||
)
|
||||
|
||||
// KMeans1D performs 1-dimensional KMeans clustering.
|
||||
@@ -89,86 +88,3 @@ func KMeans1D(data []float64, k int) (labels []int, centroids []float64) {
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
// Silhouette1D computes the silhouette score for 1D data.
|
||||
// Returns a score in [-1, 1]. Higher is better.
|
||||
// Returns -1 if the score cannot be computed (fewer than 2 unique labels).
|
||||
// Samples alone in their cluster contribute 0, matching sklearn behavior.
|
||||
//
|
||||
// Python: sklearn.metrics.silhouette_score with Euclidean distance.
|
||||
func Silhouette1D(data []float64, labels []int) float64 {
|
||||
n := len(data)
|
||||
if n <= 1 {
|
||||
return 0
|
||||
}
|
||||
|
||||
clusterCounts := make(map[int]int)
|
||||
for _, l := range labels {
|
||||
clusterCounts[l]++
|
||||
}
|
||||
|
||||
uniqueClusters := make([]int, 0, len(clusterCounts))
|
||||
for cl := range clusterCounts {
|
||||
uniqueClusters = append(uniqueClusters, cl)
|
||||
}
|
||||
|
||||
// Need at least 2 distinct labels for silhouette.
|
||||
if len(uniqueClusters) < 2 {
|
||||
return -1
|
||||
}
|
||||
sort.Ints(uniqueClusters)
|
||||
|
||||
var totalScore float64
|
||||
for i := 0; i < n; i++ {
|
||||
// sklearn convention: silhouette = 0 for samples alone in their cluster.
|
||||
if clusterCounts[labels[i]] <= 1 {
|
||||
continue
|
||||
}
|
||||
|
||||
// a_i: mean distance to other points in same cluster
|
||||
var aSum float64
|
||||
aCount := 0
|
||||
for j := 0; j < n; j++ {
|
||||
if i != j && labels[j] == labels[i] {
|
||||
aSum += math.Abs(data[i] - data[j])
|
||||
aCount++
|
||||
}
|
||||
}
|
||||
a := 0.0
|
||||
if aCount > 0 {
|
||||
a = aSum / float64(aCount)
|
||||
}
|
||||
|
||||
// b_i: min mean distance to points in other clusters
|
||||
b := math.MaxFloat64
|
||||
for _, cl := range uniqueClusters {
|
||||
if cl == labels[i] {
|
||||
continue
|
||||
}
|
||||
var bSum float64
|
||||
bCount := 0
|
||||
for j := 0; j < n; j++ {
|
||||
if labels[j] == cl {
|
||||
bSum += math.Abs(data[i] - data[j])
|
||||
bCount++
|
||||
}
|
||||
}
|
||||
if bCount > 0 {
|
||||
meanDist := bSum / float64(bCount)
|
||||
if meanDist < b {
|
||||
b = meanDist
|
||||
}
|
||||
}
|
||||
}
|
||||
if b == math.MaxFloat64 {
|
||||
b = 0
|
||||
}
|
||||
|
||||
maxAB := math.Max(a, b)
|
||||
if maxAB > 0 {
|
||||
totalScore += (b - a) / maxAB
|
||||
}
|
||||
}
|
||||
|
||||
return totalScore / float64(n)
|
||||
}
|
||||
|
||||
@@ -61,31 +61,3 @@ func TestKMeans1D(t *testing.T) {
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestSilhouette1D(t *testing.T) {
|
||||
t.Run("well-separated clusters", func(t *testing.T) {
|
||||
data := []float64{0, 1, 2, 100, 101, 102}
|
||||
labels := []int{0, 0, 0, 1, 1, 1}
|
||||
score := Silhouette1D(data, labels)
|
||||
if score < 0.8 {
|
||||
t.Errorf("well-separated score should be high, got %.3f", score)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("overlapping clusters", func(t *testing.T) {
|
||||
data := []float64{0, 1, 0, 1, 0, 1}
|
||||
labels := []int{0, 0, 0, 1, 1, 1}
|
||||
score := Silhouette1D(data, labels)
|
||||
if score > 0.5 {
|
||||
t.Errorf("overlapping score should be low, got %.3f", score)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("single cluster returns -1", func(t *testing.T) {
|
||||
data := []float64{1, 2, 3}
|
||||
labels := []int{0, 0, 0}
|
||||
if score := Silhouette1D(data, labels); score != -1 {
|
||||
t.Errorf("single cluster should return -1, got %.3f", score)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user