Spaces:
Running on Zero
Running on Zero
| 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) | |