65 lines
2.4 KiB
Python
65 lines
2.4 KiB
Python
"""OCR fallback for scanned PDFs (Tesseract, rus+eng).
|
||
|
||
Triggered by the extract worker only when `extract_text` returns < 100 chars
|
||
(a strong scan signal). Rasterizes each PDF page via pymupdf and runs
|
||
pytesseract. Both libs are imported lazily so this module imports cleanly even
|
||
in images without tesseract/pymupdf installed (the bot/analyze images).
|
||
|
||
Yandex Vision is a future alternative behind the same `ocr_pdf(path)` entry.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
from pathlib import Path
|
||
|
||
from ..logging import get_logger
|
||
from .extractor import ExtractionError
|
||
|
||
log = get_logger(__name__)
|
||
|
||
|
||
class OCRError(Exception):
|
||
"""OCR failed (Tesseract unavailable, rasterization failed, no text)."""
|
||
|
||
|
||
def ocr_pdf(path: str | Path, *, lang: str = "rus+eng") -> str:
|
||
"""Rasterize and OCR a PDF → plaintext. Raises OCRError on failure."""
|
||
p = Path(path)
|
||
try:
|
||
import pymupdf
|
||
import pytesseract
|
||
from PIL import Image
|
||
except ImportError as exc:
|
||
raise OCRError(f"OCR backend not installed: {exc.name}") from exc
|
||
|
||
try:
|
||
doc = pymupdf.open(p)
|
||
except Exception as exc: # pymupdf кидает разные типы
|
||
raise OCRError(f"Не удалось открыть PDF для OCR {p.name}: {exc}") from exc
|
||
|
||
parts: list[str] = []
|
||
try:
|
||
for page_index in range(doc.page_count):
|
||
page = doc[page_index]
|
||
pix = page.get_pixmap(dpi=300)
|
||
img = Image.frombytes("RGB", (pix.width, pix.height), pix.samples)
|
||
text = pytesseract.image_to_string(img, lang=lang)
|
||
if text:
|
||
parts.append(text)
|
||
except pytesseract.pytesseract.TesseractError as exc:
|
||
# Tesseract itself failed (missing language data, etc.) — surface as
|
||
# an OCR infra failure so the retry/DLQ path can refund appropriately.
|
||
raise OCRError(f"Tesseract engine failed for {p.name}: {exc}") from exc
|
||
except Exception as exc:
|
||
raise OCRError(f"OCR failed for {p.name}: {exc}") from exc
|
||
finally:
|
||
doc.close()
|
||
|
||
stripped = "\n".join(parts).strip()
|
||
if len(stripped) < 100:
|
||
raise ExtractionError(
|
||
f"OCR тоже дал мало текста ({len(stripped)} симв.). "
|
||
"Файл, видимо, не содержит распознаваемого текста."
|
||
)
|
||
log.info("ocr_done", file=p.name, chars=len(stripped))
|
||
return stripped
|