Coverage for src / quber / playground / backfill_boxes.py: 0%

1"""In-place tightening of stored table-cell grounding boxes. 

2 

3Documents already in the library were extracted before cell rectangles were 

4tightened at acquisition, so their stored boxes are still structural grid 

5geometry — the column's x-extent crossed with the row band. Re-extracting 

6would cost LLM passes and re-ingesting would wipe curation, so this backfill 

7corrects the one thing that is wrong: for every tableCell grounding whose 

8text can be found in the staged PDF's text layer near the stored box, the 

9row's bbox is updated to the printed extent. No other column is touched; 

10unmatched cells keep their boxes. 

11 

12Run: 

13 uv run python -m quber.playground.backfill_boxes [doc_key ...] 

14 

15With no arguments every document in the library is processed. A document 

16whose staged PDF is missing is reported and skipped, never silently. 

17""" 

18 

19from __future__ import annotations 

20 

21import argparse 

22import json 

23from pathlib import Path 

24from typing import Any, Dict, List, Tuple 

25 

26import psycopg 

27from psycopg.types.json import Jsonb 

28 

29from quber.agents.completeness import page_words 

30from quber.core.extractors.camelot.tighten import ink_extent, tighten_box 

31from quber.playground import db 

32from quber.settings import get_settings 

33 

34 

35def backfill_document(conn: psycopg.Connection, doc_id: int, pdf_path: Path) -> Tuple[int, int]: 

36 """Tighten one document's cell boxes in place. Returns (tightened, kept).""" 

37 rows = conn.execute( 

38 """SELECT ref_id, page, bbox, cell_text FROM ade_playground.groundings 

39 WHERE document_id = %s AND ref_type = 'tableCell' 

40 AND bbox IS NOT NULL AND coalesce(cell_text, '') <> ''""", 

41 (doc_id,), 

42 ).fetchall() 

43 by_page: Dict[int, List[Tuple[Any, ...]]] = {} 

44 for r in rows: 

45 by_page.setdefault(r[1], []).append(r) 

46 

47 tightened = kept = 0 

48 for page0, page_rows in by_page.items(): 

49 try: 

50 w, h, words = page_words(pdf_path, page0 + 1) 

51 except IndexError: 

52 kept += len(page_rows) 

53 continue 

54 for ref_id, _, bbox, cell_text in page_rows: 

55 box = bbox if isinstance(bbox, dict) else json.loads(bbox) 

56 points = (box["left"] * w, box["top"] * h, box["right"] * w, box["bottom"] * h) 

57 tight = tighten_box(points, cell_text, words) or ink_extent(points, words) 

58 if tight is None: 

59 kept += 1 

60 continue 

61 x1, y1, x2, y2 = tight 

62 conn.execute( 

63 """UPDATE ade_playground.groundings SET bbox = %s 

64 WHERE document_id = %s AND ref_id = %s""", 

65 ( 

66 Jsonb({"left": x1 / w, "top": y1 / h, "right": x2 / w, "bottom": y2 / h}), 

67 doc_id, 

68 ref_id, 

69 ), 

70 ) 

71 tightened += 1 

72 return tightened, kept 

73 

74 

75def main() -> int: 

76 parser = argparse.ArgumentParser(description="Tighten stored table-cell grounding boxes in place") 

77 parser.add_argument("doc_keys", nargs="*", help="doc_key values to backfill; default all") 

78 args = parser.parse_args() 

79 

80 data_dir = get_settings().playground.data_dir 

81 total_tightened = total_kept = 0 

82 missing: List[str] = [] 

83 with db.connect() as conn: 

84 docs = conn.execute( 

85 "SELECT id, doc_key, coalesce(title, filename) FROM ade_playground.documents ORDER BY id" 

86 ).fetchall() 

87 for doc_id, doc_key, label in docs: 

88 if args.doc_keys and doc_key not in args.doc_keys: 

89 continue 

90 pdf_path = data_dir / f"{doc_key}.pdf" 

91 if not pdf_path.exists(): 

92 missing.append(f"{doc_key} ({label})") 

93 continue 

94 tightened, kept = backfill_document(conn, doc_id, pdf_path) 

95 total_tightened += tightened 

96 total_kept += kept 

97 print(f" {doc_key} tightened {tightened:>5} kept {kept:>4} {label}") 

98 

99 print(f"\ntotal tightened {total_tightened}, kept {total_kept}") 

100 if missing: 

101 print(f"\nSKIPPED — staged PDF missing for {len(missing)} document(s):") 

102 for m in missing: 

103 print(f" {m}") 

104 return 0 

105 

106 

107if __name__ == "__main__": 

108 raise SystemExit(main())