"""Compare every finishing option on identical model output. Phase 1 calls the vision+grid anatomy model ONCE per page (cached to JSON) to get the raw located boxes. Phase 2 applies each finishing option as pure post-processing on that same output, so differences are attributable to the finishing, not run-to-run model variation. Produces one annotated PDF per finishing (all documents concatenated) plus side-by-side panels for the pages that actually discriminate between options. Finishings: raw anatomy boxes, no post-processing declash + remove stacked-box overlap (anti-contamination) pad6 + declash + 6pt symmetric pad (anti-clip margin) pad12 + declash + 12pt symmetric pad snap + declash + line-walking title-snap (best edges, most fragile) """ from __future__ import annotations import asyncio import json import tempfile from pathlib import Path import fitz from PIL import Image from experiments.que245.approaches import ( anatomy_prompt, declash_stacked, pad_boxes, snap_titles, ) from experiments.que245.overlay import draw_boxes, render_page from quber.agents.completeness import page_words from quber.agents.grid_locator import ( GridFlagResult, LocatedTable, flags_to_located, overlay_grid, ) ROWS, COLS = 36, 12 RENDER = Path("experiments/que245/render") CACHE = Path("experiments/que245/results/raw_located.json") DOCS = [ Path("documents/BHE_991.pdf"), Path("documents/Q2FY25-Visa-Operational-Performance-Data-FINAL.pdf"), Path("documents/TMUS_Q225_991.pdf"), ] # Pages that distinguish the options (stacked tables, footnotes, prose seams). DISCRIMINATING = [ ("BHE_991", 1), ("BHE_991", 2), ("BHE_991", 7), ("Q2FY25-Visa-Operational-Performance-Data-FINAL", 1), ("Q2FY25-Visa-Operational-Performance-Data-FINAL", 3), ("TMUS_Q225_991", 8), ] FINISHINGS = ["raw", "declash", "pad6", "pad12", "snap"] COLORS = {"raw": (220, 0, 0)} def apply_finishing(name, located, words, pw, ph): if name == "raw": return located if name == "declash": return declash_stacked(located, words, pw, ph) if name == "pad6": return declash_stacked(pad_boxes(located, pw, ph, 6.0), words, pw, ph) if name == "pad12": return declash_stacked(pad_boxes(located, pw, ph, 12.0), words, pw, ph) if name == "snap": return snap_titles(declash_stacked(located, words, pw, ph), words, pw, ph) raise ValueError(name) async def phase1_collect(): if CACHE.exists(): print(f"using cached {CACHE}") return from experiments.que245.approaches import FlexGridLocator agent = FlexGridLocator(anatomy_prompt(ROWS, COLS), rows=ROWS, cols=COLS).agent from pydantic_ai import BinaryContent out = {} for pdf in DOCS: doc = fitz.open(str(pdf)) for pno in range(doc.page_count): with tempfile.TemporaryDirectory() as t: ip = Path(t) / "p.png" doc[pno].get_pixmap(dpi=200).save(str(ip)) gridded = Path(t) / "g.png" overlay_grid(ip, gridded, ROWS, COLS) img = BinaryContent(data=gridded.read_bytes(), media_type="image/png") res = await agent.run(["Flag the tables using the grid.", img]) pw, ph, words = page_words(pdf, pno + 1) located = flags_to_located(res.output.tables, words, pw, ph, ROWS, COLS) key = f"{pdf.stem}|{pno+1}" out[key] = [ {"ordinal": t.ordinal, "title": t.title, "region": list(t.region), "grid_rows": list(t.grid_rows), "grid_cols": list(t.grid_cols), "tightened": t.tightened} for t in located ] print(f" {key}: {len(located)} tables") CACHE.parent.mkdir(parents=True, exist_ok=True) CACHE.write_text(json.dumps(out, indent=1)) print(f"saved {CACHE}") def _located_from(dicts): return [LocatedTable(**d) for d in dicts] def phase2_render(): raw = json.loads(CACHE.read_text()) # one annotated PDF per finishing, all docs concatenated for name in FINISHINGS: master = fitz.open() for pdf in DOCS: src = fitz.open(str(pdf)) for pno in range(src.page_count): page = src[pno] w, h = page.rect.width, page.rect.height key = f"{pdf.stem}|{pno+1}" located = _located_from(raw.get(key, [])) _, _, words = page_words(pdf, pno + 1) fin = apply_finishing(name, located, words, w, h) col = (0.0, 0.32, 0.92) for t in fin: x0, y0, x1, y1 = t.region rect = fitz.Rect(min(x0, x1) * w, min(y0, y1) * h, max(x0, x1) * w, max(y0, y1) * h) page.draw_rect(rect, color=col, width=1.4) page.insert_text((rect.x0 + 2, max(8, rect.y0 - 3)), f"{t.ordinal}: {t.title}", fontsize=6, color=col) master.insert_pdf(src) outp = RENDER / f"finish_{name}.pdf" master.save(str(outp)) print(f"saved {outp}") # Per discriminating page: stack the 5 finishings vertically, each cropped # to the table band (+margin) at high DPI so the few-point edge differences # are visible. from PIL import ImageDraw, ImageFont try: font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 18) except Exception: font = ImageFont.load_default() dpi = 150 scale = dpi / 72.0 for stem, pno in DISCRIMINATING: pdf = next(d for d in DOCS if d.stem == stem) _, _, words = page_words(pdf, pno) pg = fitz.open(str(pdf))[pno - 1] w, h = pg.rect.width, pg.rect.height key = f"{stem}|{pno}" located = _located_from(raw.get(key, [])) fins = {name: apply_finishing(name, located, words, w, h) for name in FINISHINGS} all_boxes = [b.region for f in fins.values() for b in f] if all_boxes: y0 = max(0.0, min(min(b[1], b[3]) for b in all_boxes) - 0.06) y1 = min(1.0, max(max(b[1], b[3]) for b in all_boxes) + 0.04) else: y0, y1 = 0.0, 0.4 top_px, bot_px = int(y0 * h * scale), int(y1 * h * scale) crops = [] for name in FINISHINGS: im = render_page(pdf, pno, dpi=dpi) im = draw_boxes(im, [(name, [t.region for t in fins[name]], (210, 20, 20))]) crops.append((name, im.crop((0, top_px, im.width, bot_px)))) cw = crops[0][1].width ch = crops[0][1].height sheet = Image.new("RGB", (cw, (ch + 26) * len(crops)), (255, 255, 255)) d = ImageDraw.Draw(sheet) for i, (name, im) in enumerate(crops): yo = i * (ch + 26) d.text((4, yo + 4), f"{name}", fill=(0, 0, 160), font=font) sheet.paste(im, (0, yo + 26)) out = RENDER / f"panel_{stem[:14]}_p{pno}.png" sheet.save(out) print(f"saved {out}") async def main(): await phase1_collect() phase2_render() if __name__ == "__main__": asyncio.run(main())