Coverage for src / quber / core / fusion / matching.py: 96%
142 statements
« prev ^ index » next coverage.py v7.14.0, created at 2026-09-23 22:14 -0400
1"""Match the Set-of-Mark/Camelot tables against docling's tables and pictures.
3One page at a time, every SoM table is overlapped against every docling table
4and picture on that page. There are only a handful of tables per page, so the
5match is all-pairs; no spatial index. Boxes from the three sources are first
6normalized to one frame — 0..1 with the page top-left as origin — so they
7compare directly:
9- docling table / picture boxes: PDF points, bottom-left origin (flip y).
10- SoM `som_region`: already normalized top-left (used as-is).
11- Camelot `bbox` (fallback when a SoM table has no region): PDF points,
12 bottom-left origin.
14An edge is drawn when either box covers the other past `MATCH_FRACTION`, so a
15large SoM region containing a small docling table and a small SoM table inside a
16large docling table both register. SoM tables and docling tables are then grouped
17into regions by connected overlap, and each region's cardinality is read off the
18matrix in `models.py`.
19"""
21from __future__ import annotations
23from typing import Dict, List, Optional, Tuple
25from docling_core.types.doc.base import CoordOrigin
26from docling_core.types.doc.document import PictureItem, TableItem
28from quber.core.extractors.base import ExtractedTable
29from quber.core.extractors.camelot.correspondence.geometry import (
30 camelot_bbox_to_norm,
31 coverage_fraction,
32)
33from quber.core.fusion.models import RegionMatch
34from quber.core.parsers.result import ParseResult, TableProvenance
36#: The one knob: minimum box-overlap coverage for two tables to be considered the
37#: same region. Carried over from the correspondence matcher's default.
38MATCH_FRACTION = 0.20
40NormBox = Tuple[float, float, float, float]
43def match_tables(
44 parse: ParseResult,
45 som_tables: List[ExtractedTable],
46 match_fraction: float = MATCH_FRACTION,
47) -> List[RegionMatch]:
48 """Match the SoM tables against the docling document, page by page.
50 Returns one `RegionMatch` per region: every SoM table lands in exactly one
51 record (replace / som_merged / docling_undercount / image_table / chart /
52 docling_miss), and every docling table no SoM table matched yields a
53 `som_miss` record.
54 """
55 document = parse.document
56 page_dims = {p.page_no: (p.width, p.height) for p in parse.pages}
57 prov_by_ref = parse.provenance_by_ref()
59 pages = sorted(
60 {t.page for t in som_tables}
61 | {t.prov[0].page_no for t in document.tables if t.prov}
62 | {p.prov[0].page_no for p in document.pictures if p.prov}
63 )
65 matches: List[RegionMatch] = []
66 for page in pages:
67 width, height = page_dims.get(page, (612.0, 792.0))
68 matches.extend(
69 _match_page(
70 page,
71 width,
72 height,
73 som_tables,
74 document.tables,
75 document.pictures,
76 prov_by_ref,
77 match_fraction,
78 )
79 )
80 return matches
83def _match_page(
84 page: int,
85 width: float,
86 height: float,
87 som_tables: List[ExtractedTable],
88 docling_tables: List[TableItem],
89 docling_pictures: List[PictureItem],
90 prov_by_ref: Dict[str, TableProvenance],
91 match_fraction: float,
92) -> List[RegionMatch]:
93 soms = [
94 (i, _som_box(som_tables[i], width, height))
95 for i in range(len(som_tables))
96 if som_tables[i].page == page
97 ]
98 soms = [(i, b) for i, b in soms if b is not None]
99 dtables = [
100 (t.self_ref, _docling_box(t, width, height))
101 for t in docling_tables
102 if t.prov and t.prov[0].page_no == page
103 ]
104 dtables = [(ref, b) for ref, b in dtables if b is not None]
105 dpics = [
106 (p.self_ref, _docling_box(p, width, height), _picture_classes(p))
107 for p in docling_pictures
108 if p.prov and p.prov[0].page_no == page
109 ]
110 dpics = [(ref, b, c) for ref, b, c in dpics if b is not None]
112 # SoM table -> docling tables it overlaps, and the best coverage seen.
113 s_to_d: Dict[int, List[str]] = {i: [] for i, _ in soms}
114 d_to_s: Dict[str, List[int]] = {ref: [] for ref, _ in dtables}
115 best_cov: Dict[int, float] = {i: 0.0 for i, _ in soms}
116 for i, sbox in soms:
117 for ref, dbox in dtables:
118 cov = max(coverage_fraction(sbox, dbox), coverage_fraction(dbox, sbox))
119 if cov >= match_fraction:
120 s_to_d[i].append(ref)
121 d_to_s[ref].append(i)
122 best_cov[i] = max(best_cov[i], cov)
124 # Group SoM tables and docling tables that overlap into regions.
125 components = _connected_regions([i for i, _ in soms], [ref for ref, _ in dtables], s_to_d)
127 matches: List[RegionMatch] = []
128 matched_soms: set[int] = set()
129 matched_dtables: set[str] = set()
130 for som_group, dtable_group in components:
131 matched_soms.update(som_group)
132 matched_dtables.update(dtable_group)
133 matches.append(_classify_region(page, som_group, dtable_group, prov_by_ref, som_tables, best_cov))
135 # SoM tables with no docling-table overlap: chart (over a picture) or a miss
136 # docling did not see at all.
137 for i, sbox in soms:
138 if i in matched_soms:
139 continue
140 pic_refs: List[str] = []
141 pic_classes: List[str] = []
142 cov = 0.0
143 for ref, pbox, classes in dpics:
144 c = max(coverage_fraction(sbox, pbox), coverage_fraction(pbox, sbox))
145 if c >= match_fraction:
146 pic_refs.append(ref)
147 pic_classes.extend(classes)
148 cov = max(cov, c)
149 if pic_refs:
150 matches.append(
151 RegionMatch(
152 page=page,
153 kind="chart",
154 som_indices=[i],
155 docling_picture_refs=pic_refs,
156 picture_classes=pic_classes,
157 overlap=cov,
158 detail="SoM table over a docling picture, no docling table",
159 )
160 )
161 else:
162 matches.append(
163 RegionMatch(
164 page=page,
165 kind="docling_miss",
166 som_indices=[i],
167 detail="SoM table with no overlapping docling table or picture",
168 )
169 )
171 # Docling tables no SoM table matched: SoM missed a table docling found.
172 for ref, _ in dtables:
173 if ref in matched_dtables:
174 continue
175 matches.append(
176 RegionMatch(
177 page=page,
178 kind="som_miss",
179 docling_table_refs=[ref],
180 detail="docling table with no overlapping SoM table",
181 )
182 )
184 return matches
187def _classify_region(
188 page: int,
189 som_group: List[int],
190 dtable_group: List[str],
191 prov_by_ref: Dict[str, TableProvenance],
192 som_tables: List[ExtractedTable],
193 best_cov: Dict[int, float],
194) -> RegionMatch:
195 """Read one region's cardinality off the matrix."""
196 c = len(som_group)
197 d = len(dtable_group)
198 cov = max((best_cov.get(i, 0.0) for i in som_group), default=0.0)
200 # A 1:1 region where docling read the table off the image and Camelot found
201 # nothing is a table rendered as an image, not a plain replace.
202 if c == 1 and d == 1:
203 prov = prov_by_ref.get(dtable_group[0])
204 som = som_tables[som_group[0]]
205 no_text_layer = prov is not None and prov.verdict in ("ocr", "empty")
206 camelot_empty = not (som.markdown or "").strip() or som.camelot_accuracy <= 0.0
207 if no_text_layer and camelot_empty:
208 return RegionMatch(
209 page=page,
210 kind="image_table",
211 som_indices=som_group,
212 docling_table_refs=dtable_group,
213 overlap=cov,
214 detail=f"docling table read by OCR (verdict={prov.verdict if prov else 'n/a'}), Camelot empty",
215 )
217 if d > c:
218 kind: str = "som_merged"
219 detail = f"docling found {d} tables, SoM {c}: SoM merged stacked tables"
220 elif c > d:
221 kind = "docling_undercount"
222 detail = f"SoM found {c} tables, docling {d}: docling dropped/merged a complex table"
223 else:
224 kind = "replace"
225 detail = f"{c} SoM table(s) agree with {d} docling table(s)"
227 return RegionMatch(
228 page=page,
229 kind=kind,
230 som_indices=som_group, # type: ignore[arg-type]
231 docling_table_refs=dtable_group,
232 overlap=cov,
233 detail=detail,
234 )
237def _connected_regions(
238 som_ids: List[int], dtable_refs: List[str], s_to_d: Dict[int, List[str]]
239) -> List[Tuple[List[int], List[str]]]:
240 """Group SoM tables and docling tables into connected overlap regions.
242 Only SoM tables that overlap at least one docling table (and the docling
243 tables they reach) form regions here; unmatched SoM tables and unmatched
244 docling tables are handled by the caller.
245 """
246 parent: Dict[str, str] = {}
248 def key_s(i: int) -> str:
249 return f"s{i}"
251 def key_d(ref: str) -> str:
252 return f"d{ref}"
254 def find(x: str) -> str:
255 parent.setdefault(x, x)
256 while parent[x] != x:
257 parent[x] = parent[parent[x]]
258 x = parent[x]
259 return x
261 def union(a: str, b: str) -> None:
262 parent.setdefault(a, a)
263 parent.setdefault(b, b)
264 parent[find(a)] = find(b)
266 for i, refs in s_to_d.items():
267 for ref in refs:
268 union(key_s(i), key_d(ref))
270 groups: Dict[str, Tuple[List[int], List[str]]] = {}
271 for i in som_ids:
272 if not s_to_d.get(i):
273 continue
274 root = find(key_s(i))
275 groups.setdefault(root, ([], []))[0].append(i)
276 for ref in dtable_refs:
277 node = key_d(ref)
278 if node not in parent:
279 continue
280 root = find(node)
281 groups.setdefault(root, ([], []))[1].append(ref)
283 return [g for g in groups.values() if g[0]]
286def _som_box(table: ExtractedTable, width: float, height: float) -> Optional[NormBox]:
287 """A SoM table's box in the normalized top-left frame.
289 Prefers the SoM region (already normalized); falls back to Camelot's bbox
290 (PDF points, bottom-left) when the region is absent.
291 """
292 if table.som_region is not None:
293 return table.som_region
294 if table.bbox is not None:
295 return camelot_bbox_to_norm(table.bbox, width, height)
296 return None
299def _docling_box(item: TableItem | PictureItem, width: float, height: float) -> Optional[NormBox]:
300 """A docling item's provenance bbox in the normalized top-left frame."""
301 if not item.prov:
302 return None
303 bbox = item.prov[0].bbox
304 if bbox.coord_origin == CoordOrigin.TOPLEFT:
305 return (bbox.l / width, bbox.t / height, bbox.r / width, bbox.b / height)
306 return camelot_bbox_to_norm((bbox.l, bbox.b, bbox.r, bbox.t), width, height)
309def _picture_classes(picture: PictureItem) -> List[str]:
310 """docling's top predicted class for a picture (empty if unclassified).
312 Predictions are ordered by descending confidence, so the first entry is
313 docling's call (e.g. `bar_chart`). Reads the current `meta.classification`
314 field, falling back to the deprecated `annotations` list for older output.
315 """
316 meta = getattr(picture, "meta", None)
317 classification = getattr(meta, "classification", None) if meta is not None else None
318 predictions = getattr(classification, "predictions", None) if classification is not None else None
319 if predictions:
320 return [predictions[0].class_name]
322 classes: List[str] = []
323 for ann in getattr(picture, "annotations", []):
324 predicted = getattr(ann, "predicted_classes", None)
325 if predicted:
326 classes.append(predicted[0].class_name)
327 return classes