diff --git a/headroom/image/compressor.py b/headroom/image/compressor.py index 67080338e..f9ba35a78 100644 --- a/headroom/image/compressor.py +++ b/headroom/image/compressor.py @@ -36,6 +36,67 @@ from .trained_router import Technique logger = logging.getLogger(__name__) +# OCR backend resolution — see issue #372. +# +# After version 1.4.x the rapidocr ecosystem split: +# * rapidocr-onnxruntime — bundled-ORT, capped at Python <3.13. +# * rapidocr 3.x — engine-agnostic core, supports 3.13+; +# requires `onnxruntime` installed alongside +# for the same ORT backend; returns a +# RapidOCROutput dataclass instead of a tuple. +# +# We try v1 first (legacy / Python <3.13 install path), fall back to +# v3 (Python 3.13+ install path), and cache the resolved tuple at +# module scope. Result is intentionally None when neither package is +# installed — OCR is an optional capability gated by `[image]` extra. +_RESOLVED_OCR: tuple[Any | None, str | None] | None = None + + +def _resolve_rapidocr() -> tuple[Any | None, str | None]: + """Return ``(RapidOCR class, api_version)`` cached on first call. + + ``api_version`` is ``"v1"`` for ``rapidocr_onnxruntime`` (tuple + result shape) and ``"v3"`` for ``rapidocr`` 3.x (dataclass result + shape). Returns ``(None, None)`` when neither package is installed. + + Detection is at runtime (not based on Python version) because a + user on Python 3.11 might choose to install the 3.x package, and + a future ABI3 ORT release may make rapidocr-onnxruntime work on + Python 3.13. The actual install state is the source of truth. + """ + global _RESOLVED_OCR + if _RESOLVED_OCR is not None: + return _RESOLVED_OCR + + try: + from rapidocr_onnxruntime import RapidOCR as _RapidOCRv1 + + _RESOLVED_OCR = (_RapidOCRv1, "v1") + return _RESOLVED_OCR + except ImportError: + pass + + try: + from rapidocr import RapidOCR as _RapidOCRv3 # type: ignore[import-not-found] + + _RESOLVED_OCR = (_RapidOCRv3, "v3") + return _RESOLVED_OCR + except ImportError: + pass + + _RESOLVED_OCR = (None, None) + return _RESOLVED_OCR + + +def _reset_resolved_ocr_for_tests() -> None: + """Test-only hook: clear the module-level resolver cache so each + test can re-monkeypatch ``sys.modules`` and exercise a fresh + resolution. Production code never calls this. + """ + global _RESOLVED_OCR + _RESOLVED_OCR = None + + @dataclass class CompressionResult: """Result of image compression.""" @@ -308,44 +369,121 @@ class ImageCompressor: def _ocr_extract(self, image_data: bytes, min_confidence: float = 0.7) -> str | None: """Extract text from image using RapidOCR. - Returns extracted text if OCR is confident, None otherwise (fallback to image). + Adapts both API generations of the rapidocr ecosystem at runtime + (issue #372): + + * ``rapidocr-onnxruntime`` 1.4.x (Python <3.13) — call returns + ``(list[(box, text, score)], elapsed)``. + * ``rapidocr`` 3.x (Python 3.13+) — call returns a + ``RapidOCROutput`` dataclass with ``.txts`` (list[str]), + ``.scores`` (list[float]), ``.boxes`` (list); each may be + ``None`` when nothing was detected. + + Returns extracted text if OCR is confident, ``None`` otherwise + (caller falls back to image-as-image). """ - try: - from rapidocr_onnxruntime import RapidOCR + ocr_cls, api_version = _resolve_rapidocr() + if ocr_cls is None: + logger.debug( + "OCR backend unavailable: neither rapidocr-onnxruntime nor " + "rapidocr installed — skipping (event=ocr_backend_missing)" + ) + return None - if not hasattr(self, "_ocr_engine"): - self._ocr_engine = RapidOCR() - - result, _ = self._ocr_engine(image_data) - if not result: - return None - - # Check average confidence - confidences = [line[2] for line in result] - avg_confidence = sum(confidences) / len(confidences) - - if avg_confidence < min_confidence: - logger.debug( - f"OCR confidence too low ({avg_confidence:.0%} < {min_confidence:.0%}), " - f"falling back to image" + if not hasattr(self, "_ocr_engine"): + try: + self._ocr_engine = ocr_cls() + except Exception as exc: + logger.warning( + "OCR engine init failed (event=ocr_engine_init_failed, api=%s): %s", + api_version, + exc, ) return None - # Combine lines into text - text = "\n".join(line[1] for line in result) - - logger.info( - f"OCR extracted {len(result)} lines ({avg_confidence:.0%} avg confidence, " - f"{len(text)} chars)" + try: + raw = self._ocr_engine(image_data) + except Exception as exc: + logger.warning( + "OCR call failed (event=ocr_call_failed, api=%s): %s", + api_version, + exc, ) - return text + return None - except ImportError: - logger.debug("rapidocr-onnxruntime not installed, skipping OCR") + if api_version == "v1": + # 1.x returns (list_of_tuples, elapsed). list may be empty + # or None when no text is detected. + try: + result, _elapsed = raw + except (TypeError, ValueError): + logger.warning( + "OCR returned unexpected v1 shape (event=ocr_unknown_api_shape, api=v1): %r", + type(raw).__name__, + ) + return None + if not result: + return None + try: + texts = [line[1] for line in result] + confidences = [line[2] for line in result] + except (IndexError, TypeError): + logger.warning( + "OCR v1 result rows missing expected (box, text, score) " + "shape (event=ocr_unknown_api_shape, api=v1)" + ) + return None + + elif api_version == "v3": + # 3.x returns RapidOCROutput with txts/scores attributes. + # Both are None when detection found nothing — coerce to []. + texts_attr = getattr(raw, "txts", None) + scores_attr = getattr(raw, "scores", None) + if texts_attr is None and scores_attr is None: + # Probe failed to detect anything — not an error. + return None + texts = list(texts_attr or []) + confidences = list(scores_attr or []) + if not texts: + return None + if len(confidences) != len(texts): + logger.warning( + "OCR v3 returned mismatched txts/scores lengths " + "(event=ocr_unknown_api_shape, api=v3, txts=%d, scores=%d)", + len(texts), + len(confidences), + ) + return None + + else: + logger.warning( + "OCR resolver returned unknown api_version (event=ocr_unknown_api_shape, api=%r)", + api_version, + ) return None - except Exception as e: - logger.warning(f"OCR failed: {e}") + + if not confidences: return None + avg_confidence = sum(confidences) / len(confidences) + if avg_confidence < min_confidence: + logger.debug( + "OCR confidence too low (event=ocr_low_confidence, " + "avg=%.2f, min=%.2f, api=%s) — falling back to image", + avg_confidence, + min_confidence, + api_version, + ) + return None + + text = "\n".join(texts) + logger.info( + "OCR extracted %d lines (event=ocr_extracted, avg_confidence=%.2f, chars=%d, api=%s)", + len(texts), + avg_confidence, + len(text), + api_version, + ) + return text def _apply_compression( self, diff --git a/pyproject.toml b/pyproject.toml index 150af2c0d..c89af3ff6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -110,10 +110,30 @@ relevance = [ "numpy>=1.24.0", ] # Image compression (ML-based routing + OCR) +# +# OCR backend uses ONNX Runtime regardless of Python version. The +# rapidocr ecosystem split into two flavors after 1.4.x: +# * rapidocr-onnxruntime 1.4.x — bundled-ORT package, capped at +# Python <3.13 by its requires-python metadata. Drop-in for our +# existing v1 tuple-shaped API call. +# * rapidocr 3.x — engine-agnostic core, supports Python 3.13+. +# Returns a RapidOCROutput dataclass (txts, scores, boxes, ...). +# Needs `onnxruntime` installed separately to use the ORT backend. +# +# `headroom/image/compressor.py` adapts both API shapes at runtime via +# a try/except cascade. See issue #372 for context. image = [ "pillow>=10.0.0", "sentencepiece>=0.1.99", # Required by SigLIP tokenizer (SiglipTokenizer) - "rapidocr-onnxruntime>=1.4.0", # ONNX-native OCR for text extraction from images (~15MB models) + # Python 3.6–3.12: keep the proven ORT-bundled package directly. + # ~15 MB ONNX models auto-downloaded on first use. + "rapidocr-onnxruntime>=1.4.0,<2; python_version<'3.13'", + # Python 3.13+: rapidocr-onnxruntime is unavailable (its wheels + # declare requires-python<3.13). Use the successor `rapidocr` 3.x + # core + `onnxruntime` engine; same ORT backend, just split into + # two packages. Total install size and inference speed unchanged. + "rapidocr>=3.0,<4; python_version>='3.13'", + "onnxruntime>=1.7,<2; python_version>='3.13'", ] # Report generation reports = [ diff --git a/tests/test_image_ocr_api_compat.py b/tests/test_image_ocr_api_compat.py new file mode 100644 index 000000000..7459021ae --- /dev/null +++ b/tests/test_image_ocr_api_compat.py @@ -0,0 +1,235 @@ +"""OCR backend API-compat regression tests for issue #372. + +The rapidocr ecosystem split after 1.4.x: + +* `rapidocr-onnxruntime` 1.4.x — Python <3.13 only; tuple result. +* `rapidocr` 3.x — Python 3.13+; `RapidOCROutput` dataclass result. + +`headroom/image/compressor.py` adapts both at runtime via +`_resolve_rapidocr` + per-version branches in `_ocr_extract`. These +tests pin both branches so a future "let me clean up the v1 path" +refactor doesn't silently break Python <3.13 users. +""" + +from __future__ import annotations + +import sys +from types import ModuleType, SimpleNamespace +from typing import Any + +import pytest + +from headroom.image import compressor as compressor_module +from headroom.image.compressor import ImageCompressor + + +def _install_fake_module(monkeypatch: pytest.MonkeyPatch, name: str, attrs: dict[str, Any]) -> None: + """Inject a synthetic module into sys.modules so import sees it.""" + mod = ModuleType(name) + for k, v in attrs.items(): + setattr(mod, k, v) + monkeypatch.setitem(sys.modules, name, mod) + + +def _hide_module(monkeypatch: pytest.MonkeyPatch, name: str) -> None: + """Force ImportError when `name` is imported.""" + monkeypatch.setitem(sys.modules, name, None) + + +@pytest.fixture(autouse=True) +def _reset_resolver_cache() -> None: + compressor_module._reset_resolved_ocr_for_tests() + yield + compressor_module._reset_resolved_ocr_for_tests() + + +# --------------------------------------------------------------------------- +# Resolver +# --------------------------------------------------------------------------- + + +def test_resolve_rapidocr_prefers_v1_when_both_available(monkeypatch: pytest.MonkeyPatch) -> None: + class _V1RapidOCR: + pass + + class _V3RapidOCR: + pass + + _install_fake_module(monkeypatch, "rapidocr_onnxruntime", {"RapidOCR": _V1RapidOCR}) + _install_fake_module(monkeypatch, "rapidocr", {"RapidOCR": _V3RapidOCR}) + + cls, api = compressor_module._resolve_rapidocr() + assert cls is _V1RapidOCR + assert api == "v1" + + +def test_resolve_rapidocr_falls_back_to_v3_when_v1_missing(monkeypatch: pytest.MonkeyPatch) -> None: + class _V3RapidOCR: + pass + + _hide_module(monkeypatch, "rapidocr_onnxruntime") + _install_fake_module(monkeypatch, "rapidocr", {"RapidOCR": _V3RapidOCR}) + + cls, api = compressor_module._resolve_rapidocr() + assert cls is _V3RapidOCR + assert api == "v3" + + +def test_resolve_rapidocr_returns_none_when_neither_installed( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _hide_module(monkeypatch, "rapidocr_onnxruntime") + _hide_module(monkeypatch, "rapidocr") + + cls, api = compressor_module._resolve_rapidocr() + assert cls is None + assert api is None + + +# --------------------------------------------------------------------------- +# _ocr_extract — v1 tuple shape +# --------------------------------------------------------------------------- + + +def _make_compressor() -> ImageCompressor: + """Build an ImageCompressor with the heavy ML deps stubbed. + + `ImageCompressor.__init__` lazy-loads the trained router; we don't + exercise it here. Constructing a bare instance bypasses the load. + """ + return ImageCompressor.__new__(ImageCompressor) + + +def test_ocr_extract_v1_tuple_shape_parses_correctly(monkeypatch: pytest.MonkeyPatch) -> None: + class _V1Engine: + def __call__(self, _image_data: bytes) -> tuple[list[tuple[Any, str, float]], float]: + return ( + [ + (None, "hello", 0.95), + (None, "world", 0.90), + ], + 0.123, + ) + + class _V1RapidOCR: + def __init__(self) -> None: ... + def __call__(self, image_data: bytes) -> Any: # noqa: D401 + return _V1Engine()(image_data) + + _install_fake_module(monkeypatch, "rapidocr_onnxruntime", {"RapidOCR": _V1RapidOCR}) + + c = _make_compressor() + text = c._ocr_extract(b"\x89PNG fake") + assert text == "hello\nworld" + + +def test_ocr_extract_v1_low_confidence_returns_none(monkeypatch: pytest.MonkeyPatch) -> None: + class _V1RapidOCR: + def __init__(self) -> None: ... + def __call__(self, _image_data: bytes) -> tuple[list[Any], float]: + return ([(None, "blurry", 0.3), (None, "smudge", 0.4)], 0.0) + + _install_fake_module(monkeypatch, "rapidocr_onnxruntime", {"RapidOCR": _V1RapidOCR}) + + c = _make_compressor() + assert c._ocr_extract(b"\x89PNG fake", min_confidence=0.7) is None + + +def test_ocr_extract_v1_empty_result_returns_none(monkeypatch: pytest.MonkeyPatch) -> None: + class _V1RapidOCR: + def __init__(self) -> None: ... + def __call__(self, _image_data: bytes) -> tuple[list[Any] | None, float]: + return ([], 0.0) + + _install_fake_module(monkeypatch, "rapidocr_onnxruntime", {"RapidOCR": _V1RapidOCR}) + + c = _make_compressor() + assert c._ocr_extract(b"\x89PNG fake") is None + + +# --------------------------------------------------------------------------- +# _ocr_extract — v3 dataclass shape +# --------------------------------------------------------------------------- + + +def test_ocr_extract_v3_dataclass_shape_parses_correctly(monkeypatch: pytest.MonkeyPatch) -> None: + class _V3RapidOCR: + def __init__(self) -> None: ... + def __call__(self, _image_data: bytes) -> Any: + # Mirror the real `rapidocr.RapidOCROutput` minimal surface. + return SimpleNamespace( + txts=["hello", "world"], + scores=[0.95, 0.90], + boxes=[None, None], + ) + + _hide_module(monkeypatch, "rapidocr_onnxruntime") + _install_fake_module(monkeypatch, "rapidocr", {"RapidOCR": _V3RapidOCR}) + + c = _make_compressor() + text = c._ocr_extract(b"\x89PNG fake") + assert text == "hello\nworld" + + +def test_ocr_extract_v3_low_confidence_returns_none(monkeypatch: pytest.MonkeyPatch) -> None: + class _V3RapidOCR: + def __init__(self) -> None: ... + def __call__(self, _image_data: bytes) -> Any: + return SimpleNamespace(txts=["blurry", "smudge"], scores=[0.3, 0.4], boxes=[None, None]) + + _hide_module(monkeypatch, "rapidocr_onnxruntime") + _install_fake_module(monkeypatch, "rapidocr", {"RapidOCR": _V3RapidOCR}) + + c = _make_compressor() + assert c._ocr_extract(b"\x89PNG fake", min_confidence=0.7) is None + + +def test_ocr_extract_v3_none_attrs_when_no_text_detected(monkeypatch: pytest.MonkeyPatch) -> None: + """Real-world v3 behavior: when detection finds nothing, RapidOCROutput + has txts=None and scores=None. Verified in dispatch smoke test against + rapidocr 3.8.1. + """ + + class _V3RapidOCR: + def __init__(self) -> None: ... + def __call__(self, _image_data: bytes) -> Any: + return SimpleNamespace(txts=None, scores=None, boxes=None) + + _hide_module(monkeypatch, "rapidocr_onnxruntime") + _install_fake_module(monkeypatch, "rapidocr", {"RapidOCR": _V3RapidOCR}) + + c = _make_compressor() + assert c._ocr_extract(b"\x89PNG fake") is None + + +def test_ocr_extract_v3_mismatched_lengths_logs_and_returns_none( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + class _V3RapidOCR: + def __init__(self) -> None: ... + def __call__(self, _image_data: bytes) -> Any: + return SimpleNamespace(txts=["a", "b", "c"], scores=[0.9, 0.8], boxes=[None] * 3) + + _hide_module(monkeypatch, "rapidocr_onnxruntime") + _install_fake_module(monkeypatch, "rapidocr", {"RapidOCR": _V3RapidOCR}) + + c = _make_compressor() + with caplog.at_level("WARNING", logger="headroom.image.compressor"): + result = c._ocr_extract(b"\x89PNG fake") + assert result is None + assert any("event=ocr_unknown_api_shape" in r.getMessage() for r in caplog.records) + + +# --------------------------------------------------------------------------- +# Backend missing +# --------------------------------------------------------------------------- + + +def test_ocr_extract_returns_none_when_no_backend_installed( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _hide_module(monkeypatch, "rapidocr_onnxruntime") + _hide_module(monkeypatch, "rapidocr") + + c = _make_compressor() + assert c._ocr_extract(b"\x89PNG fake") is None