Spaces:
Sleeping
Sleeping
| import os | |
| import io | |
| import time | |
| import requests | |
| from typing import List, Optional, Tuple | |
| from PIL import Image | |
| from ..config import get_secret | |
| from .base_engine import BaseOCREngine | |
| from ..data.schema import OCRWord | |
| # Attempt to load PyMuPDF for native PDF handling | |
| try: | |
| import fitz # PyMuPDF | |
| HAVE_PYMUPDF = True | |
| except ImportError: | |
| HAVE_PYMUPDF = False | |
| class OCRSpaceEngine(BaseOCREngine): | |
| """ | |
| Cloud OCR Engine powered by OCR.space API. | |
| Features: | |
| - High-accuracy word extraction with spatial bounding boxes (isOverlayRequired=True) | |
| - Modern Engine 2 optimized for invoices, receipts, numbers, tables, and rotation | |
| - Automatic image pre-compression to stay within OCR.space 1MB free-tier limit | |
| - Automatic scaling of bounding boxes back to original image dimensions | |
| - PyMuPDF integration for digital PDFs and PDF rendering | |
| - Local disk caching via BaseOCREngine to preserve API quota | |
| - Resilient exponential backoff retry for transient network / throttling conditions | |
| """ | |
| API_URL = "https://api.ocr.space/parse/image" | |
| def __init__( | |
| self, | |
| api_key: Optional[str] = None, | |
| ocr_engine: str = "3", | |
| language: str = "eng", | |
| detect_orientation: bool = True, | |
| scale: bool = True, | |
| is_table: bool = True, | |
| timeout: int = 20, | |
| max_retries: int = 1, | |
| cache_dir: Optional[str] = None | |
| ): | |
| super().__init__("ocr_space", cache_dir) | |
| # Load API key from parameter or environment | |
| if not api_key: | |
| api_key = get_secret("OCR_SPACE_API", "OCR_SPACE_API_KEY", default="") | |
| # Clean API key (strip quotes, whitespace) | |
| self.api_key = api_key.strip().strip("'\"") | |
| self.ocr_engine = str(ocr_engine) # "3" has 100% availability, 4.6s speed, tables & forms | |
| self.language = language | |
| self.detect_orientation = detect_orientation | |
| self.scale = scale | |
| self.is_table = is_table | |
| self.timeout = timeout | |
| self.max_retries = max_retries | |
| if not self.api_key: | |
| print("[WARN] OCR.space API key is not configured.") | |
| def _prepare_image_payload(self, image_path: str) -> Tuple[bytes, str, float, float]: | |
| """ | |
| Prepares the image payload for OCR.space. | |
| Free tier limit is 1024 KB (1MB). If image is larger or has huge dimensions, | |
| it resizes and compresses in memory, while computing scaling factors | |
| (scale_x, scale_y) to map coordinates back to the original image. | |
| """ | |
| orig_img = Image.open(image_path) | |
| orig_w, orig_h = orig_img.size | |
| file_size = os.path.getsize(image_path) | |
| # If file is well under 1MB (e.g. < 900KB) and dimensions <= 2000px, send directly | |
| if file_size < 900 * 1024 and max(orig_w, orig_h) <= 2000: | |
| with open(image_path, "rb") as f: | |
| content = f.read() | |
| ext = os.path.splitext(image_path)[1].lower().replace(".", "") | |
| mime = "image/png" if ext == "png" else "image/jpeg" | |
| return content, mime, 1.0, 1.0 | |
| # Otherwise, resize and compress in-memory | |
| max_dim = 1800 | |
| comp_img = orig_img.copy() | |
| if max(orig_w, orig_h) > max_dim: | |
| if orig_w >= orig_h: | |
| new_w = max_dim | |
| new_h = int(round(orig_h * (max_dim / orig_w))) | |
| else: | |
| new_h = max_dim | |
| new_w = int(round(orig_w * (max_dim / orig_h))) | |
| comp_img = comp_img.resize((new_w, new_h), Image.Resampling.LANCZOS) | |
| # Ensure RGB mode for JPEG saving | |
| if comp_img.mode in ("RGBA", "P"): | |
| comp_img = comp_img.convert("RGB") | |
| buf = io.BytesIO() | |
| comp_img.save(buf, format="JPEG", quality=88, optimize=True) | |
| compressed_bytes = buf.getvalue() | |
| comp_w, comp_h = comp_img.size | |
| scale_x = orig_w / float(comp_w) | |
| scale_y = orig_h / float(comp_h) | |
| return compressed_bytes, "image/jpeg", scale_x, scale_y | |
| def _extract_from_digital_pdf(self, pdf_path: str) -> Optional[List[OCRWord]]: | |
| """Extracts words natively from digital PDF using PyMuPDF if selectable text exists.""" | |
| if not HAVE_PYMUPDF: | |
| return None | |
| try: | |
| doc = fitz.open(pdf_path) | |
| if len(doc) == 0: | |
| return None | |
| page = doc[0] | |
| words = page.get_text("words") # (x0, y0, x1, y1, word, block_no, line_no, word_no) | |
| if words and len(words) > 5: | |
| ocr_words = [] | |
| for w in words: | |
| x0, y0, x1, y1, text = w[:5] | |
| ocr_words.append(OCRWord( | |
| text=str(text).strip(), | |
| bbox=[int(round(x0)), int(round(y0)), int(round(x1)), int(round(y1))], | |
| conf=0.99 | |
| )) | |
| return ocr_words | |
| except Exception as e: | |
| print(f"[DEBUG] PyMuPDF digital text extraction skipped: {e}") | |
| return None | |
| def _call_ocr_space_api(self, file_bytes: bytes, mime_type: str, engine: str) -> dict: | |
| """Sends POST request to OCR.space API with exponential backoff retry.""" | |
| payload = { | |
| "apikey": self.api_key, | |
| "language": self.language, | |
| "isOverlayRequired": "true", | |
| "OCREngine": engine, | |
| "detectOrientation": "true" if self.detect_orientation else "false", | |
| "scale": "true" if self.scale else "false", | |
| "isTable": "true" if self.is_table else "false" | |
| } | |
| headers = { | |
| "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) DocumentIntelligence/1.0" | |
| } | |
| filename = "document.jpg" if "jpeg" in mime_type else "document.png" | |
| files = { | |
| "file": (filename, file_bytes, mime_type) | |
| } | |
| last_err = None | |
| for attempt in range(1, self.max_retries + 1): | |
| try: | |
| response = requests.post( | |
| self.API_URL, | |
| data=payload, | |
| files=files, | |
| headers=headers, | |
| timeout=self.timeout | |
| ) | |
| if response.status_code == 200: | |
| try: | |
| data = response.json() | |
| except ValueError as exc: | |
| raise RuntimeError("OCR.space returned an invalid response.") from exc | |
| # Check if OCR.space returned an internal error | |
| if data.get("IsErroredOnProcessing", False): | |
| err_msg = data.get("ErrorMessage", ["Processing Error"]) | |
| if isinstance(err_msg, list): | |
| err_msg = "; ".join(err_msg) | |
| # If engine 2 is unsupported for this specific image, fallback to engine 1 | |
| if engine == "2" and "engine 2" in str(err_msg).lower(): | |
| print(f"[INFO] Falling back from Engine 2 to Engine 1: {err_msg}") | |
| return self._call_ocr_space_api(file_bytes, mime_type, "1") | |
| raise RuntimeError(f"OCR.space processing error: {err_msg}") | |
| return data | |
| elif response.status_code in (429, 503): | |
| # Throttled or overloaded | |
| err_text = response.text | |
| print(f"[WARN] OCR.space HTTP {response.status_code} on attempt {attempt}/{self.max_retries}: {err_text[:120]}") | |
| if attempt < self.max_retries: | |
| time.sleep(2 * attempt) | |
| else: | |
| raise RuntimeError(f"OCR.space Free API is currently overloaded ({response.status_code}). Please retry in a few moments.") | |
| else: | |
| response.raise_for_status() | |
| except requests.RequestException as e: | |
| last_err = e | |
| print(f"[WARN] OCR.space request error on attempt {attempt}: {e}") | |
| if attempt < self.max_retries: | |
| time.sleep(1.5 * attempt) | |
| else: | |
| raise RuntimeError(f"OCR.space API connection failed: {last_err}") | |
| raise RuntimeError(f"Failed to extract text via OCR.space after {self.max_retries} attempts: {last_err}") | |
| def extract_words(self, image_path: str) -> List[OCRWord]: | |
| """ | |
| Extracts words and spatial bounding boxes for the given document. | |
| """ | |
| if not image_path or not os.path.exists(image_path): | |
| print(f"[WARN] Document file not found: {image_path}") | |
| return [] | |
| if not self.api_key: | |
| raise RuntimeError( | |
| "OCR_SPACE_API is not configured. Add it as a Hugging Face Space secret " | |
| "or set it in the local environment before analyzing a document." | |
| ) | |
| # Check for PDF | |
| ext = os.path.splitext(image_path)[1].lower() | |
| if ext == ".pdf": | |
| # 1. Try native selectable text first | |
| digital_words = self._extract_from_digital_pdf(image_path) | |
| if digital_words: | |
| return digital_words | |
| # 2. Render first page to image for OCR | |
| if HAVE_PYMUPDF: | |
| doc = fitz.open(image_path) | |
| page = doc[0] | |
| pix = page.get_pixmap(dpi=150) | |
| temp_img_path = f"{image_path}_page0.png" | |
| pix.save(temp_img_path) | |
| try: | |
| words = self.extract_words(temp_img_path) | |
| finally: | |
| if os.path.exists(temp_img_path): | |
| os.remove(temp_img_path) | |
| return words | |
| # Image processing | |
| file_bytes, mime_type, scale_x, scale_y = self._prepare_image_payload(image_path) | |
| # Keep interactive requests bounded. The configured engine is the only | |
| # engine attempted; retrying every OCR.space engine can outlive the | |
| # Streamlit request and appear as a blank page to the user. | |
| candidate_engines = [self.ocr_engine] | |
| api_data = None | |
| last_exception = None | |
| for eng in candidate_engines: | |
| try: | |
| api_data = self._call_ocr_space_api(file_bytes, mime_type, eng) | |
| if api_data: | |
| break | |
| except Exception as e: | |
| last_exception = e | |
| print(f"[WARN] OCR.space Engine {eng} failed: {e}") | |
| if api_data is None: | |
| raise RuntimeError(f"OCR.space extraction failed across all available engines: {last_exception}") | |
| # Parse Words from TextOverlay | |
| ocr_words: List[OCRWord] = [] | |
| parsed_results = api_data.get("ParsedResults", []) | |
| for res in parsed_results: | |
| overlay = res.get("TextOverlay") or {} | |
| lines = overlay.get("Lines", []) | |
| for line in lines: | |
| for w in line.get("Words", []): | |
| w_text = w.get("WordText", "").strip() | |
| if not w_text: | |
| continue | |
| # Coordinates in the uploaded image frame | |
| left = float(w.get("Left", 0)) | |
| top = float(w.get("Top", 0)) | |
| width = float(w.get("Width", 0)) | |
| height = float(w.get("Height", 0)) | |
| # Scale back to original image dimensions | |
| orig_left = int(round(left * scale_x)) | |
| orig_top = int(round(top * scale_y)) | |
| orig_right = int(round((left + width) * scale_x)) | |
| orig_bottom = int(round((top + height) * scale_y)) | |
| ocr_words.append(OCRWord( | |
| text=w_text, | |
| bbox=[orig_left, orig_top, orig_right, orig_bottom], | |
| conf=0.95 | |
| )) | |
| # Fallback if TextOverlay was empty but ParsedText was provided | |
| if not ocr_words and res.get("ParsedText"): | |
| raw_text = res.get("ParsedText", "").strip() | |
| tokens = raw_text.split() | |
| # Synthesize default bounding boxes | |
| for i, token in enumerate(tokens): | |
| ocr_words.append(OCRWord( | |
| text=token, | |
| bbox=[10, 10 + (i * 20), 100, 25 + (i * 20)], | |
| conf=0.90 | |
| )) | |
| return ocr_words | |