From b4dd0f7a0c5293b681ff7425f7ff6b7028b8334a Mon Sep 17 00:00:00 2001 From: Jack Date: Tue, 11 Aug 2026 22:13:42 +0800 Subject: [PATCH] fix(pdf/layout): replace KMeans-only column detection with gap + KMeans hybrid (#18023) --- .../parser/pdf/layout/combined_column.go | 474 ++++++++++++++++++ .../layout/combined_column_synthetic_test.go | 109 ++++ .../parser/pdf/layout/combined_column_test.go | 283 +++++++++++ internal/deepdoc/parser/pdf/layout/layout.go | 142 +----- .../deepdoc/parser/pdf/layout/layout_test.go | 193 +------ internal/deepdoc/parser/pdf/util/kmeans.go | 84 ---- .../deepdoc/parser/pdf/util/kmeans_test.go | 28 -- 7 files changed, 874 insertions(+), 439 deletions(-) create mode 100644 internal/deepdoc/parser/pdf/layout/combined_column.go create mode 100644 internal/deepdoc/parser/pdf/layout/combined_column_synthetic_test.go create mode 100644 internal/deepdoc/parser/pdf/layout/combined_column_test.go diff --git a/internal/deepdoc/parser/pdf/layout/combined_column.go b/internal/deepdoc/parser/pdf/layout/combined_column.go new file mode 100644 index 0000000000..fe1b0fcc3c --- /dev/null +++ b/internal/deepdoc/parser/pdf/layout/combined_column.go @@ -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 +} diff --git a/internal/deepdoc/parser/pdf/layout/combined_column_synthetic_test.go b/internal/deepdoc/parser/pdf/layout/combined_column_synthetic_test.go new file mode 100644 index 0000000000..a81a0c1254 --- /dev/null +++ b/internal/deepdoc/parser/pdf/layout/combined_column_synthetic_test.go @@ -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) + } +} diff --git a/internal/deepdoc/parser/pdf/layout/combined_column_test.go b/internal/deepdoc/parser/pdf/layout/combined_column_test.go new file mode 100644 index 0000000000..5f17408f3b --- /dev/null +++ b/internal/deepdoc/parser/pdf/layout/combined_column_test.go @@ -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) + } +} diff --git a/internal/deepdoc/parser/pdf/layout/layout.go b/internal/deepdoc/parser/pdf/layout/layout.go index d898763632..387247e570 100644 --- a/internal/deepdoc/parser/pdf/layout/layout.go +++ b/internal/deepdoc/parser/pdf/layout/layout.go @@ -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) ---- diff --git a/internal/deepdoc/parser/pdf/layout/layout_test.go b/internal/deepdoc/parser/pdf/layout/layout_test.go index 4626a031a1..e9e7b16f5f 100644 --- a/internal/deepdoc/parser/pdf/layout/layout_test.go +++ b/internal/deepdoc/parser/pdf/layout/layout_test.go @@ -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) - } -} diff --git a/internal/deepdoc/parser/pdf/util/kmeans.go b/internal/deepdoc/parser/pdf/util/kmeans.go index 8bdd37a5f9..4f40edb6ea 100644 --- a/internal/deepdoc/parser/pdf/util/kmeans.go +++ b/internal/deepdoc/parser/pdf/util/kmeans.go @@ -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) -} diff --git a/internal/deepdoc/parser/pdf/util/kmeans_test.go b/internal/deepdoc/parser/pdf/util/kmeans_test.go index 11725c0304..cab126be06 100644 --- a/internal/deepdoc/parser/pdf/util/kmeans_test.go +++ b/internal/deepdoc/parser/pdf/util/kmeans_test.go @@ -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) - } - }) -}