"""QUE-270: draw table (SoM) and footnote regions as distinct boxes on a corpus. Goal ---- Expand the region identifiers so every page carries TWO kinds of box: the grid-locator's table regions (blue) and a dedicated footnote finder's grounded footnote regions (orange). Together they delimit the combined table+footnote target a downstream extraction agent would read. This script produces an annotated PDF per document plus a JSON index, run across the whole real corpus, so the boxes can be eyeballed at scale rather than reasoned about on a few pages. It reuses, unchanged: - PydanticAIGridLocator for table regions (the production locator); - the dedicated footnote-finder prompt/schema and the marker-grounded region logic validated in footnote_finder_probe.py; - PyMuPDF box drawing, mirroring src/quber/review/annotated.py. Run --- uv run python experiments/que270/annotate_regions.py # whole corpus uv run python experiments/que270/annotate_regions.py ... # specific docs """ from __future__ import annotations import asyncio import json import sys import tempfile from pathlib import Path from typing import List, Optional, Tuple import fitz from loguru import logger sys.path.insert(0, str(Path(__file__).resolve().parent)) import footnote_finder_probe as ffp # noqa: E402 from quber.agents.completeness import page_words # noqa: E402 from quber.agents.grid_locator import ( # noqa: E402 BOX_PAD_PTS, DEFAULT_COLS, DEFAULT_ROWS, GridFlag, LocatedTable, PydanticAIGridLocator, declash_stacked, flags_to_located, overlay_grid, pad_boxes, ) from quber.files.cache import init_s3_cache, resolve_document # noqa: E402 REPO = Path(__file__).resolve().parents[2] DOCS = REPO / ".cache/s3/qubera-docs" DPI = 200 PAGE_CONCURRENCY = 8 PAGE_CAP = 60 # annotate at most this many pages per document # Footnote grounding: "marker" = the bespoke marker/prose heuristic from the # probe; "prod" = the SAME geometry chain the tables use (tighten_region -> # pad_boxes -> declash against the table boxes). Set via --grounding. GROUNDING = "prod" OUT_DIR = Path(__file__).resolve().parent / "out" / "annotated" # reset per mode in main() def _render_capped(pdf: Path, out_dir: Path, dpi: int, cap: int) -> List[Path]: """Render the first `cap` pages of `pdf` to PNGs via PyMuPDF (honors the cap cheaply; a 500-page filing renders only its first `cap` pages).""" doc = fitz.open(str(pdf)) n = min(doc.page_count, cap) zoom = dpi / 72.0 mat = fitz.Matrix(zoom, zoom) paths: List[Path] = [] for i in range(n): pix = doc[i].get_pixmap(matrix=mat) path = out_dir / f"page-{i + 1:04d}.png" pix.save(str(path)) paths.append(path) doc.close() return paths TABLE_RGB = (0.0, 0.32, 0.92) # blue FOOTNOTE_RGB = (0.95, 0.45, 0.0) # orange def corpus() -> List[Path]: """Every real (non-synthetic) PDF under the qubera-docs cache, sorted.""" return sorted(p for p in DOCS.rglob("*.pdf")) async def _page_regions( locator: PydanticAIGridLocator, fn_agent: object, pdf: Path, page: int, raw_img: Path, ) -> Tuple[int, List[Tuple[int, str, Tuple[float, float, float, float]]], List[Tuple[List[str], Tuple[float, float, float, float]]]]: """Locate tables and footnote blocks on one page; return both as norm boxes. Tables come from the production locator. Footnote blocks come from the dedicated finder, each grounded to a precise region off the PDF text layer. Footnotes are only sought on pages that have at least one table. """ from pydantic_ai import BinaryContent located = await locator.locate(raw_img, pdf, page) tables = [(t.ordinal, t.title, t.region) for t in located] footnotes: List[Tuple[List[str], Tuple[float, float, float, float]]] = [] if located: page_w, page_h, words = await asyncio.to_thread(page_words, pdf, page) with tempfile.TemporaryDirectory(prefix="que270-fn-") as tmp: gridded = Path(tmp) / "gridded.png" await asyncio.to_thread(overlay_grid, raw_img, gridded, DEFAULT_ROWS, DEFAULT_COLS) image = BinaryContent(data=gridded.read_bytes(), media_type="image/png") try: result = await fn_agent.run(["Report using the grid.", image]) # type: ignore[attr-defined] blocks = result.output.blocks except Exception as exc: logger.warning("footnote finder failed page={} of {}: {}", page, pdf.name, exc) blocks = [] footnotes = _ground_footnotes(blocks, located, words, page_w, page_h) return page, tables, footnotes def _ground_footnotes( blocks: list, located_tables: List[LocatedTable], words: object, page_w: float, page_h: float, ) -> List[Tuple[List[str], Tuple[float, float, float, float]]]: """Ground footnote-block flags to precise regions. "prod": feed the flags through the production geometry chain the tables use - flags_to_located (which runs tighten_region against the text layer), pad_boxes, then declash_stacked over the combined table+footnote set so a footnote box is pulled off the table directly above it. "marker": the probe's bespoke marker/prose heuristic. """ if GROUNDING == "prod": flags = [ GridFlag( ordinal=i + 1, title=" ".join(b.markers), row_start=b.row_start, row_end=b.row_end, col_start=b.col_start, col_end=b.col_end, ) for i, b in enumerate(blocks) ] fn_located = flags_to_located(flags, words, page_w, page_h, DEFAULT_ROWS, DEFAULT_COLS) fn_located = pad_boxes(fn_located, page_w, page_h, BOX_PAD_PTS) combined = declash_stacked(list(located_tables) + list(fn_located), words, page_w, page_h) grounded = combined[len(located_tables):] return [(b.markers, g.region) for b, g in zip(blocks, grounded)] out: List[Tuple[List[str], Tuple[float, float, float, float]]] = [] for blk in blocks: box, _text, _marker = ffp._ground_block(blk, words, page_w, page_h, DEFAULT_ROWS, DEFAULT_COLS) if box is not None: out.append((blk.markers, box)) return out def _draw( doc: fitz.Document, page_no: int, tables: List[Tuple[int, str, Tuple[float, float, float, float]]], footnotes: List[Tuple[List[str], Tuple[float, float, float, float]]], ) -> None: page = doc[page_no - 1] w, h = page.rect.width, page.rect.height for ordinal, title, region in tables: rect = _rect(region, w, h) page.draw_rect(rect, color=TABLE_RGB, width=1.4) label = f"T{page_no}.{ordinal}: {title}" if title else f"T{page_no}.{ordinal}" page.insert_text((rect.x0 + 2, max(8.0, rect.y0 - 3)), label[:80], fontsize=6, color=TABLE_RGB) for markers, region in footnotes: rect = _rect(region, w, h) page.draw_rect(rect, color=FOOTNOTE_RGB, width=1.4, dashes="[3 2] 0") label = "fn " + " ".join(markers) if markers else "fn" page.insert_text((rect.x0 + 2, min(h - 2, rect.y1 + 7)), label[:60], fontsize=6, color=FOOTNOTE_RGB) def _rect(region: Tuple[float, float, float, float], w: float, h: float) -> fitz.Rect: x0, y0, x1, y1 = region return fitz.Rect(min(x0, x1) * w, min(y0, y1) * h, max(x0, x1) * w, max(y0, y1) * h) async def annotate_document(locator: PydanticAIGridLocator, fn_agent: object, pdf: Path) -> dict: logger.info("annotating {}", pdf.name) doc = fitz.open(str(pdf)) n_pages = min(doc.page_count, PAGE_CAP) with tempfile.TemporaryDirectory(prefix="que270-pages-") as tmp: images = await asyncio.to_thread(_render_capped, pdf, Path(tmp), DPI, PAGE_CAP) sem = asyncio.Semaphore(PAGE_CONCURRENCY) async def one(page: int) -> Tuple[int, list, list]: async with sem: return await _page_regions(locator, fn_agent, pdf, page, images[page - 1]) results = await asyncio.gather(*(one(p) for p in range(1, n_pages + 1)), return_exceptions=True) n_tables = n_fn = 0 fn_pages: List[int] = [] for res in results: if isinstance(res, BaseException): logger.error("page failed in {}: {}", pdf.name, res) continue page, tables, footnotes = res _draw(doc, page, tables, footnotes) n_tables += len(tables) n_fn += len(footnotes) if footnotes: fn_pages.append(page) OUT_DIR.mkdir(parents=True, exist_ok=True) out_pdf = OUT_DIR / f"{pdf.parent.name}__{pdf.stem}.annotated.pdf" doc.save(str(out_pdf)) doc.close() summary = { "doc": pdf.name, "out": str(out_pdf.relative_to(REPO)), "pages": n_pages, "tables": n_tables, "footnote_blocks": n_fn, "footnote_pages": fn_pages, } logger.info( "{}: {} pages, {} tables, {} footnote blocks on pages {}", pdf.name, n_pages, n_tables, n_fn, fn_pages, ) return summary def _resolve(target: str) -> Optional[Path]: """Local path passes through; an s3:// URI materializes through the cache.""" if target.startswith("s3://"): return resolve_document(target) return Path(target) async def main(targets: List[str]) -> None: init_s3_cache() locator = PydanticAIGridLocator() fn_agent = ffp._build_agent( ffp.footnote_only_prompt(DEFAULT_ROWS, DEFAULT_COLS), ffp.FootnoteFinderResult ) index = [] for i, target in enumerate(targets, 1): logger.info("=== doc {}/{}: {} ===", i, len(targets), target) try: pdf = _resolve(target) index.append(await annotate_document(locator, fn_agent, pdf)) except Exception as exc: logger.error("document failed {}: {}", target, exc) index.append({"doc": target, "error": str(exc)}) OUT_DIR.mkdir(parents=True, exist_ok=True) (OUT_DIR / "index.json").write_text(json.dumps(index, indent=2)) # checkpoint each doc print(json.dumps(index, indent=2)) print(f"\nAnnotated PDFs + index.json under {OUT_DIR}") if __name__ == "__main__": argv = sys.argv[1:] if "--grounding" in argv: gi = argv.index("--grounding") GROUNDING = argv[gi + 1] del argv[gi : gi + 2] OUT_DIR = Path(__file__).resolve().parent / "out" / f"annotated_{GROUNDING}" targets = argv if argv else [str(p) for p in corpus()] asyncio.run(main(targets))