Coverage for src / quber / core / figures / capture.py: 89%
139 statements
« prev ^ index » next coverage.py v7.14.0, created at 2026-09-23 22:14 -0400
1"""Turn a table read off a page image into an extracted table.
3A page is nominated partly for its tables: the parse records, for every table,
4what fraction of its cells it read off the page image rather than out of a text
5layer, and a table read off the image never reached the Set-of-Mark and Camelot
6engine either, because that engine extracts from a text layer and there was none.
7What stands in the document for such a table is the parse's own reading, with
8nothing having checked it since.
10The scan replaces that reading. It returns the table as a table — a grid with a
11box on every cell — so the grid goes straight into the same correction and
12grounding steps every other table in the document goes through, and what comes
13back is an `ExtractedTable` filled the way every other one is filled.
15Replacing rather than keeping both is deliberate. The parse's reading of these
16tables is not merely rougher, it is wrong in ways that change values: a currency
17symbol read as a letter, a thousands separator read as a decimal point. A wrong
18number carried beside a right one is how a wrong number survives.
20The parse's reading is not discarded on the way out, though. Before the grid
21reaches the correction step, every cell in it is checked against what the parse
22read in the same place, and any cell the two disagree on is read off the printed
23page by an agent. The scan has its own failure mode — it returns a thousands
24separator as a decimal point — and the parse is a second opinion on it that
25costs nothing.
27Only the tables that nominated the page are replaced. A page can print a table
28with a text layer beside one without, and the first was read properly by the
29table engine; overriding it with a scan would discard a vetted extraction for no
30reason. A nominated table the scan returned nothing over is reported.
31"""
33from __future__ import annotations
35import asyncio
36import tempfile
37from pathlib import Path
38from typing import Dict, List, Optional, Tuple
40from docling_core.types.doc.document import DocItem, DoclingDocument, PictureItem, TableItem
41from docling_core.types.doc.labels import DocItemLabel
42from loguru import logger
44from quber.agents.cell_reader import CellReader
45from quber.core.extractors.base import ExtractedTable
46from quber.core.extractors.camelot.correspondence.geometry import (
47 coverage_fraction,
48 crop_region_png,
49 table_crop_box,
50)
51from quber.core.extractors.set_of_mark.assemble import TableAssembly, assemble_table
52from quber.core.extractors.set_of_mark.extent import apply_content_regions
53from quber.core.figures.crosscheck import crosscheck_cells
54from quber.core.figures.geometry import NormBox, norm_box, point_box, prov_box
55from quber.core.figures.models import PageScan, ScannedTable
56from quber.core.figures.nominate import FURNITURE_CLASSES, picture_classes
57from quber.core.fusion.graft import READ_BY_FIELD
58from quber.files.pdf import render_page
60#: Minimum box overlap for a returned table and a parse table to be the same
61#: region, measured both ways so a table drawn wider on either side registers.
62#: The same threshold the figures are matched at.
63MATCH_FRACTION = 0.20
65#: Page rasterization for the correction and inspection agents. They read cell
66#: text off the image, which is what the table pipeline renders at.
67CAPTURE_DPI = 200
70async def capture_tables(
71 document: DoclingDocument,
72 scans: List[PageScan],
73 source: Path,
74 page_dims: Dict[int, Tuple[float, float]],
75 assembly: TableAssembly,
76 reader: Optional[CellReader] = None,
77 match_fraction: float = MATCH_FRACTION,
78 orphans: bool = False,
79) -> Tuple[List[ExtractedTable], List[str]]:
80 """Extract every nominated table the scans read, in page order.
82 Each record in `scans` is given the reference of the parse table its grid
83 replaced and the address of the table it produced, so a value traces back to
84 the element that holds it. Returns the tables and one error per nominated
85 table no returned table covered.
87 With `orphans`, a returned grid over a region the document holds nothing
88 for — no table, no picture, no figure, and no text the parse read there —
89 is captured too, through the same checks, with the cell check asking an
90 agent about every cell because there is no parse reading to check against.
91 A map's legend is the standing example: the page prints the values, the
92 parse's text layer has nothing, and the scan's grid is the only reading
93 anywhere. The caller gives such a table a home by inserting it.
94 """
95 by_ref = {t.self_ref: t for t in document.tables}
96 work: List[Tuple[PageScan, ScannedTable, str, int]] = []
97 errors: List[str] = []
98 for scan in scans:
99 width, height = page_dims.get(scan.page, (612.0, 792.0))
100 ordinals = _page_ordinals(document, scan.page)
101 # A returned table stands for one parse table. Claiming it keeps two
102 # nominated tables in the same region from both taking the same grid.
103 available = list(scan.tables)
104 for ref in scan.table_refs:
105 table_item = by_ref.get(ref)
106 scanned = _best_match(table_item, available, width, height, match_fraction)
107 if scanned is None:
108 errors.append(
109 f"page {scan.page}: table {ref} was read off the page image and the scan "
110 "returned no table over it; what stands in the document is the parse's reading"
111 )
112 continue
113 scanned.table_ref = ref
114 available.remove(scanned)
115 work.append((scan, scanned, ref, ordinals.get(ref, 1)))
117 # A table the parse filed as a picture. The claim above runs over the
118 # parse's tables and this region is not among them, so nothing took the
119 # grid. It is still a table the page prints and the scan read, so the
120 # picture stands in as its home and the grid travels the same path every
121 # other scanned table travels — including the cell check, which asks an
122 # agent about every cell when there is no parse reading to check against.
123 ordinal = len(ordinals)
124 for picture, scanned in _over_pictures(document, scan, available, width, height, match_fraction):
125 scanned.picture_ref = picture.self_ref
126 available.remove(scanned)
127 ordinal += 1
128 work.append((scan, scanned, picture.self_ref, ordinal))
130 if orphans:
131 for scanned in _orphan_grids(document, scan, available, width, height, match_fraction):
132 available.remove(scanned)
133 ordinal += 1
134 work.append((scan, scanned, "", ordinal))
136 if not work:
137 return [], errors
139 with tempfile.TemporaryDirectory(prefix="quber-scan-table-") as tmp:
140 pages = sorted({scan.page for scan, _s, _r, _o in work})
141 rendered = await asyncio.gather(
142 *(
143 asyncio.to_thread(render_page, source, page, assembly.dpi, Path(tmp) / f"page-{page:04d}.png")
144 for page in pages
145 )
146 )
147 images = {page: image for page, (image, _w, _h) in zip(pages, rendered, strict=True)}
148 tables = await asyncio.gather(
149 *(
150 _capture(
151 scanned,
152 by_ref.get(ref),
153 images[scan.page],
154 page_dims.get(scan.page, (612.0, 792.0)),
155 ordinal,
156 assembly,
157 reader,
158 )
159 for scan, scanned, ref, ordinal in work
160 )
161 )
163 for (_scan, scanned, _ref, _ordinal), table in zip(work, tables, strict=True):
164 scanned.table_id = table.table_id
165 # The scan bounds a table generously, so its box can run past the last data
166 # row. Record the true end of the content from the corrected last row.
167 await asyncio.to_thread(apply_content_regions, list(tables), str(source))
168 return list(tables), errors
171async def _capture(
172 scanned: ScannedTable,
173 table_item: Optional[TableItem],
174 page_image: Path,
175 page_dims: Tuple[float, float],
176 ordinal: int,
177 assembly: TableAssembly,
178 reader: Optional[CellReader],
179) -> ExtractedTable:
180 """One returned grid, checked against the parse, then corrected and grounded."""
181 width, height = page_dims
182 boxes = [[point_box(box, width, height) for box in row] for row in scanned.cell_boxes]
183 bbox = point_box(scanned.box, width, height)
184 logger.info(
185 "Table capture: page {} reading a {}x{} table off the page image",
186 scanned.page,
187 len(scanned.cells),
188 len(scanned.cells[0]) if scanned.cells else 0,
189 )
190 # Values are settled against the printed page before the structure
191 # correction runs, so that step works on a grid whose values are agreed.
192 cells = scanned.cells
193 if bbox is not None:
194 crop = await asyncio.to_thread(
195 crop_region_png, page_image, table_crop_box(bbox, height), assembly.dpi
196 )
197 cells, _read = await crosscheck_cells(scanned, table_item, page_dims, crop, reader)
198 return await assemble_table(
199 assembly,
200 page=scanned.page,
201 ordinal=ordinal,
202 page_image=page_image,
203 page_dims=page_dims,
204 cells=cells,
205 cell_boxes=boxes,
206 bbox=bbox,
207 som_region=norm_box(scanned.box),
208 kind="image_table",
209 # The page prints no text under this table — that is what made it one of
210 # ours. So there is no text layer to check the corrected numbers against,
211 # and checking them against the scan's own grid would reject every repair
212 # of what the scan misread.
213 ground_values=False,
214 )
217def _over_pictures(
218 document: DoclingDocument,
219 scan: PageScan,
220 available: List[ScannedTable],
221 width: float,
222 height: float,
223 match_fraction: float,
224) -> List[Tuple[PictureItem, ScannedTable]]:
225 """Each unclaimed grid paired with the picture it was printed over.
227 A picture stands in for a table only where the parse filed the region as a
228 picture and as nothing else. Two things disqualify it.
230 A parse table over the same region means the document already holds that
231 table, read from a text layer by the table engine. The scan read it too, and
232 that reading is ignored for the reason the module's own rule gives: a vetted
233 extraction is not replaced by a scan of the same region. Adding it beside the
234 table would state the same figures twice.
236 A figure over the same region means the scan called it a picture, and the
237 graft attaches that reading to the picture as a description. Taking it as a
238 table as well would put one thing in the document twice again.
239 """
240 figures = [box for box in (norm_box(chart.box) for chart in scan.figures) if box is not None]
241 parse_tables = [
242 box
243 for box in (
244 prov_box(table, width, height)
245 for table in document.tables
246 if table.prov and table.prov[0].page_no == scan.page
247 )
248 if box is not None
249 ]
250 pairs: List[Tuple[PictureItem, ScannedTable]] = []
251 for picture in document.pictures:
252 if not picture.prov or picture.prov[0].page_no != scan.page:
253 continue
254 if picture_classes(picture) and all(c in FURNITURE_CLASSES for c in picture_classes(picture)):
255 continue
256 # A reader has already read this region and the document records which.
257 # Reading it again would state the same table twice.
258 if getattr(picture.meta, READ_BY_FIELD, None):
259 continue
260 target = prov_box(picture, width, height)
261 if target is None:
262 continue
263 if any(_covers(other, target, match_fraction) for other in figures + parse_tables):
264 continue
265 remaining = [t for t in available if all(t is not p for _pic, p in pairs)]
266 scanned = _best_match(picture, remaining, width, height, match_fraction)
267 if scanned is not None:
268 pairs.append((picture, scanned))
269 return pairs
272def _covers(box: NormBox, target: NormBox, match_fraction: float) -> bool:
273 """Do the two boxes overlap enough, either way round, to be the same region?"""
274 return max(coverage_fraction(box, target), coverage_fraction(target, box)) >= match_fraction
277#: Labels that do not disqualify a region from being empty of parse text. A
278#: heading floats over a region without holding its content, and page furniture
279#: belongs to the page, not the region.
280_SPINE_LABELS = frozenset(
281 (DocItemLabel.SECTION_HEADER, DocItemLabel.PAGE_HEADER, DocItemLabel.PAGE_FOOTER, DocItemLabel.TITLE)
282)
285def _orphan_grids(
286 document: DoclingDocument,
287 scan: PageScan,
288 available: List[ScannedTable],
289 width: float,
290 height: float,
291 match_fraction: float,
292) -> List[ScannedTable]:
293 """Each unclaimed grid over a region the document holds nothing for.
295 Nothing means nothing: no parse table (the table engine read a text layer
296 there), no figure (the graft attaches that reading to a picture), no
297 picture (the picture tier above homes those), and no content text the
298 parse read inside the region. That last condition is the line between a
299 map's legend — printed values the parse's text layer never captured, where
300 the scan's grid is the only reading anywhere — and a stat panel, whose
301 strings the parse does hold and whose binding is the block grouping's job.
302 Inserting a grid over a region the parse read would state its content
303 twice.
304 """
305 figures = [box for box in (norm_box(chart.box) for chart in scan.figures) if box is not None]
306 parse_tables = [
307 box
308 for box in (
309 prov_box(table, width, height)
310 for table in document.tables
311 if table.prov and table.prov[0].page_no == scan.page
312 )
313 if box is not None
314 ]
315 pictures = [
316 box
317 for box in (
318 prov_box(picture, width, height)
319 for picture in document.pictures
320 if picture.prov and picture.prov[0].page_no == scan.page
321 )
322 if box is not None
323 ]
324 text_centers = []
325 for item in document.texts:
326 if item.label in _SPINE_LABELS or not (item.text or "").strip():
327 continue
328 if not item.prov or item.prov[0].page_no != scan.page:
329 continue
330 box = prov_box(item, width, height)
331 if box is not None:
332 text_centers.append(((box[0] + box[2]) / 2, (box[1] + box[3]) / 2))
334 found: List[ScannedTable] = []
335 for scanned in available:
336 if not scanned.cells:
337 continue
338 box = norm_box(scanned.box)
339 if box is None:
340 continue
341 if any(_covers(other, box, match_fraction) for other in figures + parse_tables + pictures):
342 continue
343 if any(box[0] <= cx <= box[2] and box[1] <= cy <= box[3] for cx, cy in text_centers):
344 continue
345 found.append(scanned)
346 return found
349def _page_ordinals(document: DoclingDocument, page: int) -> Dict[str, int]:
350 """Each table's position among the tables its page prints, counting from one."""
351 refs = [t.self_ref for t in document.tables if t.prov and t.prov[0].page_no == page]
352 return {ref: i + 1 for i, ref in enumerate(refs)}
355def _best_match(
356 item: Optional[DocItem],
357 scanned: List[ScannedTable],
358 width: float,
359 height: float,
360 match_fraction: float,
361) -> Optional[ScannedTable]:
362 """The returned table covering `item` best, or none past the threshold.
364 `item` is whatever the parse holds over that region — a table it detected as
365 one, or a picture it filed a table as.
366 """
367 if item is None:
368 return None
369 target = prov_box(item, width, height)
370 if target is None:
371 return None
373 best: Optional[ScannedTable] = None
374 best_cov = match_fraction
375 for candidate in scanned:
376 box = norm_box(candidate.box)
377 if box is None:
378 continue
379 cov = max(coverage_fraction(box, target), coverage_fraction(target, box))
380 if cov >= best_cov:
381 best, best_cov = candidate, cov
382 return best