akadmin/headroom
1
0
Fork 0
mirror of https://github.com/headroomlabs-ai/headroom.git synced 2026-08-27 14:17:10 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix: PR #372 — restore [image] extra on Python 3.13 via rapidocr 3.x adapter

Browse source 
Root cause: `headroom-ai[all]==0.20.16` fails to install on Python 3.13
because `rapidocr-onnxruntime` 1.4.0–1.4.4 wheels declare
`requires-python: <3.13,>=3.6`. After 1.4.x the rapidocr ecosystem
split: `rapidocr-onnxruntime` (bundled-ORT, capped at <3.13) vs
`rapidocr` 3.x (engine-agnostic core, supports 3.13+, returns
RapidOCROutput dataclass instead of v1's tuple).

Fix:

1. pyproject.toml — environment-marker hybrid in [image]:
   - rapidocr-onnxruntime>=1.4.0,<2; python_version<'3.13'
   - rapidocr>=3.0,<4; python_version>='3.13'
   - onnxruntime>=1.7,<2; python_version>='3.13'
   ORT remains the engine on every Python version; bundle and speed
   unchanged, just split into two packages on 3.13+.

2. headroom/image/compressor.py — runtime adapter:
   _resolve_rapidocr() tries v1 first, falls back to v3 when v1 is
   missing, returns (None, None) when neither installed. Cached at
   module scope. Detection at runtime (not Python-version-based) so
   users can install either package on any Python version.

   _ocr_extract branches on resolved api_version:
   - v1: (list[(box, text, score)], elapsed) tuple — unchanged
   - v3: RapidOCROutput dataclass with .txts / .scores / .boxes
     attrs (each may be None when nothing detected)

   Defensive None-handling, length-mismatch detection, structured
   log events for both branches.

Smoke test (real install verified before commit):
    pip install rapidocr onnxruntime pillow
    → result type: RapidOCROutput
    → fields: txts (None when empty), scores (None when empty), boxes
    Confirms the v3 None-coercion is necessary.

Tests: 11 new unit tests in tests/test_image_ocr_api_compat.py covering:
- Resolver: v1 preferred, v3 fallback, both missing
- v1 path: tuple parses, low-confidence None, empty result None
- v3 path: dataclass parses, low-confidence None, None attrs handled,
  mismatched lengths logged + None
- Backend missing: returns None gracefully

All 11 pass; `make ci-precheck` PASSED.

Closes #372.
This commit is contained in:
chopratejas 2026-05-04 08:20:01 -07:00
parent e731f87c5d
commit b154e17853
3 changed files with 423 additions and 30 deletions
Download patch file Download diff file Expand all files Collapse all files

196
headroom/image/compressor.py
View file

@ -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,

22
pyproject.toml
View file

@ -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 = [

235
tests/test_image_ocr_api_compat.py Normal file
View file

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