"""OCR engine abstraction with local PaddleOCR/EasyOCR backends. Configuration is controlled by environment variables: - ``SHELF_OCR_ENGINE=auto|paddle_v4|paddle_v5|easyocr|none`` - ``SHELF_OCR_AUTO_ORDER=paddle_v4,paddle_v5,easyocr`` - ``SHELF_PADDLE_DET_MODEL_DIR`` / ``SHELF_PADDLE_REC_MODEL_DIR`` / ``SHELF_PADDLE_CLS_MODEL_DIR`` for fully local Paddle model placement. No cloud OCR API is used. PP-OCRv5 is optional and selected only when available or explicitly requested; stable PP-OCRv4 remains in the default auto order until local challenge metrics prove v5 is better on the five labeled videos. """ from __future__ import annotations import logging import os from threading import Lock from typing import Any import numpy as np logger = logging.getLogger(__name__) _VALID_ENGINES = {"auto", "paddle_v4", "paddle_v5", "easyocr", "none"} _DEFAULT_AUTO_ORDER = ("paddle_v4", "paddle_v5", "easyocr") def normalize_ocr_engine_name(value: str | None) -> str: """Normalize OCR engine selection from config/env.""" name = ( (value or os.getenv("SHELF_OCR_ENGINE", "auto")) .strip() .lower() .replace("-", "_") ) aliases = { "paddle": "paddle_v4", "paddle4": "paddle_v4", "ppocr4": "paddle_v4", "pp_ocr_v4": "paddle_v4", "paddle5": "paddle_v5", "ppocr5": "paddle_v5", "pp_ocr_v5": "paddle_v5", "off": "none", "disabled": "none", } name = aliases.get(name, name) return name if name in _VALID_ENGINES else "auto" def _auto_order() -> list[str]: raw = os.getenv("SHELF_OCR_AUTO_ORDER", ",".join(_DEFAULT_AUTO_ORDER)) order = [normalize_ocr_engine_name(x) for x in raw.split(",")] return [ x for x in order if x in {"paddle_v4", "paddle_v5", "easyocr"} ] or list(_DEFAULT_AUTO_ORDER) def _paddle_model_dir_kwargs() -> dict[str, str]: """Collect optional local Paddle model directories from env.""" mapping = { "det_model_dir": os.getenv("SHELF_PADDLE_DET_MODEL_DIR"), "rec_model_dir": os.getenv("SHELF_PADDLE_REC_MODEL_DIR"), "cls_model_dir": os.getenv("SHELF_PADDLE_CLS_MODEL_DIR"), } return {k: v for k, v in mapping.items() if v} class OCREngine: """Lazy OCR backend wrapper. PaddleOCR is preferred for digits/prices. EasyOCR is a local fallback and is still useful for Russian product names. The wrapper avoids importing heavy OCR packages until the first inference call. """ _easyocr_lock = Lock() def __init__( self, lang: str = "en", force_easyocr: bool = False, engine: str | None = None, ) -> None: self.lang = lang self.force_easyocr = force_easyocr self.engine_name = ( "easyocr" if force_easyocr else normalize_ocr_engine_name(engine) ) self._ocr: Any | None = None self._backend: str | None = ( None # paddle_v4 | paddle_v5 | easyocr | none ) def _load(self) -> None: if self.engine_name == "none": self._backend = "none" self._ocr = None logger.info("OCR backend disabled by SHELF_OCR_ENGINE=none") return order = ( [self.engine_name] if self.engine_name != "auto" else _auto_order() ) for backend in order: if backend.startswith("paddle"): if self._try_load_paddle(backend): return elif backend == "easyocr": if self._try_load_easyocr(): return self._ocr = None self._backend = "none" logger.warning("No local OCR backend could be loaded; OCR disabled") def _try_load_paddle(self, backend: str) -> bool: try: from paddleocr import PaddleOCR except Exception as exc: logger.info("PaddleOCR import failed for %s: %s", backend, exc) return False kwargs = _paddle_model_dir_kwargs() attempts = self._paddle_constructor_attempts(backend, kwargs) for attempt in attempts: try: self._ocr = PaddleOCR(**attempt) self._backend = backend logger.info( "OCR backend: PaddleOCR %s (lang=%s)", backend, self.lang ) return True except TypeError as exc: logger.debug( "PaddleOCR constructor rejected %s: %s", attempt, exc ) except Exception as exc: logger.info( "PaddleOCR %s unavailable with %s: %s", backend, attempt, exc, ) return False def _paddle_constructor_attempts( self, backend: str, model_dirs: dict[str, str] ) -> list[dict[str, Any]]: """Return constructor kwargs from modern to legacy PaddleOCR APIs.""" version = "PP-OCRv5" if backend == "paddle_v5" else "PP-OCRv4" common: dict[str, Any] = {"lang": self.lang, **model_dirs} attempts: list[dict[str, Any]] = [ { **common, "ocr_version": version, "use_textline_orientation": True, }, { **common, "ocr_version": version, "use_angle_cls": True, "show_log": False, }, {**common, "use_textline_orientation": True}, {**common, "use_angle_cls": True, "show_log": False}, ] if backend == "paddle_v5": attempts.insert( 0, { **common, "ocr_version": "PP-OCRv5", "text_detection_model_name": "PP-OCRv5_mobile_det", "text_recognition_model_name": "PP-OCRv5_mobile_rec", "use_textline_orientation": True, }, ) return attempts def _try_load_easyocr(self) -> bool: try: import easyocr langs = ["ru", "en"] if self.lang == "ru" else ["en"] with self._easyocr_lock: self._ocr = easyocr.Reader(langs, verbose=False) self._backend = "easyocr" logger.info("OCR backend: EasyOCR (langs=%s)", langs) return True except Exception as exc: logger.info("EasyOCR unavailable: %s", exc) return False def run(self, image: np.ndarray) -> list[tuple[list, str, float]]: """Return a list of ``(box, text, confidence)`` tuples.""" if image is None or image.size == 0: return [] if self._backend is None: self._load() if self._backend == "none" or self._ocr is None: return [] try: if str(self._backend).startswith("paddle"): return self._run_paddle(image) return self._run_easyocr(image) except Exception as exc: logger.warning("OCR inference failed (%s)", exc) return [] def _run_paddle(self, image: np.ndarray) -> list[tuple[list, str, float]]: if hasattr(self._ocr, "ocr"): try: results = self._ocr.ocr(image, cls=True) except TypeError: results = self._ocr.ocr(image) elif hasattr(self._ocr, "predict"): results = self._ocr.predict(image) else: return [] return self._parse_paddle_results(results) def _parse_paddle_results( self, results: Any ) -> list[tuple[list, str, float]]: out: list[tuple[list, str, float]] = [] if not results: return out for page in results if isinstance(results, list) else [results]: if page is None: continue if isinstance(page, dict): out.extend(self._parse_paddle_dict(page)) continue for item in page or []: if item is None: continue if isinstance(item, dict): out.extend(self._parse_paddle_dict(item)) continue try: box, payload = item if not payload: continue text, conf = payload out.append((box, str(text), float(conf))) except Exception: continue return out @staticmethod def _parse_paddle_dict( page: dict[str, Any], ) -> list[tuple[list, str, float]]: texts = page.get("rec_texts") or page.get("texts") or [] scores = page.get("rec_scores") or page.get("scores") or [] boxes = ( page.get("dt_polys") or page.get("rec_polys") or page.get("boxes") or [] ) out: list[tuple[list, str, float]] = [] for idx, text in enumerate(texts): box = ( boxes[idx].tolist() if idx < len(boxes) and hasattr(boxes[idx], "tolist") else (boxes[idx] if idx < len(boxes) else []) ) conf = float(scores[idx]) if idx < len(scores) else 0.0 out.append((box, str(text), conf)) return out def _run_easyocr(self, image: np.ndarray) -> list[tuple[list, str, float]]: results = self._ocr.readtext(image, detail=1) out: list[tuple[list, str, float]] = [] for box_pts, text, conf in results: out.append((box_pts, str(text), float(conf))) return out def run_texts(self, image: np.ndarray) -> list[str]: return [text for _, text, _ in self.run(image)]