Coverage for src / quber / core / extractors / camelot / correspondence / geometry.py: 73%
81 statements
« prev ^ index » next coverage.py v7.14.0, created at 2026-09-23 22:14 -0400
1"""
2Coordinate-frame and region helpers for the correspondence extractor.
4Three frames are in play, and a mix-up between them fails silently:
6- Camelot bboxes: PDF points, origin bottom-left.
7- The detector's boxes: normalized 0..1, origin top-left.
8- PyMuPDF words and the rendered page image: points / pixels, origin
9 top-left.
11The conversions live together here, alongside the region slicing built
12on them (text-layer words inside a box, page-image crop of a box).
13Everything in this module is deterministic; no LLM, no Camelot.
14"""
16from __future__ import annotations
18import io
19from pathlib import Path
20from typing import List, Optional, Tuple
22from quber.agents.detector import DetectedTable
24WordBox = Tuple[float, float, float, float, str]
27def page_size_pts(page_image: Path, dpi: int) -> Tuple[float, float]:
28 """Page width/height in PDF points, from the rendered image and its
29 DPI. Camelot bboxes are in points; this lets us normalize them into
30 the detector's 0..1 top-left frame.
31 """
32 from PIL import Image
34 w, h = Image.open(page_image).size
35 return w * 72.0 / dpi, h * 72.0 / dpi
38def camelot_bbox_to_norm(
39 bbox: Tuple[float, float, float, float], page_w_pts: float, page_h_pts: float
40) -> Tuple[float, float, float, float]:
41 """Convert a Camelot bbox (PDF points, bottom-left origin) to the
42 detector's frame (normalized 0..1, top-left origin).
43 """
44 x1, y1, x2, y2 = bbox
45 x_left = min(x1, x2) / page_w_pts
46 x_right = max(x1, x2) / page_w_pts
47 y_bottom_pts, y_top_pts = min(y1, y2), max(y1, y2)
48 top = 1.0 - y_top_pts / page_h_pts
49 bottom = 1.0 - y_bottom_pts / page_h_pts
50 return (x_left, top, x_right, bottom)
53def norm_bbox_to_table_area(
54 bbox: Tuple[float, float, float, float], page_w_pts: float, page_h_pts: float
55) -> str:
56 """Inverse of camelot_bbox_to_norm: a normalized 0..1 top-left box ->
57 a Camelot `table_areas` string 'x1,y1,x2,y2' in PDF points (origin
58 bottom-left, so y1=top is the larger value).
59 """
60 x1, top, x2, bottom = bbox
61 left = max(0.0, min(x1, x2)) * page_w_pts
62 right = min(1.0, max(x1, x2)) * page_w_pts
63 top_pts = (1.0 - max(0.0, min(top, bottom))) * page_h_pts
64 bottom_pts = (1.0 - min(1.0, max(top, bottom))) * page_h_pts
65 return f"{left:.1f},{top_pts:.1f},{right:.1f},{bottom_pts:.1f}"
68def bbox_to_top_left(
69 bbox: Tuple[float, float, float, float], page_h_pts: float, pad: float = 6.0
70) -> Tuple[float, float, float, float]:
71 """Camelot bbox (PDF points, bottom-left origin) -> a padded (left, top,
72 right, bottom) box in the top-left point frame used by PyMuPDF words and by
73 the rendered page image."""
74 x1, y1, x2, y2 = bbox
75 return (
76 min(x1, x2) - pad,
77 page_h_pts - max(y1, y2) - pad,
78 max(x1, x2) + pad,
79 page_h_pts - min(y1, y2) + pad,
80 )
83# The table image crop reaches this many points above the data box to capture
84# the title/caption that sits just above the grid. Only the image is widened;
85# a value text-layer slice stays tight to the data box.
86CAPTION_PAD_PTS = 28.0
89def table_crop_box(
90 bbox: Tuple[float, float, float, float], page_h_pts: float
91) -> Tuple[float, float, float, float]:
92 """The image crop box (left, top, right, bottom; top-left points) for a table:
93 its data box widened upward by the caption pad. The same crop the
94 structure-correction agent and the source locator both see, so grid cells
95 read off that crop map back to page points consistently."""
96 left, top, right, bottom = bbox_to_top_left(bbox, page_h_pts)
97 return (left, max(0.0, top - CAPTION_PAD_PTS), right, bottom)
100def coverage_fraction(
101 chunk: Tuple[float, float, float, float], det: Tuple[float, float, float, float]
102) -> float:
103 """How much of the DETECTED box the chunk covers: intersection area as
104 a fraction of the detected box's area. Both rectangles are normalized
105 0..1 top-left.
107 We measure against the detected box (not the chunk) so the metric
108 answers "does this chunk fill this table" and is comparable across
109 chunks of different sizes when picking the best match for a table.
110 """
111 ix1, iy1 = max(chunk[0], det[0]), max(chunk[1], det[1])
112 ix2, iy2 = min(chunk[2], det[2]), min(chunk[3], det[3])
113 inter = max(0.0, ix2 - ix1) * max(0.0, iy2 - iy1)
114 det_area = max(1e-9, (det[2] - det[0]) * (det[3] - det[1]))
115 return inter / det_area
118def neighbor_bounded_bbox(
119 detected: List[DetectedTable], di: int, pad: float = 0.02
120) -> Optional[Tuple[float, float, float, float]]:
121 """The crop box for the table at index `di`, clamped vertically so it
122 does not bleed into the adjacent detected tables.
124 `detected` is in top-to-bottom order. The completeness audit checks
125 the table's top/bottom edges, so a crop that includes a neighbor's
126 rows makes it read the neighbor as a continuation. We split the gutter
127 to each neighbor's midpoint; only the outer edges keep the small pad.
128 """
129 d = detected[di]
130 if d.bbox is None:
131 return None
132 x1, y1, x2, y2 = d.bbox
133 top = min(y1, y2)
134 bottom = max(y1, y2)
135 prev_bbox = detected[di - 1].bbox if di > 0 else None
136 if prev_bbox is not None:
137 prev_bottom = max(prev_bbox[1], prev_bbox[3])
138 top = max(top, (top + prev_bottom) / 2)
139 else:
140 top = top - pad
141 next_bbox = detected[di + 1].bbox if di + 1 < len(detected) else None
142 if next_bbox is not None:
143 next_top = min(next_bbox[1], next_bbox[3])
144 bottom = min(bottom, (bottom + next_top) / 2)
145 else:
146 bottom = bottom + pad
147 return (x1, top, x2, bottom)
150def region_text_in_bbox(words: List[WordBox], region: Tuple[float, float, float, float]) -> str:
151 """Text-layer words inside the region (top-left point frame), grouped into
152 lines top-to-bottom, left-to-right. This is the value 'well' for the table,
153 scoped so no adjacent table can bleed in."""
154 left, top, right, bottom = region
155 ins = [w for w in words if w[0] >= left and w[2] <= right and w[1] >= top and w[3] <= bottom]
156 ins.sort(key=lambda w: (round(w[1] / 3), w[0]))
157 lines: List[str] = []
158 cur_y: Optional[float] = None
159 cur: List[str] = []
160 for _x0, y0, _x1, _y1, txt in ins:
161 if cur_y is None or abs(y0 - cur_y) > 4:
162 if cur:
163 lines.append(" ".join(cur))
164 cur, cur_y = [txt], y0
165 else:
166 cur.append(txt)
167 if cur:
168 lines.append(" ".join(cur))
169 return "\n".join(lines)
172def crop_region_png(page_image: Path, region: Tuple[float, float, float, float], dpi: int) -> bytes:
173 """Crop the rendered page PNG to the table region (top-left points -> pixels
174 at `dpi`). This cropped image is the structural arbiter for the vetting LLM."""
175 from PIL import Image
177 s = dpi / 72.0
178 left, top, right, bottom = region
179 img = Image.open(page_image)
180 crop = img.crop((int(max(0.0, left * s)), int(max(0.0, top * s)), int(right * s), int(bottom * s)))
181 buf = io.BytesIO()
182 crop.save(buf, format="PNG")
183 return buf.getvalue()