Coverage for src / quber / core / extractors / camelot / tighten.py: 98%
95 statements
« prev ^ index » next coverage.py v7.14.0, created at 2026-09-23 22:14 -0400
1"""Tighten Camelot cell rectangles to the words that print inside them.
3Camelot reports each cell as structural grid geometry: the column's x-extent
4crossed with the row band. On dense filings that geometry hugs the print, but
5on sparse layouts the separators land mid-whitespace, so a cell's rectangle
6runs far wider than its value and can sit half a line off the glyphs — a
7highlight drawn from it cuts through the number instead of surrounding it.
9The PDF itself records the rectangle of every word it prints (the text
10layer). A cell whose text can be found among the words near its structural
11rectangle takes those words' extent as its box instead. A cell whose text has
12no match there — an image-based table with no text layer, OCR-corrected text
13that diverges from the print — keeps the structural rectangle; no box is ever
14invented.
16Words arrive as `quber.agents.completeness.page_words` returns them:
17(x0, y0, x1, y1, text) in PDF points, top-left origin — the same text layer
18Camelot parses. `tighten_box` works in that frame. `tighten_cell_boxes`
19accepts and returns Camelot's frame — PDF points, bottom-left origin — so the
20acquisition layer stores exactly what Camelot's contract promises, only
21tighter.
22"""
24from __future__ import annotations
26from typing import List, Optional, Sequence, Tuple
28Rect = Tuple[float, float, float, float] # x1, y1, x2, y2
29Word = Tuple[float, float, float, float, str]
31# Slack in points when gathering candidate words: separators can land a hair
32# inside the print horizontally, and the row band can sit up to its own height
33# off the glyph line vertically (tighten_box expands by the band height).
34PAD_X = 4.0
35# Words whose vertical centers differ by less than this many points read as
36# one printed line when sorting into reading order.
37LINE_QUANT = 4.0
39# Typographic variants that make identical text read as different strings:
40# curly quotes against straight ones, the dash family against the hyphen.
41# Camelot and the correction agents emit ASCII where the print is set in
42# typographic glyphs, so both sides normalize before comparing.
43PUNCT_VARIANTS = str.maketrans(
44 {"‘": "'", "’": "'", "“": '"', "”": '"', "‐": "-", "‑": "-", "‒": "-", "–": "-", "—": "-", "−": "-"}
45)
48def squeeze(text: str) -> str:
49 """Comparison form: whitespace removed, typographic punctuation folded."""
50 return "".join(text.split()).translate(PUNCT_VARIANTS)
53def _union(run: Sequence[Word]) -> Rect:
54 return (
55 min(w[0] for w in run),
56 min(w[1] for w in run),
57 max(w[2] for w in run),
58 max(w[3] for w in run),
59 )
62def _candidates(box: Rect, words: Sequence[Word]) -> List[Tuple[Word, str]]:
63 """Words near the structural box, in reading order, with comparison text.
65 The margin is one box height vertically (camelot's row band drifts off
66 the glyph line) and one box height horizontally (a glued '$' can print in
67 the neighboring column's band).
68 """
69 x1, y1, x2, y2 = box
70 band = y2 - y1
71 margin_x = max(PAD_X, band)
72 cand = [
73 (w, t)
74 for w, t in ((w, squeeze(w[4])) for w in words)
75 if t
76 and x1 - margin_x <= (w[0] + w[2]) / 2 <= x2 + margin_x
77 and y1 - band <= (w[1] + w[3]) / 2 <= y2 + band
78 ]
79 cand.sort(key=lambda wt: (round((wt[0][1] + wt[0][3]) / 2 / LINE_QUANT), wt[0][0]))
80 return cand
83def _token_runs(cand: List[Tuple[Word, str]], token: str) -> List[Tuple[int, int]]:
84 """Every contiguous candidate run whose text equals `token`,
85 as (first, last) index pairs."""
86 out: List[Tuple[int, int]] = []
87 for i in range(len(cand)):
88 joined = ""
89 for j in range(i, len(cand)):
90 joined += cand[j][1]
91 if len(joined) >= len(token):
92 if joined == token:
93 out.append((i, j))
94 break
95 return out
98def tighten_box(box: Rect, text: str, words: Sequence[Word]) -> Optional[Rect]:
99 """The printed extent of `text` near the structural `box`, or None.
101 Two passes over the nearby words. First the whole cell text as one
102 contiguous run in reading order — the normal case; of several matching
103 runs the one nearest the box center wins. Failing that, each
104 whitespace-separated token matched as its own contiguous run, all tokens
105 in reading order — which grounds glued cells whose fragments print apart
106 ('$' left-aligned, its value right-aligned) and labels whose superscript
107 markers disturb the reading order of the joined string. No pass matching
108 means None; nothing is invented.
109 """
110 target = squeeze(text)
111 if not target:
112 return None
113 cand = _candidates(box, words)
115 runs: List[List[Word]] = []
116 for i in range(len(cand)):
117 joined = ""
118 for j in range(i, len(cand)):
119 joined += cand[j][1]
120 if len(joined) >= len(target):
121 if joined == target:
122 runs.append([wt[0] for wt in cand[i : j + 1]])
123 break
124 if runs:
125 x1, y1, x2, y2 = box
126 cx, cy = (x1 + x2) / 2, (y1 + y2) / 2
128 def offset(run: Sequence[Word]) -> float:
129 rx = (min(w[0] for w in run) + max(w[2] for w in run)) / 2
130 ry = (min(w[1] for w in run) + max(w[3] for w in run)) / 2
131 return (rx - cx) ** 2 + (ry - cy) ** 2
133 return _union(min(runs, key=offset))
135 tokens = [squeeze(t) for t in text.split()]
136 tokens = [t for t in tokens if t]
137 if len(tokens) < 2:
138 return None
139 # No ordering constraint between tokens: a superscript marker sorts off
140 # its label's line, and a glued '$' prints a column away. Each token takes
141 # its run nearest the box center; the candidate window bounds the damage
142 # a repeated token could do.
143 bx1, by1, bx2, by2 = box
144 cx, cy = (bx1 + bx2) / 2, (by1 + by2) / 2
145 picked: List[Word] = []
146 for token in tokens:
147 found = _token_runs(cand, token)
148 if not found:
149 return None
151 def run_offset(ij: Tuple[int, int]) -> float:
152 run = [wt[0] for wt in cand[ij[0] : ij[1] + 1]]
153 rx = (min(w[0] for w in run) + max(w[2] for w in run)) / 2
154 ry = (min(w[1] for w in run) + max(w[3] for w in run)) / 2
155 return (rx - cx) ** 2 + (ry - cy) ** 2
157 i, j = min(found, key=run_offset)
158 picked.extend(wt[0] for wt in cand[i : j + 1])
159 return _union(picked)
162def ink_extent(box: Rect, words: Sequence[Word]) -> Optional[Rect]:
163 """The extent of the words printed inside the structural `box`, or None.
165 The last resort for a cell whose text has no match in the text layer —
166 a synthesized header label, OCR-corrected text. The band cannot be
167 grounded to specific words, but it should never claim more of the page
168 than the ink it actually holds, so it shrinks to the words it contains.
169 A band holding no words (image tables) returns None and stays as it is.
170 """
171 x1, y1, x2, y2 = box
172 inside = [
173 w for w in words if x1 <= (w[0] + w[2]) / 2 <= x2 and y1 <= (w[1] + w[3]) / 2 <= y2 and w[4].strip()
174 ]
175 if not inside:
176 return None
177 ux1, uy1, ux2, uy2 = _union(inside)
178 # Never grow: the clamp intersects the band, it does not replace it.
179 return (max(ux1, x1), max(uy1, y1), min(ux2, x2), min(uy2, y2))
182def tighten_cell_boxes(
183 cells: Sequence[Sequence[str]],
184 cell_boxes: List[List[Optional[Rect]]],
185 words: Sequence[Word],
186 page_height: float,
187) -> List[List[Optional[Rect]]]:
188 """Cell boxes with every matchable cell shrunk to its printed words.
190 `cell_boxes` arrive and return in Camelot's frame (PDF points,
191 bottom-left origin), shaped like `cells`. A cell with no box, no text, or
192 no match keeps its entry unchanged.
193 """
194 out: List[List[Optional[Rect]]] = []
195 for i, row in enumerate(cell_boxes):
196 out_row: List[Optional[Rect]] = []
197 for j, box in enumerate(row):
198 text = cells[i][j] if i < len(cells) and j < len(cells[i]) else ""
199 if box is None or not text.strip():
200 out_row.append(box)
201 continue
202 bx1, by1, bx2, by2 = box
203 flipped = (bx1, page_height - by2, bx2, page_height - by1)
204 tight = tighten_box(flipped, text, words) or ink_extent(flipped, words)
205 if tight is None:
206 out_row.append(box)
207 else:
208 tx1, ty1, tx2, ty2 = tight
209 out_row.append((tx1, page_height - ty2, tx2, page_height - ty1))
210 out.append(out_row)
211 return out