Coverage for src / quber / core / ocr / well.py: 58%

1"""OCR grounding well: page text read from a rendered image, with provenance. 

2 

3A scanned page or an image-bound figure has no native text layer, so a value 

4extracted from it has nothing to ground against. The well fills that gap. It 

5runs RapidOCR over a rendered page image and keeps every recognized fragment 

6with its bounding box and confidence. For a render of an original PDF page the 

7box is stored in PDF points, so a fragment grounds a value from a scanned 

8region the same way a native text cell grounds a table cell. 

9 

10The well is evidence, never a correction target. Fragments are stored exactly 

11as the engine read them; anything that interprets or repairs them happens in 

12the consumer, against the fragment ids, so the trace from a value back to the 

13pixels survives every downstream step. 

14""" 

15 

16from __future__ import annotations 

17 

18import subprocess 

19import tempfile 

20from pathlib import Path 

21from typing import Any 

22 

23from pydantic import BaseModel 

24 

25POINTS_PER_INCH = 72.0 

26 

27 

28class OcrFragment(BaseModel): 

29 """One recognized text fragment and where it sits on the page.""" 

30 

31 id: str 

32 page: int 

33 text: str 

34 #: Left, top, right, bottom. PDF points when the well was built from a PDF 

35 #: render; pixels when built from a standalone image (dpi is None then). 

36 bbox: tuple[float, float, float, float] 

37 confidence: float 

38 

39 

40class PageWell(BaseModel): 

41 """Every fragment read from one page image.""" 

42 

43 page: int 

44 image: str 

45 dpi: int | None 

46 width: float 

47 height: float 

48 fragments: list[OcrFragment] 

49 

50 def fragments_in_box( 

51 self, 

52 bbox: tuple[float, float, float, float], 

53 min_overlap: float = 0.5, 

54 ) -> list[OcrFragment]: 

55 """Fragments whose area overlaps the region by at least min_overlap.""" 

56 left, top, right, bottom = bbox 

57 hits = [] 

58 for frag in self.fragments: 

59 fl, ft, fr, fb = frag.bbox 

60 inter_w = max(0.0, min(fr, right) - max(fl, left)) 

61 inter_h = max(0.0, min(fb, bottom) - max(ft, top)) 

62 area = max((fr - fl) * (fb - ft), 1e-6) 

63 if (inter_w * inter_h) / area >= min_overlap: 

64 hits.append(frag) 

65 return hits 

66 

67 

68class DocumentWell(BaseModel): 

69 """The wells for every page read from one source document.""" 

70 

71 source: str 

72 engine: str 

73 pages: list[PageWell] 

74 

75 def page(self, page_no: int) -> PageWell: 

76 for well in self.pages: 

77 if well.page == page_no: 

78 return well 

79 raise KeyError(f"no well for page {page_no}") 

80 

81 

82def make_engine(device_id: int = 0) -> Any: 

83 """RapidOCR configured the way the tuned-financial preset runs it. 

84 

85 docling propagates the CUDA device to RapidOCR's paddle and torch engines 

86 but not the default onnxruntime engine, so the CUDA execution provider is 

87 named explicitly here, matching TunedFinancialParser. 

88 """ 

89 from rapidocr import RapidOCR 

90 

91 return RapidOCR( 

92 params={ 

93 "EngineConfig.onnxruntime.use_cuda": True, 

94 "EngineConfig.onnxruntime.cuda_ep_cfg.device_id": device_id, 

95 } 

96 ) 

97 

98 

99def read_image( 

100 image_path: Path, 

101 page_no: int, 

102 dpi: int | None = None, 

103 engine: Any = None, 

104) -> PageWell: 

105 """Read one page image into a well. 

106 

107 When dpi is given the fragment boxes are converted from pixels to PDF 

108 points, so they land in the coordinate frame of the original page. 

109 """ 

110 from PIL import Image 

111 

112 if engine is None: 

113 engine = make_engine() 

114 with Image.open(image_path) as img: 

115 width_px, height_px = img.size 

116 result = engine(str(image_path)) 

117 scale = POINTS_PER_INCH / dpi if dpi else 1.0 

118 

119 fragments: list[OcrFragment] = [] 

120 boxes = getattr(result, "boxes", None) 

121 txts = getattr(result, "txts", None) or [] 

122 scores = getattr(result, "scores", None) or [] 

123 if boxes is not None: 

124 for i, (box, text, score) in enumerate(zip(boxes, txts, scores, strict=False)): 

125 xs = [float(p[0]) for p in box] 

126 ys = [float(p[1]) for p in box] 

127 fragments.append( 

128 OcrFragment( 

129 id=f"p{page_no}.f{i}", 

130 page=page_no, 

131 text=str(text), 

132 bbox=( 

133 min(xs) * scale, 

134 min(ys) * scale, 

135 max(xs) * scale, 

136 max(ys) * scale, 

137 ), 

138 confidence=float(score), 

139 ) 

140 ) 

141 return PageWell( 

142 page=page_no, 

143 image=str(image_path), 

144 dpi=dpi, 

145 width=width_px * scale, 

146 height=height_px * scale, 

147 fragments=fragments, 

148 ) 

149 

150 

151def build_well( 

152 pdf: Path, 

153 pages: list[int], 

154 dpi: int = 300, 

155 engine: Any = None, 

156) -> DocumentWell: 

157 """Render the named pages of a PDF and read each into a well.""" 

158 if engine is None: 

159 engine = make_engine() 

160 page_wells: list[PageWell] = [] 

161 with tempfile.TemporaryDirectory() as tmp: 

162 for page_no in pages: 

163 prefix = Path(tmp) / f"page{page_no:04d}" 

164 subprocess.run( 

165 [ 

166 "pdftoppm", 

167 "-r", 

168 str(dpi), 

169 "-png", 

170 "-f", 

171 str(page_no), 

172 "-l", 

173 str(page_no), 

174 str(pdf), 

175 str(prefix), 

176 ], 

177 check=True, 

178 capture_output=True, 

179 ) 

180 rendered = sorted(Path(tmp).glob(f"page{page_no:04d}*.png")) 

181 if not rendered: 

182 raise FileNotFoundError(f"pdftoppm produced no image for page {page_no}") 

183 page_wells.append(read_image(rendered[0], page_no, dpi=dpi, engine=engine)) 

184 return DocumentWell(source=str(pdf), engine="rapidocr", pages=page_wells)