doc_intelligence/src/ocr/ocr_space_engine.py

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