398 lines
14 KiB
Python
398 lines
14 KiB
Python
"""Rectangle-based table detection (pdf-inspector-inspired union-find).
|
||
|
||
MIT-licensed algorithm family (Firecrawl pdf-inspector detect_rects): cluster
|
||
cell-sized rectangles by spatial overlap, then build a grid when regular.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
from dataclasses import dataclass
|
||
|
||
from app.services.convert.idm.model import BBox, Block, BlockType, TextSpan
|
||
from app.services.convert.layout import table_plausibility
|
||
from app.services.convert.layout.glyphs import Line
|
||
|
||
|
||
@dataclass
|
||
class PageRect:
|
||
x: float
|
||
y: float
|
||
w: float
|
||
h: float
|
||
|
||
@property
|
||
def x1(self) -> float:
|
||
return self.x + self.w
|
||
|
||
@property
|
||
def y1(self) -> float:
|
||
return self.y + self.h
|
||
|
||
|
||
# Share of the lines inside a rectangle cluster that must land in a cell
|
||
# before the grid is believed. Below this the rectangles describe something
|
||
# other than a table — a form's boxes, a chart's plot area — and the
|
||
# ruling/gap heuristics downstream do better.
|
||
MIN_PLACEMENT = 0.75
|
||
|
||
|
||
class _UnionFind:
|
||
def __init__(self, n: int) -> None:
|
||
self.p = list(range(n))
|
||
self.sz = [1] * n
|
||
|
||
def find(self, i: int) -> int:
|
||
while self.p[i] != i:
|
||
self.p[i] = self.p[self.p[i]]
|
||
i = self.p[i]
|
||
return i
|
||
|
||
def union(self, a: int, b: int) -> None:
|
||
ra, rb = self.find(a), self.find(b)
|
||
if ra == rb:
|
||
return
|
||
if self.sz[ra] < self.sz[rb]:
|
||
ra, rb = rb, ra
|
||
self.p[rb] = ra
|
||
self.sz[ra] += self.sz[rb]
|
||
|
||
|
||
def _overlap(a: PageRect, b: PageRect, tol: float) -> bool:
|
||
return not (
|
||
a.x1 + tol < b.x
|
||
or b.x1 + tol < a.x
|
||
or a.y1 + tol < b.y
|
||
or b.y1 + tol < a.y
|
||
)
|
||
|
||
|
||
def _cluster_rects(rects: list[PageRect], *, tol: float = 2.0) -> list[list[int]]:
|
||
n = len(rects)
|
||
if n == 0:
|
||
return []
|
||
uf = _UnionFind(n)
|
||
# Bucket by coarse grid to avoid O(n^2) on huge pages
|
||
for i in range(n):
|
||
for j in range(i + 1, n):
|
||
if _overlap(rects[i], rects[j], tol):
|
||
uf.union(i, j)
|
||
if uf.sz[uf.find(i)] > 800:
|
||
break
|
||
groups: dict[int, list[int]] = {}
|
||
for i in range(n):
|
||
r = uf.find(i)
|
||
groups.setdefault(r, []).append(i)
|
||
return [g for g in groups.values() if len(g) >= 4]
|
||
|
||
|
||
def _merge_close(values: list[float], gap: float) -> list[float]:
|
||
"""Collapse coordinates within *gap* of each other into one boundary."""
|
||
if not values:
|
||
return []
|
||
out = [values[0]]
|
||
for v in values[1:]:
|
||
if abs(v - out[-1]) <= gap:
|
||
out[-1] = (out[-1] + v) / 2.0
|
||
else:
|
||
out.append(v)
|
||
return out
|
||
|
||
|
||
def _bucket(edges: list[float], value: float) -> int | None:
|
||
"""Index of the band ``[edges[i], edges[i+1])`` holding *value*."""
|
||
if len(edges) < 2 or value < edges[0] or value > edges[-1]:
|
||
return None
|
||
lo, hi = 0, len(edges) - 2
|
||
while lo <= hi:
|
||
mid = (lo + hi) // 2
|
||
if value < edges[mid]:
|
||
hi = mid - 1
|
||
elif value >= edges[mid + 1]:
|
||
lo = mid + 1
|
||
else:
|
||
return mid
|
||
return max(0, min(len(edges) - 2, lo))
|
||
|
||
|
||
def _grid_from_cluster(
|
||
rects: list[PageRect],
|
||
idxs: list[int],
|
||
lines: list[Line],
|
||
) -> tuple[list[list[str]], float, set[int]] | None:
|
||
"""Build a cell grid from a rectangle cluster and fill it from *lines*.
|
||
|
||
The grid is defined by the rectangles' **edges**, not their centres. Edges
|
||
partition the table area into bands, so every point inside falls in exactly
|
||
one cell and no tolerance has to be guessed; matching on nearest centre
|
||
with a fixed ±14pt/±40pt window silently discarded any line that sat
|
||
between two centres, which is how whole data rows went missing.
|
||
|
||
Text is placed span by span rather than line by line. A row of a ruled
|
||
table often reaches the layout as one line — "Low Vision 5 2 3" — because
|
||
the gaps between its cells are narrower than the column-gutter threshold;
|
||
assigning that whole line to the column under its midpoint puts five
|
||
values in one cell and leaves four empty.
|
||
|
||
Returns the grid, a confidence, and the indices of the lines it consumed,
|
||
so the caller can return everything else to the page instead of dropping
|
||
it inside the table's bounding box.
|
||
"""
|
||
cluster = [rects[i] for i in idxs]
|
||
# Filter page-background giants
|
||
areas = [max(r.w * r.h, 1.0) for r in cluster]
|
||
med = sorted(areas)[len(areas) // 2]
|
||
cells = [r for r in cluster if r.w * r.h <= med * 8 and r.w > 4 and r.h > 4]
|
||
if len(cells) < 4:
|
||
return None
|
||
|
||
col_edges = _merge_close(sorted({round(v, 1) for c in cells for v in (c.x, c.x1)}), 6.0)
|
||
row_edges = _merge_close(sorted({round(v, 1) for c in cells for v in (c.y, c.y1)}), 4.0)
|
||
return fill_grid(lines, col_edges, row_edges)
|
||
|
||
|
||
def _merge_spans_text(spans: list[TextSpan]) -> str:
|
||
"""Merge glyph/word spans into clean cell text, avoiding spurious spaces.
|
||
|
||
PDF extractors often split runs across kerning pairs, decimal points, and punctuation
|
||
(e.g. ['Catego', 'ry'], ['34', '.', '5%'], ['Co', 'm', 'pleted']). Naive joining
|
||
with spaces corrupts numbers and words.
|
||
"""
|
||
if not spans:
|
||
return ""
|
||
if len(spans) == 1:
|
||
return (spans[0].text or "").strip()
|
||
|
||
NO_LEADING_SPACE = {".", ",", ":", ";", "!", "?", "%", ")", "]", "}", "/", "-", "–", "—", "="}
|
||
NO_TRAILING_SPACE = {"(", "[", "{", "$", "£", "€", "/", "-", "–", "—", "="}
|
||
|
||
ordered = sorted(spans, key=lambda s: float(getattr(s, "x", 0.0) or 0.0))
|
||
result: list[str] = []
|
||
prev_span: TextSpan | None = None
|
||
|
||
for s in ordered:
|
||
raw = s.text or ""
|
||
txt = raw.strip()
|
||
if not txt:
|
||
if raw and result and not result[-1].endswith(" "):
|
||
result.append(" ")
|
||
continue
|
||
|
||
if not result:
|
||
result.append(txt)
|
||
prev_span = s
|
||
continue
|
||
|
||
cur_x = float(getattr(s, "x", 0.0) or 0.0)
|
||
prev_x = float(getattr(prev_span, "x", 0.0) or 0.0) if prev_span else 0.0
|
||
prev_w = float(getattr(prev_span, "w", 0.0) or 0.0) if prev_span else 0.0
|
||
prev_end = prev_x + prev_w
|
||
gap = cur_x - prev_end
|
||
prev_text = prev_span.text or "" if prev_span else ""
|
||
|
||
needs_space = False
|
||
if prev_text.endswith(" ") or raw.startswith(" "):
|
||
# Some walkers include the *next* text-matrix advance in a span,
|
||
# e.g. ``"= "``. If the next glyph starts before that span's
|
||
# reported right edge, the trailing blank is metadata rather than
|
||
# a visual separator (``n= 1`` must remain ``n=1``). A real word
|
||
# space starts at, or just after, the previous span's edge.
|
||
needs_space = cur_x >= prev_end - 0.5
|
||
elif txt in NO_LEADING_SPACE or (len(txt) > 0 and txt[0] in NO_LEADING_SPACE):
|
||
needs_space = False
|
||
elif prev_text.strip() and prev_text.strip()[-1] in NO_TRAILING_SPACE:
|
||
needs_space = False
|
||
elif prev_span and cur_x > 0 and prev_end > 0:
|
||
fsize = float(getattr(s, "font_size", 0.0) or getattr(prev_span, "font_size", 0.0) or 10.0)
|
||
space_thresh = max(1.8, fsize * 0.22)
|
||
if gap >= space_thresh:
|
||
needs_space = True
|
||
else:
|
||
needs_space = False
|
||
else:
|
||
needs_space = True
|
||
|
||
if needs_space and not result[-1].endswith(" "):
|
||
result.append(" ")
|
||
result.append(txt)
|
||
prev_span = s
|
||
|
||
return "".join(result).strip()
|
||
|
||
|
||
def fill_grid(
|
||
lines: list[Line],
|
||
col_edges: list[float],
|
||
row_edges: list[float],
|
||
) -> tuple[list[list[str]], float, set[int]] | None:
|
||
"""Place *lines* into the cells defined by ascending edge coordinates.
|
||
|
||
Shared by the rectangle and ruling detectors: once a grid's boundaries are
|
||
known, filling it is the same problem however the boundaries were found.
|
||
"""
|
||
n_cols, n_rows = len(col_edges) - 1, len(row_edges) - 1
|
||
if n_cols < 2 or n_rows < 2 or n_rows * n_cols > 2000:
|
||
return None
|
||
|
||
# row_edges ascend in PDF space (y up); rows read top-down.
|
||
grid: list[list[list[str]]] = [[[] for _ in range(n_cols)] for _ in range(n_rows)]
|
||
consumed: set[int] = set()
|
||
enclosed = 0
|
||
placed = 0
|
||
|
||
for idx, ln in enumerate(lines):
|
||
if not (ln.text or "").strip():
|
||
continue
|
||
mx = (ln.x0 + ln.x1) / 2.0
|
||
if not (col_edges[0] - 2 <= mx <= col_edges[-1] + 2):
|
||
continue
|
||
ri = _bucket(row_edges, ln.y)
|
||
if ri is None:
|
||
continue
|
||
enclosed += 1
|
||
row = n_rows - 1 - ri
|
||
# Distribute the line's spans across the columns they actually sit in.
|
||
parts: dict[int, list[TextSpan]] = {}
|
||
for span in ln.spans or []:
|
||
if not (span.text or "").strip():
|
||
continue
|
||
sx = float(span.x or ln.x0)
|
||
centre = sx + float(span.w or 0) / 2.0
|
||
ci = _bucket(col_edges, centre)
|
||
if ci is None:
|
||
ci = 0 if centre < col_edges[0] else n_cols - 1
|
||
parts.setdefault(ci, []).append(span)
|
||
if len(parts) <= 1:
|
||
ci = next(iter(parts.keys())) if parts else _bucket(col_edges, mx)
|
||
if ci is not None:
|
||
if parts and ci in parts:
|
||
t = _merge_spans_text(parts[ci])
|
||
else:
|
||
t = (ln.text or "").strip()
|
||
if t:
|
||
grid[row][ci].append(t)
|
||
else:
|
||
for ci, chunk in parts.items():
|
||
t = _merge_spans_text(chunk)
|
||
if t:
|
||
grid[row][ci].append(t)
|
||
consumed.add(idx)
|
||
placed += 1
|
||
|
||
if enclosed and placed / enclosed < MIN_PLACEMENT:
|
||
return None
|
||
|
||
flat = [[" ".join(cell).strip() for cell in row] for row in grid]
|
||
# Drop bands that carry nothing: edge sets include hairline separators,
|
||
# which would otherwise produce a blank row or column between every pair.
|
||
keep_rows = [i for i, row in enumerate(flat) if any(c for c in row)]
|
||
keep_cols = [j for j in range(n_cols) if any(flat[i][j] for i in range(n_rows))]
|
||
if len(keep_rows) < 2 or len(keep_cols) < 2:
|
||
return None
|
||
flat = [[flat[i][j] for j in keep_cols] for i in keep_rows]
|
||
|
||
filled = sum(1 for row in flat for c in row if c)
|
||
if filled < 4 or placed < 3:
|
||
return None
|
||
# Confidence is what the grid's own content says about it, not a count of
|
||
# non-empty cells: eight filled cells used to score 0.95 whether they held
|
||
# a price list or the labels of a figure.
|
||
verdict = table_plausibility.assess(flat)
|
||
if not verdict:
|
||
return None
|
||
return flat, verdict.score, consumed
|
||
|
||
|
||
def detect_tables_from_rects(
|
||
lines: list[Line],
|
||
page_rects: list[PageRect] | list[dict] | None,
|
||
*,
|
||
start_order: int = 0,
|
||
min_conf: float = 0.60,
|
||
) -> tuple[list[Block], float, set[int]]:
|
||
"""Table blocks from rectangle clusters, plus the line indices consumed.
|
||
|
||
The third value is what lets the caller keep every line the grid did not
|
||
take. Deciding "inside the table" from the block's bounding box instead
|
||
deletes a caption or a footnote that happens to sit between two rows.
|
||
"""
|
||
if not page_rects or not lines:
|
||
return [], 0.0, set()
|
||
rects: list[PageRect] = []
|
||
for r in page_rects:
|
||
if isinstance(r, PageRect):
|
||
rects.append(r)
|
||
else:
|
||
rects.append(
|
||
PageRect(
|
||
x=float(r.get("x", 0)),
|
||
y=float(r.get("y", 0)),
|
||
w=float(r.get("w", r.get("width", 0))),
|
||
h=float(r.get("h", r.get("height", 0))),
|
||
)
|
||
)
|
||
# Keep cell-ish sizes
|
||
rects = [r for r in rects if 8 < r.w < 400 and 8 < r.h < 200]
|
||
if len(rects) < 6:
|
||
return [], 0.0, set()
|
||
|
||
best_grid = None
|
||
best_conf = 0.0
|
||
best_used: set[int] = set()
|
||
for group in _cluster_rects(rects):
|
||
got = _grid_from_cluster(rects, group, lines)
|
||
if got and got[1] > best_conf:
|
||
best_grid, best_conf, best_used = got
|
||
|
||
if not best_grid or best_conf < min_conf:
|
||
return [], best_conf, set()
|
||
|
||
used = [lines[i] for i in sorted(best_used)] or lines
|
||
ys = [ln.y for ln in used]
|
||
xs0 = [ln.x0 for ln in used]
|
||
xs1 = [ln.x1 for ln in used]
|
||
block = Block(
|
||
type=BlockType.table,
|
||
cells=best_grid,
|
||
table_confidence=best_conf,
|
||
bbox=BBox(
|
||
x=min(xs0) if xs0 else 0,
|
||
y=min(ys) if ys else 0,
|
||
w=(max(xs1) - min(xs0)) if xs0 and xs1 else 0,
|
||
h=(max(ys) - min(ys) + 12) if ys else 0,
|
||
),
|
||
reading_order=start_order,
|
||
text="",
|
||
)
|
||
return [block], best_conf, best_used
|
||
|
||
|
||
def extract_rects_from_display_list(ops: list[dict] | None) -> list[PageRect]:
|
||
"""Pull rectangle operators from C++/engine display list ops."""
|
||
if not ops:
|
||
return []
|
||
out: list[PageRect] = []
|
||
for op in ops:
|
||
typ = str(op.get("type") or op.get("op") or "").lower()
|
||
if typ not in ("rect", "re", "rectangle", "fill_rect", "stroke_rect"):
|
||
continue
|
||
if "args" in op and isinstance(op["args"], (list, tuple)) and len(op["args"]) >= 4:
|
||
try:
|
||
x = float(op["args"][0])
|
||
y = float(op["args"][1])
|
||
w = float(op["args"][2])
|
||
h = float(op["args"][3])
|
||
except (ValueError, TypeError):
|
||
x = y = w = h = 0.0
|
||
else:
|
||
x = float(op.get("x", op.get("x0", 0)))
|
||
y = float(op.get("y", op.get("y0", 0)))
|
||
w = float(op.get("w", op.get("width", 0)))
|
||
h = float(op.get("h", op.get("height", 0)))
|
||
if w <= 0 and "x1" in op:
|
||
w = float(op["x1"]) - x
|
||
if h <= 0 and "y1" in op:
|
||
h = abs(float(op["y1"]) - y)
|
||
if w > 2 and abs(h) > 2:
|
||
out.append(PageRect(x=x, y=y, w=w, h=abs(h)))
|
||
return out
|