Sorin Savu
57 lines
1.8 KiB
Python
57 lines
1.8 KiB
Python
"""Pluggable OCR.
|
|
|
|
The engine is chosen by the OCR_ENGINE env var:
|
|
|
|
tesseract (default) -- pytesseract + Pillow + the system `tesseract` binary
|
|
none -- OCR disabled; images are marked 'skipped'
|
|
|
|
Swapping in another engine later (e.g. a Claude-vision call) is just another
|
|
branch in `run_ocr` behind the same signature. `run_ocr` raises OCRUnavailable
|
|
when the configured engine's dependencies are missing (the worker then marks the
|
|
row 'skipped' rather than 'failed'), and lets real OCR errors propagate.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import io
|
|
import os
|
|
|
|
|
|
class OCRUnavailable(RuntimeError):
|
|
"""The configured OCR engine can't run (missing binary / library)."""
|
|
|
|
|
|
def engine() -> str:
|
|
return os.environ.get("OCR_ENGINE", "tesseract").lower()
|
|
|
|
|
|
def enabled() -> bool:
|
|
return engine() != "none"
|
|
|
|
|
|
def run_ocr(data: bytes, mime: str) -> str:
|
|
"""Return extracted text for an image. Raises OCRUnavailable if the engine
|
|
is missing, or any other exception on a genuine OCR failure."""
|
|
eng = engine()
|
|
if eng == "none":
|
|
raise OCRUnavailable("OCR_ENGINE=none")
|
|
if eng == "tesseract":
|
|
return _tesseract(data)
|
|
raise OCRUnavailable(f"unknown OCR_ENGINE={eng!r}")
|
|
|
|
|
|
def _tesseract(data: bytes) -> str:
|
|
try:
|
|
import pytesseract
|
|
from PIL import Image
|
|
except ImportError as e:
|
|
raise OCRUnavailable(f"pytesseract/Pillow not installed: {e}") from e
|
|
try:
|
|
img = Image.open(io.BytesIO(data))
|
|
except Exception as e: # not a decodable image
|
|
raise OCRUnavailable(f"not a decodable image: {e}") from e
|
|
try:
|
|
return pytesseract.image_to_string(img)
|
|
except pytesseract.TesseractNotFoundError as e:
|
|
raise OCRUnavailable("tesseract binary not found on PATH") from e
|