mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-07 03:48:44 +08:00
372 lines
9.8 KiB
Go
372 lines
9.8 KiB
Go
package table
|
|
|
|
import (
|
|
"sort"
|
|
"strings"
|
|
|
|
pdf "ragflow/internal/deepdoc/parser/pdf/type"
|
|
)
|
|
|
|
// rb is a row-box entry that holds cell data during grid construction.
|
|
type rb struct {
|
|
row, col int
|
|
txt string
|
|
x0, y0, x1, y1 float64
|
|
label string
|
|
}
|
|
|
|
// GroupBoxesByRC groups text boxes into a cell grid by R/C annotations.
|
|
// Matches Python's construct_table: sort by R, sort by C within each row,
|
|
// merge nearby columns by X proximity.
|
|
func GroupBoxesByRC(boxes []pdf.TextBox) [][]pdf.TSRCell {
|
|
if len(boxes) == 0 {
|
|
return nil
|
|
}
|
|
// If no real R/C annotations (maxR <= 0), fall back to YX coordinate
|
|
// grouping — matching Python's construct_table when all R=-1.
|
|
maxR := 0
|
|
for _, b := range boxes {
|
|
if b.R > maxR {
|
|
maxR = b.R
|
|
}
|
|
}
|
|
if maxR <= 0 {
|
|
return GroupBoxesByYX(boxes)
|
|
}
|
|
// Sort by R index first (Python: sort_R_firstly), then Y, then X.
|
|
sort.Slice(boxes, func(i, j int) bool {
|
|
if boxes[i].R != boxes[j].R {
|
|
return boxes[i].R < boxes[j].R
|
|
}
|
|
if boxes[i].Top != boxes[j].Top {
|
|
return boxes[i].Top < boxes[j].Top
|
|
}
|
|
return boxes[i].X0 < boxes[j].X0
|
|
})
|
|
|
|
// Compress R indices: Python's sort_R_firstly grouping.
|
|
rowMap, compressed := compressRowIndices(boxes)
|
|
|
|
// Collect boxes per row.
|
|
cmap, _ := collectBoxesPerRow(boxes, rowMap)
|
|
|
|
// Compress C indices per row.
|
|
cCompressed, cMaxCol := compressColIndices(boxes, rowMap, compressed)
|
|
|
|
// Build grid.
|
|
return buildGrid(cmap, cCompressed, cMaxCol, compressed)
|
|
}
|
|
|
|
// GroupBoxesByYX groups boxes into a cell grid by Y/X coordinates,
|
|
// matching Python's construct_table which uses sort_R_firstly and
|
|
// sort_C_firstly when R/C annotations are absent. Falls back from
|
|
// GroupBoxesByRC when boxes lack R/C annotations.
|
|
func GroupBoxesByYX(boxes []pdf.TextBox) [][]pdf.TSRCell {
|
|
if len(boxes) == 0 {
|
|
return nil
|
|
}
|
|
// Sort by (page, top, x0) — same as Python sort_R_firstly with R=-1.
|
|
sort.Slice(boxes, func(i, j int) bool {
|
|
if boxes[i].PageNumber != boxes[j].PageNumber {
|
|
return boxes[i].PageNumber < boxes[j].PageNumber
|
|
}
|
|
if boxes[i].Top != boxes[j].Top {
|
|
return boxes[i].Top < boxes[j].Top
|
|
}
|
|
return boxes[i].X0 < boxes[j].X0
|
|
})
|
|
|
|
// Group into rows by Y proximity (Python's row grouping).
|
|
type rowGroup struct {
|
|
boxes []pdf.TextBox
|
|
top, btm float64
|
|
}
|
|
var rowGroups []rowGroup
|
|
rowGroups = append(rowGroups, rowGroup{
|
|
boxes: []pdf.TextBox{boxes[0]},
|
|
top: boxes[0].Top,
|
|
btm: boxes[0].Bottom,
|
|
})
|
|
for i := 1; i < len(boxes); i++ {
|
|
prev := &rowGroups[len(rowGroups)-1]
|
|
// Python: same row if top < prev.btm (Y overlaps) and same page.
|
|
if boxes[i].PageNumber == prev.boxes[0].PageNumber && boxes[i].Top < prev.btm {
|
|
prev.boxes = append(prev.boxes, boxes[i])
|
|
if boxes[i].Top < prev.top {
|
|
prev.top = boxes[i].Top
|
|
}
|
|
if boxes[i].Bottom > prev.btm {
|
|
prev.btm = boxes[i].Bottom
|
|
}
|
|
} else {
|
|
rowGroups = append(rowGroups, rowGroup{
|
|
boxes: []pdf.TextBox{boxes[i]},
|
|
top: boxes[i].Top,
|
|
btm: boxes[i].Bottom,
|
|
})
|
|
}
|
|
}
|
|
|
|
// Within each row, group into columns by X proximity.
|
|
rows := make([][]pdf.TSRCell, len(rowGroups))
|
|
for ri, rg := range rowGroups {
|
|
// Sort by X0.
|
|
sort.Slice(rg.boxes, func(i, j int) bool {
|
|
return rg.boxes[i].X0 < rg.boxes[j].X0
|
|
})
|
|
// Group by X overlap.
|
|
var cols []struct {
|
|
boxes []pdf.TextBox
|
|
x1 float64
|
|
}
|
|
cols = append(cols, struct {
|
|
boxes []pdf.TextBox
|
|
x1 float64
|
|
}{
|
|
boxes: []pdf.TextBox{rg.boxes[0]},
|
|
x1: rg.boxes[0].X1,
|
|
})
|
|
for i := 1; i < len(rg.boxes); i++ {
|
|
prev := &cols[len(cols)-1]
|
|
if rg.boxes[i].X0 < prev.x1 {
|
|
prev.boxes = append(prev.boxes, rg.boxes[i])
|
|
if rg.boxes[i].X1 > prev.x1 {
|
|
prev.x1 = rg.boxes[i].X1
|
|
}
|
|
} else {
|
|
cols = append(cols, struct {
|
|
boxes []pdf.TextBox
|
|
x1 float64
|
|
}{
|
|
boxes: []pdf.TextBox{rg.boxes[i]},
|
|
x1: rg.boxes[i].X1,
|
|
})
|
|
}
|
|
}
|
|
rows[ri] = make([]pdf.TSRCell, len(cols))
|
|
for ci, col := range cols {
|
|
var sb strings.Builder
|
|
for _, b := range col.boxes {
|
|
t := strings.TrimSpace(b.Text)
|
|
if t == "" {
|
|
continue
|
|
}
|
|
if sb.Len() > 0 {
|
|
sb.WriteByte(' ')
|
|
}
|
|
sb.WriteString(t)
|
|
}
|
|
rows[ri][ci].Text = sb.String()
|
|
}
|
|
}
|
|
return rows
|
|
}
|
|
|
|
// cellPosFromBox returns the position coordinates and label for a cell
|
|
// derived from a text box. Header cells use HLeft/HRight/HTop/HBott
|
|
// for spanning-aware positions; regular cells use the box's own bounds.
|
|
func cellPosFromBox(b pdf.TextBox) (x0, y0, x1, y1 float64, label string) {
|
|
x0, y0, x1, y1 = b.X0, b.Top, b.X1, b.Bottom
|
|
if b.H > 0 {
|
|
label = "table header"
|
|
if b.HLeft != 0 || b.HRight != 0 {
|
|
if b.HLeft != 0 {
|
|
x0 = b.HLeft
|
|
}
|
|
if b.HRight != 0 {
|
|
x1 = b.HRight
|
|
}
|
|
}
|
|
if b.HTop != 0 {
|
|
y0 = b.HTop
|
|
}
|
|
if b.HBott != 0 {
|
|
y1 = b.HBott
|
|
}
|
|
} else if b.SP > 0 {
|
|
label = "table spanning cell"
|
|
}
|
|
return
|
|
}
|
|
|
|
// cellLabelFromBox returns the TSR label for a box based on H/SP annotations.
|
|
// Used when merging multiple boxes into one cell — preserves the spanning label.
|
|
func cellLabelFromBox(b pdf.TextBox) string {
|
|
if b.H > 0 {
|
|
return "table header"
|
|
}
|
|
if b.SP > 0 {
|
|
return "table spanning cell"
|
|
}
|
|
return ""
|
|
}
|
|
|
|
// compressRowIndices compresses R values into contiguous row indices.
|
|
// Returns rowMap (original R → compressed index) and the maximum compressed index.
|
|
// Boxes must already be sorted by R, Y, X.
|
|
func compressRowIndices(boxes []pdf.TextBox) (map[int]int, int) {
|
|
rowMap := make(map[int]int) // original R → compressed row index
|
|
compressed := 0
|
|
rowMap[boxes[0].R] = 0
|
|
lastR := boxes[0].R
|
|
for i := 1; i < len(boxes); i++ {
|
|
if boxes[i].R != lastR {
|
|
compressed++
|
|
rowMap[boxes[i].R] = compressed
|
|
lastR = boxes[i].R
|
|
} else {
|
|
rowMap[boxes[i].R] = compressed
|
|
}
|
|
}
|
|
return rowMap, compressed
|
|
}
|
|
|
|
// collectBoxesPerRow collects boxes into row groups, merging boxes in the same cell.
|
|
// Returns cmap (row → col → entry) and maxCols (max column index per row).
|
|
func collectBoxesPerRow(boxes []pdf.TextBox, rowMap map[int]int) (map[int]map[int]*rb, map[int]int) {
|
|
cmap := make(map[int]map[int]*rb) // row → col → entry
|
|
maxCols := make(map[int]int)
|
|
for _, b := range boxes {
|
|
t := strings.TrimSpace(b.Text)
|
|
// Keep boxes with SP/H annotations even if text is empty —
|
|
// their coordinates are needed for colspan/rowspan calculation.
|
|
if t == "" && b.H <= 0 && b.SP <= 0 {
|
|
continue
|
|
}
|
|
r := rowMap[b.R]
|
|
c := b.C
|
|
if cmap[r] == nil {
|
|
cmap[r] = make(map[int]*rb)
|
|
}
|
|
x0, y0, x1, y1, label := cellPosFromBox(b)
|
|
if v, ok := cmap[r][c]; ok {
|
|
if t != "" {
|
|
v.txt += " " + t
|
|
}
|
|
// Merge spanning coordinates (use widest extent).
|
|
if b.H > 0 || b.SP > 0 {
|
|
v.label = cellLabelFromBox(b)
|
|
if v.x0 > x0 {
|
|
v.x0 = x0
|
|
}
|
|
if v.y0 > y0 {
|
|
v.y0 = y0
|
|
}
|
|
if v.x1 < x1 {
|
|
v.x1 = x1
|
|
}
|
|
if v.y1 < y1 {
|
|
v.y1 = y1
|
|
}
|
|
}
|
|
} else {
|
|
cmap[r][c] = &rb{r, c, t, x0, y0, x1, y1, label}
|
|
}
|
|
if c > maxCols[r] {
|
|
maxCols[r] = c
|
|
}
|
|
}
|
|
return cmap, maxCols
|
|
}
|
|
|
|
// rowBox is a helper for compressColIndices.
|
|
type rowBox struct {
|
|
c, idx int
|
|
x0, x1 float64
|
|
txt string
|
|
}
|
|
|
|
// compressColIndices compresses column indices per row based on X0 ordering and overlap.
|
|
// Returns cCompressed (row → original C → compressed C) and cMaxCol (max compressed C per row).
|
|
func compressColIndices(boxes []pdf.TextBox, rowMap map[int]int, compressed int) (map[int]map[int]int, map[int]int) {
|
|
cCompressed := make(map[int]map[int]int) // row → (original C → compressed col)
|
|
cMaxCol := make(map[int]int)
|
|
for ri := 0; ri <= compressed; ri++ {
|
|
// Collect all boxes in this row, sorted by X0.
|
|
var rowBoxes []rowBox
|
|
for i, b := range boxes {
|
|
if rowMap[b.R] == ri && (strings.TrimSpace(b.Text) != "" || b.H > 0 || b.SP > 0) {
|
|
rowBoxes = append(rowBoxes, rowBox{c: b.C, idx: i, x0: b.X0, x1: b.X1, txt: b.Text})
|
|
}
|
|
}
|
|
sort.Slice(rowBoxes, func(i, j int) bool { return rowBoxes[i].x0 < rowBoxes[j].x0 })
|
|
// Assign compressed column by X-order (disjoint X → new col).
|
|
cMap := make(map[int]int) // original C → compressed col
|
|
right := 0.0
|
|
nCols := 0
|
|
for _, rb := range rowBoxes {
|
|
if len(cMap) == 0 || rb.x0 >= right {
|
|
cMap[rb.c] = nCols
|
|
nCols++
|
|
right = rb.x1
|
|
} else {
|
|
// Overlapping X → merge into last column.
|
|
cMap[rb.c] = nCols - 1
|
|
if rb.x1 > right {
|
|
right = rb.x1
|
|
}
|
|
}
|
|
}
|
|
cCompressed[ri] = cMap
|
|
cMaxCol[ri] = nCols - 1
|
|
}
|
|
return cCompressed, cMaxCol
|
|
}
|
|
|
|
// buildGrid builds the final cell grid from the collected and compressed data.
|
|
func buildGrid(cmap map[int]map[int]*rb, cCompressed map[int]map[int]int, cMaxCol map[int]int, compressed int) [][]pdf.TSRCell {
|
|
rows := make([][]pdf.TSRCell, compressed+1)
|
|
for ri := 0; ri <= compressed; ri++ {
|
|
maxC := cMaxCol[ri]
|
|
rows[ri] = make([]pdf.TSRCell, maxC+1)
|
|
for ci, v := range cmap[ri] {
|
|
cci := cCompressed[ri][ci]
|
|
if cci <= maxC {
|
|
if rows[ri][cci].Text == "" {
|
|
rows[ri][cci].Text = v.txt
|
|
rows[ri][cci].X0 = v.x0
|
|
rows[ri][cci].Y0 = v.y0
|
|
rows[ri][cci].X1 = v.x1
|
|
rows[ri][cci].Y1 = v.y1
|
|
rows[ri][cci].Label = v.label
|
|
} else {
|
|
// Multiple originals map to same compressed cell — merge deterministically.
|
|
if v.txt != "" {
|
|
rows[ri][cci].Text += " " + v.txt
|
|
}
|
|
if v.x0 < rows[ri][cci].X0 {
|
|
rows[ri][cci].X0 = v.x0
|
|
}
|
|
if v.y0 < rows[ri][cci].Y0 {
|
|
rows[ri][cci].Y0 = v.y0
|
|
}
|
|
if v.x1 > rows[ri][cci].X1 {
|
|
rows[ri][cci].X1 = v.x1
|
|
}
|
|
if v.y1 > rows[ri][cci].Y1 {
|
|
rows[ri][cci].Y1 = v.y1
|
|
}
|
|
if rows[ri][cci].Label == "" && v.label != "" {
|
|
rows[ri][cci].Label = v.label
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
return rows
|
|
}
|
|
|
|
func BoxesHaveAnnotations(boxes []pdf.TextBox) bool {
|
|
maxR, maxC := 0, 0
|
|
for _, b := range boxes {
|
|
if b.R > maxR {
|
|
maxR = b.R
|
|
}
|
|
if b.C > maxC {
|
|
maxC = b.C
|
|
}
|
|
}
|
|
// True if at least 2 rows or 2 cols (R/C are 0-based, so maxR>0 means ≥2 rows).
|
|
return maxR > 0 || maxC > 0
|
|
}
|