Coverage for src / quber / core / extractors / set_of_mark / split.py: 79%
112 statements
« prev ^ index » next coverage.py v7.14.0, created at 2026-09-23 22:14 -0400
1"""Split a Set-of-Mark region that fused two or more stacked tables into one.
3The Set-of-Mark locator occasionally marks two vertically-stacked,
4independently-titled tables as a single region. The fused region extracts with
5correct values but a broken structure: the two tables' period labels, headers
6and footnotes get crossed or lost. `split_table` is a post-pass over one
7already-captured table that refines such a region back into the separate tables
8it should have been.
10It is self-contained: same type in, same type out. The expected table count and
11the per-table boundaries are handed in by the caller (their provenance is a
12parsed document's table geometry); this function never parses the document
13itself. It reaches the source PDF only through `table.source`, and only to
14re-extract values within a tighter box.
16Two independent count signals decide whether to split: the handed-in `count`
17and a vision count probe run here over the region image. They must agree on more
18than one table; on agreement of one, on any disagreement, or on any failure of
19the re-extraction, the region is returned exactly as it came in. The cut between
20sub-regions is deterministic geometry taken from the supplied boundaries, never
21an estimated coordinate, and the sub-regions always tile the full original
22region so a footnote sitting below a table's tight box is never dropped.
23"""
25from __future__ import annotations
27import asyncio
28import tempfile
29from pathlib import Path
30from typing import List, Optional, Sequence, Tuple
32import fitz
33from loguru import logger
35from quber.agents.llm_client import LLMClient
36from quber.core.extractors.base import ExtractedTable, MergedCellBox, grid_fingerprint, grounded_grid
37from quber.core.extractors.camelot.acquire import grid_to_markdown
38from quber.core.extractors.camelot.correspondence.correction import correct_structure
39from quber.core.extractors.camelot.correspondence.geometry import (
40 crop_region_png,
41 norm_bbox_to_table_area,
42)
43from quber.core.extractors.camelot.correspondence.recovery import camelot_targeted
44from quber.core.extractors.set_of_mark.merge_grounding import (
45 find_dropped_header_text,
46 locate_markers,
47 log_ungrounded,
48 log_ungrounded_cells,
49 resolve_corrected_grid,
50 resolve_merges,
51)
53# A box normalized to 0..1 with the page top-left as origin (x1, y1, x2, y2);
54# the same frame as ExtractedTable.som_region. Boundaries handed in must share
55# this frame so the cut geometry lines up with the region being subdivided.
56NormBox = Tuple[float, float, float, float]
59async def split_table(
60 table: ExtractedTable,
61 count: int,
62 boundaries: Sequence[NormBox],
63 llm: LLMClient,
64 *,
65 dpi: int = 200,
66 correct_sem: Optional[asyncio.Semaphore] = None,
67) -> List[ExtractedTable]:
68 """Refine one captured table into the separate tables it fused, or pass it through.
70 `count` is the expected number of tables in the region and `boundaries` are
71 their per-table boxes (one per expected table, same normalized top-left frame
72 as `table.som_region`). The region is split into `count` tables only when the
73 vision probe agrees the region holds more than one table and every sub-region
74 re-extracts cleanly; otherwise the input table is returned as a one-element
75 list, unchanged.
76 """
77 region = table.som_region
78 # Nothing to refine unless more than one table is expected, every expected
79 # table has a boundary to cut on, and we have the region plus a source to
80 # re-read. Any of these missing means the region stands as captured.
81 if count <= 1 or region is None or table.source is None or len(boundaries) != count:
82 return [table]
84 sub_regions = subdivide(region, boundaries)
85 if sub_regions is None or len(sub_regions) != count:
86 return [table]
88 source = Path(table.source)
89 sem = correct_sem or asyncio.Semaphore(1)
91 # Rendering and every LLM/Camelot call below happen only on the split path;
92 # a region that is never split incurs none of this.
93 with tempfile.TemporaryDirectory(prefix="quber-split-") as tmp:
94 try:
95 page_image, page_w, page_h = await asyncio.to_thread(
96 render_page, source, table.page, dpi, Path(tmp)
97 )
98 except Exception as exc:
99 logger.error("page {}: split render failed: {}", table.page, exc)
100 return [table]
102 vision_count = await vision_count_probe(llm, page_image, region, page_w, page_h, dpi)
103 # The handed-in count and the vision probe must agree on the same N > 1.
104 # Anything else keeps the region as Camelot already extracted it.
105 if vision_count != count:
106 return [table]
108 out: List[ExtractedTable] = []
109 for ordinal, sub in enumerate(sub_regions, start=1):
110 captured = await reextract_subregion(
111 sub, source, table.page, ordinal, llm, page_image, page_w, page_h, dpi, sem
112 )
113 if captured is None:
114 # A sub-region with no recoverable grid would drop content if
115 # emitted, so the whole split is abandoned and the fused region
116 # is kept intact rather than shipped partial.
117 logger.warning(
118 "page {}: sub-region {} of {} produced no grid; keeping fused region",
119 table.page,
120 ordinal,
121 count,
122 )
123 return [table]
124 # A sub-table extends its parent's address with its own sub-index.
125 if table.table_id:
126 captured.table_id = f"{table.table_id}-s{ordinal}"
127 out.append(captured)
129 # A sub-table's crop pads only a little above its own grid, so the unit
130 # caption sitting above the whole group is seen by the top sub-table and
131 # missed by the ones below it. Carry the last caption seen downward: each
132 # sub-table with no caption of its own inherits the one above it, and a
133 # sub-table that captured its own caption keeps it and carries that one to
134 # the sub-tables below. The fill only flows down, so a sub-table above the
135 # first caption stays empty.
136 carried = ""
137 for sub in out:
138 if sub.units:
139 carried = sub.units
140 elif carried:
141 sub.units = carried
142 return out
145def subdivide(region: NormBox, boundaries: Sequence[NormBox]) -> Optional[List[NormBox]]:
146 """Tile `region` into vertical bands, one per boundary, cut in the gaps between them.
148 Boundaries are ordered top-to-bottom by their vertical center; the cut for
149 each adjacent pair sits at the midpoint of the blank gap between them. Every
150 band spans the region's full width, and the bands together cover the region
151 from its top edge to its bottom edge exactly — so content between a tight
152 boundary box and the next cut (a table's footnote) stays inside its band
153 rather than being dropped. Returns None if a boundary center falls outside
154 the region or the cuts would not increase top-to-bottom.
155 """
156 rx1, ry1, rx2, ry2 = region
157 left, right = min(rx1, rx2), max(rx1, rx2)
158 top, bottom = min(ry1, ry2), max(ry1, ry2)
160 def center_y(box: NormBox) -> float:
161 return (min(box[1], box[3]) + max(box[1], box[3])) / 2.0
163 boxes = sorted(boundaries, key=center_y)
164 for box in boxes:
165 if not (top <= center_y(box) <= bottom):
166 return None
168 seams: List[float] = []
169 for upper, lower in zip(boxes, boxes[1:], strict=False):
170 upper_bottom = max(upper[1], upper[3])
171 lower_top = min(lower[1], lower[3])
172 seams.append((upper_bottom + lower_top) / 2.0)
174 edges = [top, *seams, bottom]
175 if any(edges[i] >= edges[i + 1] for i in range(len(edges) - 1)):
176 return None
177 return [(left, edges[i], right, edges[i + 1]) for i in range(len(edges) - 1)]
180async def vision_count_probe(
181 llm: LLMClient,
182 page_image: Path,
183 region: NormBox,
184 page_w: float,
185 page_h: float,
186 dpi: int,
187) -> int:
188 """Ask vision how many independent tables are viewable inside the region.
190 Crops the rendered page to the region and counts tables in just that crop,
191 so a multi-table page does not inflate the count. Returns 0 on any failure,
192 which reads as disagreement and keeps the region unsplit.
193 """
194 region_pts = (region[0] * page_w, region[1] * page_h, region[2] * page_w, region[3] * page_h)
195 try:
196 crop_png = await asyncio.to_thread(crop_region_png, page_image, region_pts, dpi)
197 crop_path = page_image.parent / "split-probe.png"
198 crop_path.write_bytes(crop_png)
199 return await llm.count_tables(crop_path)
200 except Exception as exc:
201 logger.error("split vision count probe failed: {}", exc)
202 return 0
205async def reextract_subregion(
206 sub_region: NormBox,
207 source: Path,
208 page: int,
209 ordinal: int,
210 llm: LLMClient,
211 page_image: Path,
212 page_w: float,
213 page_h: float,
214 dpi: int,
215 correct_sem: asyncio.Semaphore,
216) -> Optional[ExtractedTable]:
217 """Re-read one sub-region from the source PDF and clean its structure.
219 Region-constrained Camelot reads the source within the sub-region box (the
220 only source of values), then the same structure-correction step the pipeline
221 runs per table fixes the header layout and recovers the title, caption and
222 footnotes from the page image. Returns None if Camelot finds no grid in the
223 box, so the caller can abandon the split rather than emit an empty table.
224 """
225 area = norm_bbox_to_table_area(sub_region, page_w, page_h)
226 try:
227 cand = await asyncio.to_thread(camelot_targeted, str(source), page, area, ordinal)
228 except Exception as exc:
229 logger.error("page {}: sub-region {} Camelot re-extract raised: {}", page, ordinal, exc)
230 return None
231 if cand is None:
232 return None
234 correction = await correct_structure(
235 cand.cells, page_image, cand.bbox, str(source), page, llm, correct_sem, dpi
236 )
237 cell_grid = grounded_grid(cand.cells, cand.cell_boxes, page_w, page_h)
238 merged: List[MergedCellBox] = []
239 dropped_text: List[str] = []
240 corrected_grid = cell_grid
241 if correction is not None:
242 if correction.cell_merges:
243 merged = resolve_merges(correction.cell_merges, cell_grid, correction.markdown)
244 log_ungrounded(merged, page)
245 corrected_grid = resolve_corrected_grid(
246 correction.markdown, merged, cell_grid, correction.footnote_refs, correction.region_text
247 )
248 log_ungrounded_cells(corrected_grid, page)
249 dropped_text = find_dropped_header_text(
250 corrected_grid,
251 cell_grid,
252 " ".join(
253 [correction.title, correction.caption, correction.units]
254 + [f"{f.marker} {f.text}".strip() for f in correction.footnotes]
255 ),
256 page,
257 )
258 if correction is not None:
259 return ExtractedTable(
260 content_fingerprint=grid_fingerprint(cand.cells),
261 title=correction.title,
262 caption=correction.caption,
263 markdown=correction.markdown,
264 footnotes=correction.footnotes,
265 footnote_refs=correction.footnote_refs,
266 footnote_marks=locate_markers(
267 correction.footnote_marks,
268 correction.footnote_refs,
269 correction.markdown,
270 correction.footnotes,
271 f"{correction.title} {correction.caption}",
272 ),
273 page=page,
274 bbox=cand.bbox,
275 flavor=cand.flavor,
276 source=str(source),
277 camelot_accuracy=cand.accuracy,
278 llm_corrected=correction.llm_corrected,
279 som_region=sub_region,
280 units=correction.units,
281 header_rows=correction.header_rows,
282 cell_grid=cell_grid,
283 corrected_grid=corrected_grid,
284 merged_cells=merged,
285 dropped_text=dropped_text,
286 )
287 return ExtractedTable(
288 content_fingerprint=grid_fingerprint(cand.cells),
289 markdown=grid_to_markdown(cand.cells),
290 page=page,
291 bbox=cand.bbox,
292 flavor=cand.flavor,
293 source=str(source),
294 camelot_accuracy=cand.accuracy,
295 som_region=sub_region,
296 cell_grid=cell_grid,
297 corrected_grid=corrected_grid,
298 )
301def render_page(source: Path, page: int, dpi: int, out_dir: Path) -> Tuple[Path, float, float]:
302 """Render one PDF page to a PNG at `dpi`; return its path and page size in points."""
303 doc = fitz.open(str(source))
304 try:
305 pg = doc[page - 1]
306 width, height = pg.rect.width, pg.rect.height
307 pix = pg.get_pixmap(dpi=dpi)
308 path = out_dir / f"page-{page:04d}.png"
309 pix.save(str(path))
310 finally:
311 doc.close()
312 return path, width, height