| """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 |
|
|
| |
| _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 |
|
|
