Coverage for src / quber / core / extractors / camelot / correspondence / matching.py: 36%

1""" 

2Chunk-to-table correspondence: assign Camelot chunks to detected tables 

3by 2D bbox overlap, deterministically — no LLM reads the cell values. 

4 

5Camelot's chunk boxes are precise PDF geometry; the detector's boxes are 

6rough, so each detected table takes the chunk that covers it most 

7(argmax), not every chunk that clips its band — that ignores the spill a 

8rough boundary causes. Because the overlap is 2D this covers both 

9convergence axes (vertically stacked tables Camelot ran together, and 

10side-by-side tables it fused column-wise). 

11""" 

12 

13from __future__ import annotations 

14 

15from dataclasses import dataclass 

16from typing import Dict, List, Optional, Tuple 

17 

18from quber.agents.detector import DetectedTable 

19from quber.core.extractors.camelot.acquire import CamelotCandidate 

20from quber.core.extractors.camelot.correspondence.geometry import ( 

21 camelot_bbox_to_norm, 

22 coverage_fraction, 

23) 

24 

25# A Camelot chunk is assigned to a detected table when their normalized 

26# rectangles overlap by at least this fraction of the chunk's area. The 

27# detector's boxes are rough, so the threshold is loose: it catches a 

28# chunk that sits mostly inside one band (~1.0), a tall chunk split across 

29# stacked tables (~1/N each), and a wide chunk fused across side-by-side 

30# tables (~0.4-0.5 each), while rejecting incidental slivers. 

31MATCH_FRACTION = 0.20 

32 

33 

34@dataclass(frozen=True) 

35class ChunkAssignment: 

36 """One flavor's argmax assignment of chunks to detected tables. 

37 

38 `best_chunk` maps a detected-table index to the index of the chunk 

39 that covers it most. `tables_for_chunk` is the inversion, for every 

40 chunk: a chunk that is the best match for two-or-more tables is a 

41 combined chunk (kept whole, flagged, never split). `orphans` are the 

42 chunks that own no table, in the order the chunks were given — the 

43 only safe source for completeness extension, since pulling a chunk 

44 that owns another table would pollute this one. 

45 """ 

46 

47 best_chunk: Dict[int, int] 

48 tables_for_chunk: Dict[int, List[int]] 

49 orphans: List[CamelotCandidate] 

50 

51 

52def assign_chunks( 

53 detected: List[DetectedTable], 

54 remaining: List[int], 

55 chunks: List[CamelotCandidate], 

56 page_w_pts: float, 

57 page_h_pts: float, 

58) -> ChunkAssignment: 

59 """Assign each unresolved detected table (indices in `remaining`) to 

60 its best-overlap chunk. A table whose best coverage falls below 

61 MATCH_FRACTION, or whose detector box is missing, gets no assignment. 

62 `chunks` should already be in page order so `orphans` comes out in 

63 page order too. 

64 """ 

65 chunk_norms: Dict[int, Optional[Tuple[float, float, float, float]]] = { 

66 ci: (camelot_bbox_to_norm(c.bbox, page_w_pts, page_h_pts) if c.bbox is not None else None) 

67 for ci, c in enumerate(chunks) 

68 } 

69 best_chunk: Dict[int, int] = {} 

70 for di in remaining: 

71 dbb = detected[di].bbox 

72 if dbb is None: 

73 continue 

74 best_ci, best_cov = None, 0.0 

75 for ci, cnorm in chunk_norms.items(): 

76 if cnorm is None: 

77 continue 

78 cov = coverage_fraction(cnorm, dbb) 

79 if cov > best_cov: 

80 best_cov, best_ci = cov, ci 

81 if best_ci is not None and best_cov >= MATCH_FRACTION: 

82 best_chunk[di] = best_ci 

83 

84 tables_for_chunk: Dict[int, List[int]] = {ci: [] for ci in range(len(chunks))} 

85 for di in sorted(best_chunk): 

86 tables_for_chunk[best_chunk[di]].append(di) 

87 

88 orphans = [chunks[ci] for ci in range(len(chunks)) if not tables_for_chunk[ci]] 

89 return ChunkAssignment(best_chunk=best_chunk, tables_for_chunk=tables_for_chunk, orphans=orphans) 

90 

91 

92def chunk_top(cand: CamelotCandidate) -> float: 

93 """Sort key for top-to-bottom page order. Camelot bbox is in PDF 

94 points with a bottom-left origin, so a larger top-y sits higher on 

95 the page; we negate it. Chunks without a bbox sort last. 

96 """ 

97 if cand.bbox is None: 

98 return float("inf") 

99 return -cand.bbox[3] 

100 

101 

102def assemble_cells(chunks: List[CamelotCandidate]) -> List[List[str]]: 

103 """Concatenate chunk cell grids in page order, values untouched. The 

104 structured grid is the canonical intermediate: assembly, the 

105 completeness audit and the fill all work on it, and markdown is 

106 rendered once at the output. Camelot stays the source of truth. 

107 """ 

108 rows: List[List[str]] = [] 

109 for c in chunks: 

110 rows.extend([list(r) for r in c.cells]) 

111 return rows