Coverage for src / quber / core / extractors / camelot / recapture.py: 90%

1"""Detect and repair Camelot capture drops before structure correction. 

2 

3Camelot occasionally drops printed values during cell assignment even though 

4its text extraction captured them: a totals row typeset on a raised baseline 

5splits into two row bands at the default grouping tolerance, and the numbers' 

6vertical center lands exactly on the band boundary, failing the strict 

7containment test on both sides — the row's label and dollar signs survive, 

8its numbers vanish from the grid. 

9 

10The repair loop keeps authority deterministic and puts the model only where 

11judgement is needed: 

12 

131. DETECT (deterministic): numeric tokens present in the table region's text 

14 layer but absent from the grid. Same numeric normalization as the 

15 correction-stage grounding guard. 

162. RECOMMEND (agent): the capture advisor sees the cropped table image, the 

17 extracted grid, and the dropped tokens with their printed lines, and 

18 proposes one bounded retry adjustment. 

193. RETRY + ACCEPT (deterministic): re-run the region pass with the advised 

20 knobs. The retry replaces the original ONLY if it recovers every missing 

21 numeric and loses none the original had. Anything else keeps the original 

22 grid, so a wrong recommendation costs one extra Camelot pass and nothing 

23 else. 

24 

25Every detection and every verdict is logged — a capture drop on a financial 

26document must never be silent. 

27""" 

28 

29from __future__ import annotations 

30 

31import asyncio 

32from pathlib import Path 

33from typing import List, Optional, Set 

34 

35from loguru import logger 

36 

37from quber.agents.capture_advisor import CaptureAdvisor 

38from quber.agents.completeness import page_words 

39from quber.core.extractors.camelot.acquire import CamelotCandidate, grid_to_markdown 

40from quber.core.extractors.camelot.correspondence.correction import numeric_keys 

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

42 WordBox, 

43 bbox_to_top_left, 

44 crop_region_png, 

45 table_crop_box, 

46) 

47from quber.core.extractors.camelot.correspondence.recovery import camelot_targeted 

48 

49# How far outside the Camelot bbox (PDF points) a word may sit and still count 

50# as inside the table for drop detection. Tight on purpose: a generous pad 

51# would pull a neighbouring table's numbers into the missing set and flag 

52# healthy captures. 

53DETECT_PAD_PTS = 2.0 

54 

55 

56def dropped_numeric_keys(cells: List[List[str]], region_words: List[WordBox]) -> Set[str]: 

57 """Numeric tokens in the region's text layer that appear in no grid cell.""" 

58 layer = numeric_keys(" ".join(w[4] for w in region_words)) 

59 grid = numeric_keys(" ".join(c for row in cells for c in row)) 

60 return layer - grid 

61 

62 

63def printed_lines(region_words: List[WordBox], missing: Set[str]) -> List[str]: 

64 """For each missing token, the full printed line it sits on — the evidence 

65 the advisor grounds its diagnosis in.""" 

66 lines: List[str] = [] 

67 for key in sorted(missing): 

68 for w in region_words: 

69 if numeric_keys(w[4]) == {key}: 

70 line = " ".join(x[4] for x in region_words if abs(x[1] - w[1]) < 3) 

71 lines.append(f"'{key}' on printed line: {line}") 

72 break 

73 else: 

74 lines.append(f"'{key}' (word not isolated on a line)") 

75 return lines 

76 

77 

78async def repair_capture( 

79 cand: CamelotCandidate, 

80 source: str, 

81 page: int, 

82 page_image: Optional[Path], 

83 dpi: int, 

84 ordinal: int, 

85 advisor: Optional[CaptureAdvisor], 

86) -> CamelotCandidate: 

87 """Return `cand`, or an advised retry that provably captures more. 

88 

89 No-op when the advisor is off, the candidate has no bbox, or no drop is 

90 detected. The retry is scoped to the candidate's own bbox — the exact 

91 frame the drop was measured in. 

92 """ 

93 if advisor is None or cand.bbox is None or page_image is None: 

94 return cand 

95 _pw, page_h, words = await asyncio.to_thread(page_words, Path(source), page) 

96 left, top, right, bottom = bbox_to_top_left(cand.bbox, page_h, pad=DETECT_PAD_PTS) 

97 region_words = [w for w in words if w[0] >= left and w[2] <= right and w[1] >= top and w[3] <= bottom] 

98 missing = dropped_numeric_keys(cand.cells, region_words) 

99 if not missing: 

100 return cand 

101 logger.warning( 

102 "page {}: capture drop detected — {} numeric token(s) in the table's text layer " 

103 "missing from the grid: {}", 

104 page, 

105 len(missing), 

106 sorted(missing), 

107 ) 

108 

109 crop = await asyncio.to_thread(crop_region_png, page_image, table_crop_box(cand.bbox, page_h), dpi) 

110 advice = await advisor.recommend( 

111 crop, grid_to_markdown(cand.cells), printed_lines(region_words, missing), page 

112 ) 

113 if advice is None: 

114 return cand 

115 logger.info( 

116 "page {}: capture advisor recommends flavor={} row_tol={} column_tol={}: {}", 

117 page, 

118 advice.flavor, 

119 advice.row_tol, 

120 advice.column_tol, 

121 advice.diagnosis, 

122 ) 

123 

124 x1, y1, x2, y2 = cand.bbox 

125 area = f"{min(x1, x2):.1f},{max(y1, y2):.1f},{max(x1, x2):.1f},{min(y1, y2):.1f}" 

126 try: 

127 retry = await asyncio.to_thread( 

128 camelot_targeted, 

129 source, 

130 page, 

131 area, 

132 ordinal, 

133 advice.flavor, 

134 advice.row_tol, 

135 advice.column_tol, 

136 ) 

137 except Exception as exc: 

138 logger.warning("page {}: capture-repair retry failed; keeping original grid: {}", page, exc) 

139 return cand 

140 if retry is None: 

141 logger.warning("page {}: capture-repair retry found no grid; keeping original", page) 

142 return cand 

143 

144 original = numeric_keys(" ".join(c for row in cand.cells for c in row)) 

145 retried = numeric_keys(" ".join(c for row in retry.cells for c in row)) 

146 recovered = missing & retried 

147 lost = original - retried 

148 if recovered == missing and not lost: 

149 logger.info( 

150 "page {}: capture repair ACCEPTED — recovered {} token(s), lost none", 

151 page, 

152 len(recovered), 

153 ) 

154 return retry 

155 logger.warning( 

156 "page {}: capture repair REJECTED (recovered {}/{}, lost {}); keeping original grid", 

157 page, 

158 len(recovered), 

159 len(missing), 

160 len(lost), 

161 ) 

162 return cand