1"""OCR of images and scanned PDF pages with Tesseract: text, lines and words with confidence and boxes."""
2from __future__ import annotations
3
4import re
5from typing import Iterator
6
7import pypdfium2 as pdfium
8import pytesseract
9from PIL import Image, ImageOps, UnidentifiedImageError
10
11try:
12 import pillow_heif
13 pillow_heif.register_heif_opener()
14except ImportError:
15 pass
16
17Image.MAX_IMAGE_PIXELS = 200_000_000
18PDF_DPI = 300
19MIN_SIDE = 600
20PSM = {'auto': 3, 'single-block': 6, 'single-line': 7, 'sparse': 11}
21
22
23class OcrError(Exception):
24 pass
25
26
27def detect_kind(path: str) -> str:
28 with open(path, 'rb') as f:
29 head = f.read(16)
30 if head.startswith(b'%PDF'):
31 return 'pdf'
32 try:
33 with Image.open(path) as im:
34 im.verify()
35 except (UnidentifiedImageError, OSError, SyntaxError, ValueError) as e:
36 raise OcrError('the file is not a supported image or PDF') from e
37 return 'image'
38
39
40def parse_pages(spec: str, total: int) -> list[int]:
41 """'1-3,5,8-' -> zero-based page indexes. Empty = all."""
42 spec = (spec or '').replace(' ', '')
43 if not spec:
44 return list(range(total))
45 out = []
46 for part in spec.split(','):
47 m = re.fullmatch(r'(\d*)-(\d*)|(\d+)', part)
48 if not m:
49 raise OcrError(f'bad page range "{part}"')
50 if m.group(3):
51 a = b = int(m.group(3))
52 else:
53 a, b = int(m.group(1) or 1), int(m.group(2) or total)
54 out.extend(i - 1 for i in range(max(a, 1), min(b, total) + 1) if i - 1 not in out)
55 return out
56
57
58def iter_pages(path: str, kind: str, pages: str, limit: int | None) -> Iterator[tuple[int, int, Image.Image]]:
59 """Yields (page number from 1, total pages, image)."""
60 if kind == 'pdf':
61 try:
62 pdf = pdfium.PdfDocument(path)
63 except pdfium.PdfiumError as e:
64 raise OcrError('the PDF is damaged or password-protected') from e
65 try:
66 total = len(pdf)
67 for n, i in enumerate(parse_pages(pages, total)):
68 if limit is not None and n >= limit:
69 return
70 yield i + 1, total, pdf[i].render(scale=PDF_DPI / 72).to_pil()
71 finally:
72 pdf.close()
73 return
74 with Image.open(path) as im:
75 total = getattr(im, 'n_frames', 1)
76 for n, i in enumerate(parse_pages(pages, total)):
77 if limit is not None and n >= limit:
78 return
79 im.seek(i)
80 yield i + 1, total, ImageOps.exif_transpose(im.copy())
81
82
83def prepare(img: Image.Image) -> Image.Image:
84 if img.mode in ('RGBA', 'LA', 'P'):
85 bg = Image.new('RGB', img.size, 'white')
86 rgba = img.convert('RGBA')
87 bg.paste(rgba, mask=rgba.split()[-1])
88 img = bg
89 img = img.convert('L')
90 if min(img.size) < MIN_SIDE:
91 img = img.resize((img.width * 2, img.height * 2), Image.LANCZOS)
92 return img
93
94
95def rotate_upright(img: Image.Image) -> tuple[Image.Image, int]:
96 try:
97 osd = pytesseract.image_to_osd(img, output_type=pytesseract.Output.DICT)
98 except pytesseract.TesseractError:
99 return img, 0
100 angle = int(osd.get('rotate', 0))
101 return (img.rotate(-angle, expand=True) if angle else img), angle
102
103
104def ocr(img: Image.Image, langs: str, psm: str = 'auto', words: bool = False) -> dict:
105 scale = 1.0
106 prepared = prepare(img)
107 if prepared.width != img.width:
108 scale = img.width / prepared.width
109 try:
110 data = pytesseract.image_to_data(prepared, lang=langs, config=f'--psm {PSM.get(psm, 3)}',
111 output_type=pytesseract.Output.DICT)
112 except pytesseract.TesseractError as e:
113 raise OcrError(f'OCR failed: {str(e)[:200]}') from e
114 lines, word_list, blocks = {}, [], []
115 for i, text in enumerate(data['text']):
116 text = (text or '').strip()
117 conf = float(data['conf'][i])
118 if not text or conf < 0:
119 continue
120 box = [round(v * scale) for v in (data['left'][i], data['top'][i], data['width'][i], data['height'][i])]
121 key = (data['block_num'][i], data['par_num'][i], data['line_num'][i])
122 line = lines.setdefault(key, {'words': [], 'confs': [], 'boxes': []})
123 line['words'].append(text)
124 line['confs'].append(conf)
125 line['boxes'].append(box)
126 if words:
127 word_list.append({'text': text, 'confidence': round(conf, 1), 'bbox': box})
128 out_lines, paragraphs, last_par = [], [], None
129 for key in sorted(lines):
130 line = lines[key]
131 x0 = min(b[0] for b in line['boxes'])
132 y0 = min(b[1] for b in line['boxes'])
133 x1 = max(b[0] + b[2] for b in line['boxes'])
134 y1 = max(b[1] + b[3] for b in line['boxes'])
135 text = ' '.join(line['words'])
136 out_lines.append({'text': text, 'confidence': round(sum(line['confs']) / len(line['confs']), 1),
137 'bbox': [x0, y0, x1 - x0, y1 - y0]})
138 if key[:2] != last_par:
139 paragraphs.append([])
140 last_par = key[:2]
141 paragraphs[-1].append(text)
142 all_confs = [c for line in lines.values() for c in line['confs']]
143 result = {
144 'text': '\n\n'.join('\n'.join(p) for p in paragraphs),
145 'meanConfidence': round(sum(all_confs) / len(all_confs), 1) if all_confs else None,
146 'lines': out_lines,
147 'width': img.width,
148 'height': img.height,
149 }
150 if words:
151 result['words'] = word_list
152 return result
