File size: 2,059 Bytes
dde961d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 | """OmniParse AI — OCR pipeline: Tesseract (primary) + Google Vision (fallback)."""
import io, base64
from PIL import Image as PILImage
import requests
from config import GOOGLE_VISION_KEY
# Lazy imports
_pt = None; _pdf2img = None
def _get_pt():
global _pt
if _pt is None:
try:
import pytesseract as t
_pt = t
except ImportError: pass
return _pt
def _get_pdf2img():
global _pdf2img
if _pdf2img is None:
try:
from pdf2image import convert_from_path as c
_pdf2img = c
except ImportError: pass
return _pdf2img
def pdf_to_images(path: str, dpi: int = 200) -> list:
c = _get_pdf2img()
if not c: return []
try: return c(path, dpi=dpi)
except Exception as e: print(f"[WARN] pdf2image: {e}"); return []
def load_image(path: str):
try: return PILImage.open(path).convert("RGB")
except Exception as e: print(f"[WARN] open image: {e}"); return None
def ocr_tesseract(img, lang: str = "eng") -> str:
pt = _get_pt()
if not pt: return ""
try: return pt.image_to_string(img.convert("L"), lang=lang)
except Exception as e: print(f"[WARN] tesseract: {e}"); return ""
def ocr_google_vision(img) -> str:
if not GOOGLE_VISION_KEY: return ""
try:
buf = io.BytesIO(); img.save(buf, format="PNG")
b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
payload = {"requests":[{"image":{"content":b64},"features":[{"type":"TEXT_DETECTION"}]}]}
r = requests.post(
f"https://vision.googleapis.com/v1/images:annotate?key={GOOGLE_VISION_KEY}",
json=payload, timeout=15)
r.raise_for_status()
return r.json()["responses"][0].get("fullTextAnnotation",{}).get("text","")
except Exception as e: print(f"[WARN] vision: {e}"); return ""
def run_ocr(img) -> str:
text = ocr_tesseract(img)
if len(text.strip()) < 100 and GOOGLE_VISION_KEY:
vt = ocr_google_vision(img)
if len(vt.strip()) > len(text.strip()): text = vt
return text
|