Coverage for src / quber / core / figures / correct.py: 93%
75 statements
« prev ^ index » next coverage.py v7.14.0, created at 2026-09-23 22:14 -0400
1"""Settle what a figure leaves on the page once the scan has read the figure.
3Two answers come back from one look at the page: the text around the figure that
4is only its furniture, which is removed, and the footnote markers its labels
5carry, which are recorded on the picture.
7The parse lifts what it can off a picture and files the fragments as text. Where
8it files them under the picture, the graft removes them with the picture's
9children. Where it files them under the page body instead — which it does
10page to page with no pattern — they survive, and the document then holds a
11chart's axis printed as forty separate records beside the reading that
12supersedes them. Indexed for retrieval, each one carries the chart's heading, so
13a gridline labelled 3.50 becomes a chunk that answers a question about the
14weighted average risk rating with a number off the ruler.
16The candidates are the plain text items overlapping a figure the scan read. That
17is a bound on what can be removed rather than a decision about it: a heading, a
18caption, a footnote or a page mark is never a candidate, whatever it overlaps,
19and neither is a page the scan returned no figure for.
21Deciding among the candidates is left to the agent, which is shown the page.
22Nothing here inspects the text.
24The markers are asked for on every page the scan returned a figure on, whether or
25not that page has candidate text. A chart printing "(1,2)" on its title leaves
26nothing behind to sweep, and it still carries a qualification the description
27alone does not state.
28"""
30from __future__ import annotations
32import asyncio
33import tempfile
34from pathlib import Path
35from typing import Dict, List, Optional, Sequence, Tuple
37from docling_core.types.doc.document import DoclingDocument, PictureItem, PictureMeta, TextItem
38from docling_core.types.doc.labels import DocItemLabel
39from loguru import logger
41from quber.agents.figure_correction import Figure, FigureCorrector, Fragment, Marker, Note
42from quber.core.figures.geometry import NormBox, norm_box, prov_box
43from quber.core.figures.graft import FOOTNOTE_MARKS_FIELD, FOOTNOTES_FIELD
44from quber.core.figures.models import FigureRecord, PageScan, RemovedText
45from quber.files.pdf import render_page
47#: Page rasterization for the correction agent. It judges a fragment by where it
48#: sits on the page, which is the same reading the table pipeline renders for.
49CORRECTION_DPI = 200
52async def correct_figures(
53 document: DoclingDocument,
54 scans: Sequence[PageScan],
55 page_dims: Dict[int, Tuple[float, float]],
56 source: Path,
57 finder: Optional[FigureCorrector] = None,
58) -> List[RemovedText]:
59 """Sweep each page's figure furniture from `document` and record its markers.
61 Each removal is returned with its text and its box, because a completeness
62 check reading the source PDF cannot otherwise tell this from a silent loss:
63 both leave printed text the document does not hold.
65 Mutates `document` in place. Does nothing when the sweep is switched off, when
66 no page carries a figure, or when a page's call fails.
67 """
68 if finder is None:
69 return []
71 work = [
72 (scan, figures, candidates(document, scan, page_dims))
73 for scan in scans
74 for figures in [[figure for figure in scan.figures if figure.text]]
75 if figures
76 ]
77 if not work:
78 return []
80 with tempfile.TemporaryDirectory(prefix="quber-correction-") as tmp:
81 rendered = await asyncio.gather(
82 *(
83 asyncio.to_thread(
84 render_page, source, scan.page, CORRECTION_DPI, Path(tmp) / f"page-{scan.page:04d}.png"
85 )
86 for scan, _f, _c in work
87 )
88 )
89 verdicts = await asyncio.gather(
90 *(
91 finder.correct_figure(
92 Path(image).read_bytes(),
93 [Figure(index=i, description=f.text) for i, f in enumerate(figures)],
94 [Fragment(index=i, text=item.text) for i, item in enumerate(eligible)],
95 )
96 for (_scan, figures, eligible), (image, _w, _h) in zip(work, rendered, strict=True)
97 )
98 )
100 doomed = []
101 removed: List[RemovedText] = []
102 for (scan, figures, eligible), correction in zip(work, verdicts, strict=True):
103 width, height = page_dims.get(scan.page, (612.0, 792.0))
104 _record_markers(document, figures, correction.markers, correction.notes, scan.page)
105 marked = [(eligible[f.index], f.reason) for f in correction.furniture]
106 if marked:
107 logger.info(
108 "Figure correction: page {} removing {} of {} fragment(s) the scan's figures already read",
109 scan.page,
110 len(marked),
111 len(eligible),
112 )
113 for item, reason in marked:
114 doomed.append(item)
115 removed.append(
116 RemovedText(
117 page=scan.page,
118 text=item.text,
119 box=prov_box(item, width, height),
120 reason=reason,
121 )
122 )
124 if not doomed:
125 return []
126 # Boxes are read before the delete: removing an item renumbers the ones after
127 # it, and a reference resolved afterwards would name a different element.
128 document.delete_items(node_items=doomed)
129 return removed
132def candidates(
133 document: DoclingDocument, scan: PageScan, page_dims: Dict[int, Tuple[float, float]]
134) -> List[TextItem]:
135 """The plain text items on `scan`'s page that overlap a figure it returned."""
136 boxes = [box for box in (norm_box(chart.box) for chart in scan.figures) if box is not None]
137 if not boxes:
138 return []
139 width, height = page_dims.get(scan.page, (612.0, 792.0))
140 found: List[TextItem] = []
141 for item in document.texts:
142 if item.label != DocItemLabel.TEXT or not (item.text or "").strip():
143 continue
144 if not item.prov or item.prov[0].page_no != scan.page:
145 continue
146 # A fragment already filed under a picture is removed with the picture's
147 # children, so it is not put to the agent as well.
148 if isinstance(item.parent.resolve(document) if item.parent else None, PictureItem):
149 continue
150 box = prov_box(item, width, height)
151 if box is not None and any(_overlaps(figure, box) for figure in boxes):
152 found.append(item)
153 return found
156def _record_markers(
157 document: DoclingDocument,
158 figures: Sequence[FigureRecord],
159 markers: Sequence[Marker],
160 notes: Sequence[Note],
161 page: int,
162) -> None:
163 """Write each figure's markers and printed notes onto its picture.
165 A figure the graft never placed has no picture to carry them. The run already
166 reports that figure as unplaced, so the markers are logged and dropped rather
167 than reported a second time.
168 """
169 if not markers and not notes:
170 return
171 by_ref = {picture.self_ref: picture for picture in document.pictures}
172 for index, figure in enumerate(figures):
173 mine = [m for m in markers if m.figure == index]
174 read = [n for n in notes if n.figure == index]
175 if not mine and not read:
176 continue
177 picture = by_ref.get(figure.picture_ref or "")
178 if picture is None:
179 logger.warning(
180 "Figure markers: page {} figure {} carries {} marker(s) and sits on no picture",
181 page,
182 index,
183 len(mine),
184 )
185 continue
186 base = picture.meta or PictureMeta()
187 picture.meta = base.model_copy(
188 update={
189 FOOTNOTE_MARKS_FIELD: [{"marker": m.marker, "kind": m.kind, "label": m.label} for m in mine],
190 FOOTNOTES_FIELD: [{"marker": n.marker, "text": n.text} for n in read],
191 }
192 )
193 logger.info(
194 "Figure markers: page {} picture {} carries {} with {} note(s) read off the page",
195 page,
196 picture.self_ref,
197 ", ".join(f"{m.marker} ({m.kind})" for m in mine) or "no markers",
198 len(read),
199 )
202def _overlaps(figure: NormBox, text: NormBox) -> bool:
203 """Do the two boxes share any area at all?
205 Any overlap, with no fraction to tune. Measured on three decks: every text
206 item on a scanned page either sits well inside a figure or misses it
207 entirely, and not one landed in between, so a threshold decided nothing and
208 only added a constant fitted to whichever deck it was read off.
209 """
210 wide = min(figure[2], text[2]) - max(figure[0], text[0])
211 tall = min(figure[3], text[3]) - max(figure[1], text[1])
212 return wide > 0 and tall > 0