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: 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