Coverage for src / quber / core / extractors / camelot / acquire.py: 60%
122 statements
« prev ^ index » next coverage.py v7.14.0, created at 2026-09-23 22:14 -0400
1"""
2Camelot acquisition layer: get candidate cell grids out of a PDF.
4This module knows nothing about LLMs or the pipeline graph. It renders
5page images, runs both Camelot flavors in spawn-context subprocesses,
6and converts each Camelot DataFrame into a serializable
7`CamelotCandidate`.
9Both Camelot flavors run in parallel:
11- `lattice` — visible grid lines.
12- `stream` — whitespace-based; more permissive.
14There is deliberately no lattice-first / fall-back-to-stream gate: on
15pages where tables are aligned by whitespace alone, lattice still
16returns grids — empty shells — so "fall back only when lattice finds
17nothing" let those shells through as false positives. `is_content_empty`
18catches the shells instead.
19"""
21from __future__ import annotations
23import multiprocessing as mp
24from concurrent.futures import ProcessPoolExecutor
25from concurrent.futures import TimeoutError as FutureTimeoutError
26from pathlib import Path
27from typing import TYPE_CHECKING, Any, List, Literal, Optional, Tuple
29from loguru import logger
30from pydantic import BaseModel, Field
32if TYPE_CHECKING:
33 from pandas import DataFrame
35Flavor = Literal["lattice", "stream"]
37CAMELOT_FLAVOR_TIMEOUT_S = 60.0
40def render_pages(source: Path, dpi: int, out_dir: Path) -> List[Path]:
41 from pdf2image import convert_from_path
43 images = convert_from_path(str(source), dpi=dpi, fmt="png", output_folder=str(out_dir))
44 paths: List[Path] = []
45 for i, img in enumerate(images, start=1):
46 path = out_dir / f"page-{i:04d}.png"
47 img.save(path, "PNG")
48 paths.append(path)
49 return paths
52class CamelotCandidate(BaseModel):
53 """Serializable per-Camelot-table output. Crosses the spawn-process
54 boundary; only picklable primitives, no Camelot internals.
56 `cells` is the raw grid (rows of cell strings) straight from Camelot's
57 DataFrame; `markdown` is its rendering. The correspondence flow works
58 on `cells` so it can assemble, audit and fill at the grid level and
59 render markdown once at the end; the legacy classifier/unifier path
60 consumes `markdown`.
61 """
63 candidate_id: str = Field(description="Stable id assigned by the parent; flavor-page-idx")
64 flavor: Flavor
65 page: int = Field(ge=1)
66 bbox: Optional[Tuple[float, float, float, float]] = None
67 accuracy: float = Field(default=0.0, ge=0.0, le=100.0)
68 cells: List[List[str]] = Field(
69 default_factory=list, description="Camelot's raw cell grid (rows of strings)"
70 )
71 cell_boxes: List[List[Optional[Tuple[float, float, float, float]]]] = Field(
72 default_factory=list,
73 description=(
74 "Per-cell geometry aligned 1:1 with `cells`: same row/column shape, each entry the "
75 "Camelot cell box (x1, y1, x2, y2) in PDF points with a bottom-left page origin, or "
76 "None where Camelot exposed no box for that cell. Empty for candidates not built "
77 "directly from a Camelot grid (e.g. assembled or recovered candidates)."
78 ),
79 )
80 markdown: str = Field(description="Camelot's raw markdown for this table; rendering of `cells`")
83def grid_to_markdown(cells: List[List[str]]) -> str:
84 """Render a cell grid (rows of strings) to a markdown table. Row 0 is
85 the header; rows are padded to the widest row so the column count is
86 uniform. Single rendering point for the structured-grid flow.
87 """
88 if not cells:
89 return ""
90 width = max(len(r) for r in cells)
92 def render(row: List[str]) -> str:
93 padded = [str(c) for c in row] + [""] * (width - len(row))
94 return "| " + " | ".join(padded) + " |"
96 lines = [render(cells[0]), "| " + " | ".join("---" for _ in range(width)) + " |"]
97 lines.extend(render(row) for row in cells[1:])
98 return "\n".join(lines)
101def column_letter(index: int) -> str:
102 """0-based column index to a spreadsheet letter: 0->A, 25->Z, 26->AA."""
103 out = ""
104 index += 1
105 while index:
106 index, rem = divmod(index - 1, 26)
107 out = chr(ord("A") + rem) + out
108 return out
111def grid_to_addressed_markdown(cells: List[List[str]]) -> str:
112 """Render the grid with each non-empty cell's address printed inside it.
114 Every cell with text carries an inline tag — `[B3] 1,637` means data
115 column B (A=0), grid row 3 (1-based) — so the agent READS an address off
116 the label sitting next to the text; it never counts rows, columns, or
117 pipes. Blank cells stay blank. The tags are reference only — the agent's
118 corrected output must never contain them.
119 """
120 if not cells:
121 return ""
122 width = max(len(r) for r in cells)
124 def tag(r: int, c: int, text: str) -> str:
125 return f"[{column_letter(c)}{r}] {text}" if str(text).strip() else ""
127 def render_addressed(r: int, row: List[str]) -> str:
128 padded = [str(c) for c in row] + [""] * (width - len(row))
129 return "| " + " | ".join(tag(r, c, v) for c, v in enumerate(padded)) + " |"
131 lines = [render_addressed(1, cells[0]), "| " + " | ".join("---" for _ in range(width)) + " |"]
132 lines.extend(render_addressed(i, row) for i, row in enumerate(cells[1:], start=2))
133 return "\n".join(lines)
136def df_to_cells(df: DataFrame) -> List[List[str]]:
137 """Camelot DataFrame -> raw cell grid (rows of strings)."""
138 return [[str(c) for c in row] for row in df.values.tolist()]
141def cells_to_boxes(
142 raw_cells: Any, grid: List[List[str]]
143) -> List[List[Optional[Tuple[float, float, float, float]]]]:
144 """Per-cell boxes aligned to `grid`'s shape from Camelot's Cell objects.
146 `raw_cells` is Camelot's `table.cells` — rows of `Cell` objects, each with
147 x1/y1/x2/y2 in PDF points (bottom-left origin). `table.df` (hence `grid`) is
148 built from that same grid, so index i,j lines up; still, we clamp to `grid`'s
149 shape and fill None for any position Camelot did not cover, so `cell_boxes`
150 is always exactly the same shape as `cells`.
151 """
152 rows = list(raw_cells) if raw_cells is not None else []
153 boxes: List[List[Optional[Tuple[float, float, float, float]]]] = []
154 for i, grid_row in enumerate(grid):
155 raw_row = list(rows[i]) if i < len(rows) else []
156 out_row: List[Optional[Tuple[float, float, float, float]]] = []
157 for j in range(len(grid_row)):
158 cell = raw_row[j] if j < len(raw_row) else None
159 if cell is not None and all(hasattr(cell, a) for a in ("x1", "y1", "x2", "y2")):
160 out_row.append((float(cell.x1), float(cell.y1), float(cell.x2), float(cell.y2)))
161 else:
162 out_row.append(None)
163 boxes.append(out_row)
164 return boxes
167def df_to_markdown(df: DataFrame) -> str:
168 return grid_to_markdown(df_to_cells(df))
171def camelot_worker(source_str: str, flavor: Flavor) -> List[CamelotCandidate]:
172 """Top-level entrypoint for the spawned process. Imports camelot
173 fresh in the child so OpenCV initialization stays inside the worker
174 and never crosses a fork boundary.
175 """
176 import camelot
178 from quber.agents.completeness import page_words
179 from quber.core.extractors.camelot.tighten import tighten_cell_boxes
181 out: List[CamelotCandidate] = []
182 tables = camelot.read_pdf(source_str, pages="all", flavor=flavor) # pyright: ignore[reportPrivateImportUsage]
183 # Word rectangles per page, read once and shared by every table on it.
184 words_cache: dict[int, Tuple[float, List[Any]]] = {}
185 for idx, table in enumerate(tables):
186 raw_page = getattr(table, "page", 1)
187 page = int(raw_page) if raw_page is not None else 1
188 # camelot exposes the table box only as the private `_bbox`; read it
189 # via getattr so the access is not flagged as private use.
190 raw_bbox = getattr(table, "_bbox", None)
191 bbox = tuple(raw_bbox) if raw_bbox else None
192 report = getattr(table, "parsing_report", {}) or {}
193 accuracy = float(report.get("accuracy", 0.0) or 0.0)
194 cells = df_to_cells(table.df)
195 cell_boxes = cells_to_boxes(getattr(table, "cells", None), cells)
196 if page not in words_cache:
197 _, page_h, words = page_words(Path(source_str), page)
198 words_cache[page] = (page_h, words)
199 page_h, words = words_cache[page]
200 cell_boxes = tighten_cell_boxes(cells, cell_boxes, words, page_h)
201 out.append(
202 CamelotCandidate(
203 candidate_id=f"{flavor}-p{page}-{idx}",
204 flavor=flavor,
205 page=page,
206 bbox=bbox,
207 accuracy=accuracy,
208 cells=cells,
209 cell_boxes=cell_boxes,
210 markdown=grid_to_markdown(cells),
211 )
212 )
213 return out
216def run_camelot_flavors_parallel(
217 source: Path,
218 timeout_s: float = CAMELOT_FLAVOR_TIMEOUT_S,
219) -> List[CamelotCandidate]:
220 """Run lattice + stream concurrently in spawn-context subprocesses.
222 A flavor that times out or raises is logged ERROR and contributes
223 zero candidates; the other flavor's output is still returned. If
224 BOTH fail, returns an empty list — the caller decides what to do.
225 """
226 ctx = mp.get_context("spawn")
227 candidates: List[CamelotCandidate] = []
228 flavors: List[Flavor] = ["lattice", "stream"]
230 with ProcessPoolExecutor(max_workers=2, mp_context=ctx) as pool:
231 futures = {flavor: pool.submit(camelot_worker, str(source), flavor) for flavor in flavors}
232 for flavor, future in futures.items():
233 try:
234 candidates.extend(future.result(timeout=timeout_s))
235 except FutureTimeoutError:
236 logger.error(
237 "camelot {flavor} timed out after {timeout}s on {source}",
238 flavor=flavor,
239 timeout=timeout_s,
240 source=source.name,
241 )
242 future.cancel()
243 except Exception as exc:
244 logger.error(
245 "camelot {flavor} raised on {source}: {exc}",
246 flavor=flavor,
247 source=source.name,
248 exc=exc,
249 )
251 if not candidates:
252 logger.error("camelot produced 0 candidates from BOTH flavors on {}", source.name)
253 return candidates
256def is_content_empty(markdown: str) -> bool:
257 """A Camelot output is content-empty when every non-delimiter cell
258 is whitespace. Lattice's false-positive shells on whitespace-aligned
259 PDFs look like a 1-row grid of empty pipes; this catches them
260 before they reach the classifier or correction stage.
261 """
262 if not markdown.strip():
263 return True
264 for line in markdown.splitlines():
265 cells = [c.strip() for c in line.strip().strip("|").split("|")]
266 if any(c and not set(c) <= {"-", " "} for c in cells):
267 return False
268 return True