ocr/easyocr_engine.py

"""
InvoiceForge AI — ocr/easyocr_engine.py
EasyOCR fallback engine wrapper.

Used as a confidence-weighted secondary engine when PaddleOCR confidence
is below threshold. Runs CPU-only by default.
"""

from __future__ import annotations

import logging
from typing import Optional

import cv2
import numpy as np

logger = logging.getLogger(__name__)

EASYOCR_LANGS = ["en"]
CONFIDENCE_THRESHOLD = 0.40  # tokens below this are not trusted


class EasyOCREngine:
    """
    Singleton wrapper around EasyOCR Reader.

    EasyOCR is initialised lazily on first use. The model weights are
    cached in ~/.EasyOCR/ by default.

    Usage:
        engine = EasyOCREngine.instance()
        tokens = engine.run_ocr(img)
    """

    _reader: Optional[object] = None

    @classmethod
    def instance(cls) -> "EasyOCREngine":
        obj = cls.__new__(cls)
        return obj

    def _get_reader(self) -> object:
        """Lazily initialise EasyOCR Reader."""
        if EasyOCREngine._reader is None:
            logger.info("Initialising EasyOCR reader (langs=%s)…", EASYOCR_LANGS)
            try:
                import easyocr  # type: ignore

                EasyOCREngine._reader = easyocr.Reader(
                    EASYOCR_LANGS,
                    gpu=False,
                    verbose=False,
                )
                logger.info("EasyOCR reader initialised.")
            except Exception as exc:
                logger.error("EasyOCR init failed: %s", exc)
                raise
        return EasyOCREngine._reader

    def run_ocr(self, img: np.ndarray) -> list[dict]:
        """
        Run EasyOCR on the given image.

        Args:
            img: BGR or grayscale numpy array.

        Returns:
            List of dicts: {text, confidence, bbox, x, y, engine}
        """
        reader = self._get_reader()

        # EasyOCR accepts BGR; ensure correct dtype
        if img.dtype != np.uint8:
            img = img.astype(np.uint8)

        try:
            raw_result = reader.readtext(img, detail=1)  # type: ignore[union-attr]
        except Exception as exc:
            logger.warning("EasyOCR inference failed: %s", exc)
            return []

        tokens: list[dict] = []
        for detection in raw_result:
            bbox_raw, text, conf = detection
            if conf < CONFIDENCE_THRESHOLD:
                continue
            # bbox_raw: list of [x, y] corners
            xs = [p[0] for p in bbox_raw]
            ys = [p[1] for p in bbox_raw]
            tokens.append(
                {
                    "text": text.strip(),
                    "confidence": float(conf),
                    "bbox": bbox_raw,
                    "x": float(np.mean(xs)),
                    "y": float(np.mean(ys)),
                    "engine": "easyocr",
                }
            )

        tokens.sort(key=lambda t: (t["y"], t["x"]))
        logger.debug("EasyOCR: %d tokens extracted.", len(tokens))
        return tokens

    def run_ocr_on_region(
        self,
        img: np.ndarray,
        x1: int,
        y1: int,
        x2: int,
        y2: int,
    ) -> list[dict]:
        """
        Run EasyOCR on a specific rectangular region of the image.
        Coordinates are adjusted back to full-image space.

        Args:
            img:        Full image BGR array.
            x1, y1:    Top-left of region.
            x2, y2:    Bottom-right of region.

        Returns:
            List of token dicts with full-image coordinates.
        """
        h, w = img.shape[:2]
        x1 = max(0, x1)
        y1 = max(0, y1)
        x2 = min(w, x2)
        y2 = min(h, y2)
        crop = img[y1:y2, x1:x2]
        if crop.size == 0:
            return []

        tokens = self.run_ocr(crop)
        # Offset coordinates to full-image space
        for tok in tokens:
            tok["x"] += x1
            tok["y"] += y1
            tok["bbox"] = [
                [p[0] + x1, p[1] + y1] for p in tok["bbox"]
            ]
        return tokens


# ─────────────────────────────────────────────────────────────────────────────
# MODULE-LEVEL SINGLETON
# ─────────────────────────────────────────────────────────────────────────────

_engine: EasyOCREngine | None = None


def get_easyocr_engine() -> EasyOCREngine:
    """Return the module-level EasyOCREngine singleton."""
    global _engine
    if _engine is None:
        _engine = EasyOCREngine.instance()
    return _engine