agentic-mobile-control/perception/ocr.py

from __future__ import annotations

import os
import tempfile
import logging
from collections.abc import Iterable
from dataclasses import dataclass
from numbers import Real
from pathlib import Path
from typing import Any

from core.models import Bounds, SceneElement

logger = logging.getLogger(__name__)


@dataclass(frozen=True)
class OCRBox:
    text: str
    bounds: Bounds
    confidence: float | None = None

    def to_scene_element(self, element_id: str) -> SceneElement:
        return SceneElement(
            id=element_id,
            type="text",
            text=self.text,
            bounds=self.bounds,
            confidence=self.confidence,
            source="ocr",
        )


class PaddleOCREngine:
    def __init__(self, **kwargs: Any) -> None:
        # Device screenshots are already upright, flat digital captures, not
        # photographed paper documents, so PaddleOCR's document-preprocessing
        # models (orientation classification + UVDoc unwarping) have nothing
        # real to correct. Left at their library defaults (True), they still
        # run, geometrically warp the image, and detect/recognize text against
        # that warped image, returning box coordinates that no longer line up
        # with the original screenshot. Callers can still opt back in via an
        # explicit kwarg.
        kwargs.setdefault("use_doc_orientation_classify", False)
        kwargs.setdefault("use_doc_unwarping", False)
        self.kwargs = kwargs
        self._engine: Any | None = None

    def extract(self, image: bytes | str | Path) -> list[OCRBox]:
        engine = self._load()
        image_input, temp_path = _image_input(image)
        try:
            if hasattr(engine, "predict"):
                raw = engine.predict(input=image_input)
            else:
                raw = engine.ocr(image_input, cls=True)
            return parse_paddle_result(raw)
        finally:
            if temp_path:
                temp_path.unlink(missing_ok=True)

    def _load(self) -> Any:
        if self._engine is None:
            from paddleocr import PaddleOCR

            self._engine = PaddleOCR(**self.kwargs)
        return self._engine


def run_ocr(
    image: bytes | str | Path,
    *,
    engine: PaddleOCREngine | None = None,
    strict: bool = False,
) -> list[SceneElement]:
    try:
        boxes = (engine or PaddleOCREngine()).extract(image)
    except Exception:
        if strict:
            raise
        logger.warning("OCR extraction failed; continuing without OCR", exc_info=True)
        return []
    return [box.to_scene_element(f"ocr-{index:03d}") for index, box in enumerate(boxes)]


def parse_paddle_result(raw: Any) -> list[OCRBox]:
    boxes: list[OCRBox] = []
    for item in _flatten_pages(raw):
        parsed = _parse_line(item)
        if parsed:
            boxes.append(parsed)
    return boxes


def _image_input(image: bytes | str | Path) -> tuple[str, Path | None]:
    if isinstance(image, bytes):
        handle = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
        try:
            handle.write(image)
        finally:
            handle.close()
        return handle.name, Path(handle.name)
    return os.fspath(image), None


def _flatten_pages(raw: Any) -> Iterable[Any]:
    if raw is None:
        return []
    if isinstance(raw, dict):
        return _dict_lines(raw)
    if isinstance(raw, list):
        flattened: list[Any] = []
        for page in raw:
            if isinstance(page, dict):
                flattened.extend(_dict_lines(page))
            elif _looks_like_ocr_line(page):
                flattened.append(page)
            elif isinstance(page, list):
                flattened.extend(page)
        return flattened
    json_attr = getattr(raw, "json", None)
    if callable(json_attr):
        return _flatten_pages(json_attr())
    if isinstance(json_attr, (dict, list)):
        return _flatten_pages(json_attr)
    return []


def _dict_lines(data: dict[str, Any]) -> list[Any]:
    result = data.get("res")
    payload = result if isinstance(result, dict) else data
    texts = _first_nonempty(payload, "rec_texts", "texts")
    scores = _first_nonempty(payload, "rec_scores", "scores")
    boxes = _first_nonempty(payload, "rec_boxes", "rec_polys", "dt_polys")
    return [
        {
            "text": text,
            "confidence": scores[index] if index < len(scores) else None,
            "points": boxes[index] if index < len(boxes) else None,
        }
        for index, text in enumerate(texts)
    ]


def _looks_like_ocr_line(value: Any) -> bool:
    return isinstance(value, (list, tuple)) and len(value) >= 2


def _parse_line(line: Any) -> OCRBox | None:
    if isinstance(line, dict):
        text = line.get("text")
        points = _first_nonempty(line, "points", "box", "bounds")
        confidence = line.get("confidence")
        if not _has_items(text) or not _has_items(points):
            return None
        return OCRBox(
            str(text), _bounds_from_points(points), _float_or_none(confidence)
        )

    if not _looks_like_ocr_line(line):
        return None

    points = line[0]
    text_payload = line[1]
    if isinstance(text_payload, (list, tuple)) and text_payload:
        text = text_payload[0]
        confidence = text_payload[1] if len(text_payload) > 1 else None
    else:
        text = text_payload
        confidence = None
    if not text:
        return None
    return OCRBox(str(text), _bounds_from_points(points), _float_or_none(confidence))


def _bounds_from_points(points: Any) -> Bounds:
    if isinstance(points, dict):
        return Bounds.from_dict(points)
    if _has_length(points, 4) and all(isinstance(value, Real) for value in points):
        x1, y1, x2, y2 = [float(value) for value in points]
        return Bounds(x1, y1, x2 - x1, y2 - y1)

    xs: list[float] = []
    ys: list[float] = []
    for point in points:
        if isinstance(point, dict):
            xs.append(float(point["x"]))
            ys.append(float(point["y"]))
        else:
            xs.append(float(point[0]))
            ys.append(float(point[1]))
    return Bounds(min(xs), min(ys), max(xs) - min(xs), max(ys) - min(ys))


def _float_or_none(value: Any) -> float | None:
    if value is None:
        return None
    return float(value)


def _first_nonempty(data: dict[str, Any], *keys: str) -> Any:
    for key in keys:
        value = data.get(key)
        if _has_items(value):
            return value
    return []


def _has_items(value: Any) -> bool:
    if value is None:
        return False
    try:
        return len(value) > 0
    except TypeError:
        return bool(value)


def _has_length(value: Any, length: int) -> bool:
    try:
        return len(value) == length
    except TypeError:
        return False