""" ocr_engine.py ───────────── PaddleOCR inference service with: - Angle classifier enabled (handles rotated/upside-down text lines) - Confidence-aware result passing (raw boxes + scores to the formatter) - Tesseract fallback for low-confidence pages """ import os import asyncio from typing import Dict, List, Optional, Tuple import numpy as np # Disable GPU and CPU MKL-DNN optimizations to prevent container segmentation faults os.environ["CUDA_VISIBLE_DEVICES"] = "-1" os.environ["FLAGS_use_mkldnn"] = "0" os.environ["FLAGS_use_xdnn"] = "0" os.environ["OMP_NUM_THREADS"] = "1" from paddleocr import PaddleOCR from src.config import settings # Threshold: if average page confidence falls below this, run Tesseract fallback TESS_FALLBACK_THRESHOLD = 0.65 class OCREngineService: def __init__(self): # In-memory cache of instantiated PaddleOCR engines keyed by language self._engines: Dict[str, PaddleOCR] = {} # Mutex lock to serialize CPU inference and avoid resource thrashing self._lock = asyncio.Lock() def get_engine(self, lang: str) -> PaddleOCR: """Retrieves or initialises a PaddleOCR engine for the given language.""" lang = lang.lower().strip() if lang not in settings.SUPPORTED_LANGUAGES: lang = settings.DEFAULT_LANGUAGE if lang not in self._engines: self._engines[lang] = PaddleOCR( use_angle_cls=True, # Detects 0°/180° text orientation (fixes upside-down lines) lang=lang, show_log=False, enable_mkldnn=False, det_limit_side_len=3000, # Preserve high-res details for tiny text (default 960) det_db_box_thresh=0.4, # Lower threshold — detect faint/thin boxes (was 0.5) det_db_score_mode="slow", # Accurate boundary evaluation (default 'fast') rec_db_score_mode="slow", use_space_char=True, # Recognise spaces between words cls_thresh=0.9, # High bar for angle flip — only flip when very confident ) return self._engines[lang] async def _detect_language(self, img_np: np.ndarray, loop) -> str: """ Uses Tesseract OSD to detect the writing system/script, and maps it to a supported language. """ try: import pytesseract import cv2 import re # Tesseract works best on grayscale gray = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY) osd_data = await loop.run_in_executor( None, lambda: pytesseract.image_to_osd(gray, config="--psm 0") ) # Parse Script Name (e.g. 'Script: Latin', 'Script: Arabic', etc.) match = re.search(r"Script:\s+(\w+)", osd_data) if match: script = match.group(1).lower().strip() # Map script to PaddleOCR language code script_map = { "latin": "en", # Latin alphabet covers English, Spanish, Portuguese, French, German, Italian "devanagari": "hi", # Hindi "arabic": "ar", # Arabic "han": "ch", # Chinese (Simplified) "cyrillic": "ru", # Russian "telugu": "te", # Telugu "tamil": "ta", # Tamil "japanese": "japan", # Japanese "korean": "korean" # Korean } detected = script_map.get(script, "en") print(f"OSD detected script '{script}', mapped to language code '{detected}'") return detected except Exception as e: print(f"OSD language detection failed: {e}. Falling back to default 'en'.") return "en" async def run_ocr_raw(self, img_np: np.ndarray, lang: str): """ Runs PaddleOCR and returns the raw result list (bounding boxes + text + confidence). The caller (routes.py) passes this to the markdown formatter. """ async with self._lock: loop = asyncio.get_running_loop() # Resolve auto language detection if requested if lang.lower().strip() in ("auto", "detect"): resolved_lang = await self._detect_language(img_np, loop) else: resolved_lang = lang ocr = self.get_engine(resolved_lang) raw = await loop.run_in_executor( None, lambda: ocr.ocr(img_np, cls=True) ) # raw is [[result_page1], ...] — we only handle single-page images page = raw[0] if raw else [] # Optionally run Tesseract fallback if average confidence is low if _should_run_tesseract(page): tess_lines = await _run_tesseract_fallback(img_np, loop) page = _merge_tesseract_lines(page, tess_lines) return page # ───────────────────────────────────────────── # Tesseract fallback helpers # ───────────────────────────────────────────── def _should_run_tesseract(page_result) -> bool: """ Returns True if the average PaddleOCR confidence score for the page is below TESS_FALLBACK_THRESHOLD, indicating the page may have zones of text that PaddleOCR struggled with. """ if not page_result: return True # Empty result — definitely try fallback scores = [item[1][1] for item in page_result if item and item[1]] if not scores: return True avg_conf = sum(scores) / len(scores) return avg_conf < TESS_FALLBACK_THRESHOLD async def _run_tesseract_fallback(img_np: np.ndarray, loop) -> List[str]: """ Runs pytesseract in sparse-text mode as a fallback for difficult pages. psm 11 = Sparse text. Find as much text as possible in no particular order. oem 1 = LSTM neural net engine only. """ try: import pytesseract import cv2 # Tesseract works best on grayscale gray = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY) config = "--oem 1 --psm 11" text = await loop.run_in_executor( None, lambda: pytesseract.image_to_string(gray, config=config, lang="eng") ) # Split into lines, strip whitespace, drop empty lines = [l.strip() for l in text.splitlines() if l.strip()] return lines except ImportError: # pytesseract not installed — silently skip return [] except Exception as e: print(f"Tesseract fallback failed: {e}") return [] def _merge_tesseract_lines(paddle_page, tess_lines: List[str]): """ Merge Tesseract-recovered lines into the PaddleOCR result. Strategy: collect all text already captured by PaddleOCR, then append any Tesseract lines that are not already present as a 'synthetic' low-priority item (no bounding box, confidence 0.60). This ensures the formatter sees them without duplicating content. """ if not tess_lines: return paddle_page existing_texts = set() if paddle_page: for item in paddle_page: if item and item[1]: existing_texts.add(item[1][0].strip().lower()) merged = list(paddle_page) if paddle_page else [] for line in tess_lines: normalised = line.strip().lower() if normalised and normalised not in existing_texts: # Synthetic entry: no real bounding box — place at the bottom of the page # The formatter handles this gracefully (falls back to plain text for lines without boxes) merged.append([ [[0, 9999], [100, 9999], [100, 10000], [0, 10000]], (line, 0.60) ]) existing_texts.add(normalised) return merged # Global singleton ocr_service = OCREngineService()