import io import time import logging from typing import List, Dict, Any, Tuple from PIL import Image, ImageOps import numpy as np from .config import settings logger = logging.getLogger("document_intelligence.ocr") class OCRService: """Modular OCR service supporting PaddleOCR with automatic fallbacks.""" def __init__(self, engine: str = "auto"): self.engine_name = engine self._paddle_ocr = None self._rapid_ocr = None self.active_engine = "rapidocr" self._initialize_engine() def _initialize_engine(self): """Pre-load and pre-warm OCR engines at application startup.""" # 1. Initialize RapidOCR (ONNX Runtime - fast, portable, robust) try: from rapidocr_onnxruntime import RapidOCR self._rapid_ocr = RapidOCR() # Pre-warm with dummy numpy array to load ONNX weights into RAM dummy = np.zeros((64, 64, 3), dtype=np.uint8) self._rapid_ocr(dummy) self.active_engine = "rapidocr" logger.info("RapidOCR engine initialized and pre-warmed successfully.") except Exception as e: logger.warning("RapidOCR initialization failed: %s", e) # 2. Initialize PaddleOCR if requested try: from paddleocr import PaddleOCR try: self._paddle_ocr = PaddleOCR(use_angle_cls=True, lang="en") except Exception: self._paddle_ocr = PaddleOCR(lang="en") self.active_engine = "paddleocr" logger.info("PaddleOCR engine initialized successfully.") except Exception as e: logger.info("PaddleOCR not available or initialization skipped: %s", e) def validate_and_preprocess_image( self, image_input: Any ) -> Tuple[Image.Image, np.ndarray]: """Validate format, size, dimensions and prepare numpy array for OCR.""" image: Optional[Image.Image] = None if isinstance(image_input, Image.Image): image = image_input elif isinstance(image_input, np.ndarray): image = Image.fromarray(image_input) elif isinstance(image_input, dict): # Gradio 5/6 image editor / sketch dictionary format for key in ("composite", "background", "image"): val = image_input.get(key) if isinstance(val, Image.Image): image = val break elif isinstance(val, np.ndarray): image = Image.fromarray(val) break if image is None and image_input: first_val = next(iter(image_input.values())) if isinstance(first_val, Image.Image): image = first_val elif isinstance(first_val, np.ndarray): image = Image.fromarray(first_val) elif isinstance(image_input, bytes): image = Image.open(io.BytesIO(image_input)) elif isinstance(image_input, str): image = Image.open(image_input) if image is None: raise ValueError(f"Could not parse valid image from input of type: {type(image_input)}") # Convert to RGB if image.mode != "RGB": image = image.convert("RGB") # Correct EXIF orientation if present image = ImageOps.exif_transpose(image) # Dimension resize if extreme max_dim = settings.MAX_IMAGE_DIMENSION width, height = image.size if width > max_dim or height > max_dim: logger.info("Resizing image from (%d, %d) to fit max dim %d", width, height, max_dim) image.thumbnail((max_dim, max_dim), Image.Resampling.LANCZOS) np_image = np.array(image) return image, np_image def extract_text(self, image_input: Any) -> Dict[str, Any]: """ Run OCR on the given image. Returns a dictionary with: - full_text: string representation of detected text - lines: list of detected text lines with confidence & coordinates - spatial_hints: formatted string describing spatial positioning - latency_seconds: execution time - engine: OCR engine used """ start_time = time.time() image, np_image = self.validate_and_preprocess_image(image_input) lines: List[Dict[str, Any]] = [] used_engine = self.active_engine # 1. Try RapidOCR (primary fast & reliable ONNX engine) if self._rapid_ocr: try: result, _ = self._rapid_ocr(np_image) if result: for item in result: box = item[0] text = item[1] confidence = item[2] lines.append({ "text": str(text).strip(), "confidence": float(confidence), "box": box }) if lines: used_engine = "rapidocr" except Exception as e: logger.warning("RapidOCR execution error: %s", e) # 2. Try PaddleOCR if lines still empty if not lines and self._paddle_ocr: try: np_bgr = np_image[:, :, ::-1] if len(np_image.shape) == 3 else np_image results = self._paddle_ocr.ocr(np_bgr, cls=True) if results and results[0]: for item in results[0]: if item and len(item) >= 2: box = item[0] text, confidence = item[1] lines.append({ "text": str(text).strip(), "confidence": float(confidence), "box": box }) if lines: used_engine = "paddleocr" except Exception as e: logger.warning("PaddleOCR execution error: %s", e) # 3. If lines are still empty and pytesseract exists, try it if not lines: try: import pytesseract ocr_data = pytesseract.image_to_data(image, output_type=pytesseract.Output.DICT) n_boxes = len(ocr_data['text']) for i in range(n_boxes): text = ocr_data['text'][i].strip() conf = float(ocr_data['conf'][i]) if text and conf > 0: lines.append({ "text": text, "confidence": conf / 100.0, "box": [ [ocr_data['left'][i], ocr_data['top'][i]], [ocr_data['left'][i] + ocr_data['width'][i], ocr_data['top'][i]], [ocr_data['left'][i] + ocr_data['width'][i], ocr_data['top'][i] + ocr_data['height'][i]], [ocr_data['left'][i], ocr_data['top'][i] + ocr_data['height'][i]] ] }) if lines: used_engine = "pytesseract" except Exception: pass # Format full text extracted_text_lines = [item["text"] for item in lines if item["text"]] full_text = "\n".join(extracted_text_lines) # Construct spatial hint spatial_hints_list = [] for idx, item in enumerate(lines[:50]): # cap to top 50 lines for prompt compactness spatial_hints_list.append(f"Line {idx+1}: '{item['text']}' (confidence: {item['confidence']:.2f})") spatial_hints = "\n".join(spatial_hints_list) latency = time.time() - start_time logger.info( "OCR complete [%s]: extracted %d lines in %.3f seconds", used_engine, len(lines), latency ) return { "full_text": full_text, "lines": lines, "spatial_hints": spatial_hints, "latency_seconds": round(latency, 4), "engine": used_engine, "image_size": image.size } # Global singleton instance ocr_service = OCRService(engine=settings.OCR_ENGINE)