ocr-foundry 0.2.0__py3-none-any.whl

khocr_gen/__init__.py

@@ -0,0 +1,11 @@
+"""OCR Foundry - synthetic OCR training data generator."""
+
+from importlib.metadata import PackageNotFoundError, version
+
+try:
+    __version__ = version("ocr-foundry")
+except PackageNotFoundError:
+    try:
+        __version__ = version("khocr-gen")
+    except PackageNotFoundError:
+        __version__ = "0.2.0"

khocr_gen/_pil_compat.py

@@ -0,0 +1,25 @@
+"""PIL compatibility shims for resampling across Pillow versions."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from PIL import Image
+
+
+def _resolve_resample(name: str) -> Any:
+    resampling = getattr(Image, "Resampling", None)
+    if resampling is not None:
+        value = getattr(resampling, name, None)
+        if value is not None:
+            return value
+
+    legacy_value = getattr(Image, name, None)
+    if legacy_value is not None:
+        return legacy_value
+
+    raise AttributeError(f"PIL.Image does not define resample mode '{name}'")
+
+
+RESAMPLE_BILINEAR: Any = _resolve_resample("BILINEAR")
+RESAMPLE_BICUBIC: Any = _resolve_resample("BICUBIC")

khocr_gen/_rust_accel.py

@@ -0,0 +1,204 @@
+"""Rust acceleration backend for khocr-gen.
+
+This module wraps the native `_rust_accel` extension (built via maturin/PyO3).
+If the native extension is unavailable, Python augmentation and text splitting fallbacks are used.
+
+Usage::
+
+    from khocr_gen._rust_accel import (
+        apply_sauvola, apply_blur,
+        split_text_spans,
+        estimate_bg,
+        HAS_RUST_ACCEL,
+    )
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+_LOGGER = logging.getLogger("khocr_gen._rust_accel")
+
+# Try to load the native Rust extension
+
+_native: Any = None
+
+try:
+    import _rust_accel as _native
+
+    HAS_RUST_ACCEL = True
+except ImportError:
+    HAS_RUST_ACCEL = False
+    _LOGGER.info("Rust extension not available, using pure-Python fallbacks.")
+
+
+def _log_available() -> None:
+    if HAS_RUST_ACCEL:
+        try:
+            ver: str = getattr(_native, "__version__", "unknown")
+        except Exception:
+            ver = "unknown"
+        _LOGGER.info("Rust acceleration enabled (native _rust_accel v%s)", ver)
+    else:
+        _LOGGER.info("Rust acceleration NOT available.")
+
+
+# Augmentation functions
+
+# Each entry: (function_name, has_rgb_variant)
+_AUG_FUNCTIONS: list[tuple[str, bool]] = [
+    ("apply_sauvola", False),
+    ("apply_geo_warp", False),
+    ("apply_vertical_crop", False),
+    ("apply_blur", False),
+    ("apply_salt_pepper", False),
+    ("apply_background_texture", False),
+    ("apply_jpeg_compression", False),
+    ("apply_rotation", True),
+    ("apply_lowdpi", False),
+    ("apply_oversample", False),
+    ("apply_perspective", False),
+    ("apply_elastic", False),
+    ("apply_random_crop", False),
+    ("apply_online_blur", False),
+    ("apply_online_noise", False),
+    ("apply_reverse", True),
+    ("apply_brightness_contrast", True),
+    ("apply_pixelation", False),
+    ("apply_gradient_illumination", False),
+    ("apply_morphological", False),
+    ("apply_anisotropic_dilation", False),
+    ("apply_hsv", True),  # RGB-only in Rust
+]
+
+
+def _get_rust_aug_fn(name: str) -> Callable[..., Any] | None:
+    """Return the native Rust augmentation function by name, or None."""
+    if not HAS_RUST_ACCEL:
+        return None
+    assert _native is not None
+    return getattr(_native, f"{name}_rust", None)
+
+
+def _make_aug_fallback(name: str) -> Callable[..., Any]:
+    """Create a fallback that delegates to the Python augmentation module."""
+
+    def _fallback(img: np.ndarray, intensity: float, **kwargs: Any) -> np.ndarray:
+        from khocr_gen.augmentation import AUG_METHODS
+
+        fn = AUG_METHODS.get(name)
+        if fn is None:
+            return img
+        return fn(img, intensity, **kwargs)
+
+    _fallback.__name__ = name
+    _fallback.__qualname__ = name
+    return _fallback
+
+
+def _export_aug_functions() -> None:
+    """Populate the module globals with augmentation functions.
+
+    For each augmentation method, the Rust variant is preferred when available;
+    otherwise the pure-Python fallback is used.
+
+    Grayscale-only methods dispatch directly.
+    Methods with an RGB variant (`brightness_contrast`, `hsv`) are wrapped to handle both grayscale and RGB input arrays.
+    """
+    g = globals()
+
+    for name, has_rgb in _AUG_FUNCTIONS:
+        rust_fn = _get_rust_aug_fn(name)
+        rust_rgb_fn: Callable[..., Any] | None = (
+            _get_rust_aug_fn(f"{name}_rgb") if has_rgb else None
+        )
+
+        if has_rgb and HAS_RUST_ACCEL and (rust_fn is not None or rust_rgb_fn is not None):
+            # Wrap to auto-detect channel count and dispatch.
+            # Rust RGB functions return flat (h, w*3) 2D arrays; reshape to (h, w, 3).
+            _gray_fn = rust_fn
+            _rgb_fn = rust_rgb_fn
+
+            # Capture the fallback factory by value so we can safely delete it later
+            _fallback_factory = _make_aug_fallback
+
+            def _make_rgb_wrapper(
+                gray_fn: Callable[..., Any] | None,
+                rgb_fn: Callable[..., Any] | None,
+                fallback_name: str,
+                fallback_factory: Callable[..., Any],
+            ) -> Callable[..., Any]:
+                def _wrapper(img: np.ndarray, intensity: float, **kwargs: Any) -> np.ndarray:
+                    if img.ndim == 3 and img.shape[2] == 3:
+                        if rgb_fn is not None:
+                            h, w = img.shape[:2]
+                            result = rgb_fn(img, intensity)
+                            # Rust RGB fns return flat (h, w*3); reshape to (h, w, 3)
+                            if result.ndim == 2 and result.shape[1] == w * 3:
+                                return result.reshape(h, w, 3)
+                            return result
+                        if gray_fn is not None:
+                            channels = [
+                                gray_fn(np.ascontiguousarray(img[:, :, c]), intensity)
+                                for c in range(3)
+                            ]
+                            return np.stack(channels, axis=-1)
+                        return fallback_factory(fallback_name)(img, intensity)
+                    if gray_fn is not None:
+                        return gray_fn(img, intensity)
+                    return fallback_factory(fallback_name)(img, intensity)
+
+                return _wrapper
+
+            g[name] = _make_rgb_wrapper(_gray_fn, _rgb_fn, name, _fallback_factory)
+        elif rust_fn is not None:
+            g[name] = rust_fn
+        else:
+            g[name] = _make_aug_fallback(name)
+
+
+_export_aug_functions()
+
+# Clean up the helpers so they don't pollute the module namespace.
+# _make_aug_fallback and _AUG_FUNCTIONS are retained for closures referencing them.
+del _export_aug_functions
+
+
+# Text span splitting
+
+if HAS_RUST_ACCEL:
+    assert _native is not None
+    split_text_spans = _native.split_text_spans  # type: ignore[assignment]
+else:
+
+    def split_text_spans(text: str) -> list[tuple[str, str]]:  # type: ignore[no-redef]
+        """Split text into contiguous script spans (pure-Python fallback)."""
+        from khocr_gen.rendering import ImageRenderer
+
+        return ImageRenderer._split_text_into_spans(text)
+
+
+# Utilities
+
+if HAS_RUST_ACCEL:
+    assert _native is not None
+    estimate_bg = _native.estimate_bg_rust  # type: ignore[assignment]
+else:
+
+    def estimate_bg(img: np.ndarray) -> int:  # type: ignore[no-redef]
+        """Estimate background colour from image border (pure-Python fallback)."""
+        from khocr_gen.augmentation import _estimate_bg
+
+        return int(_estimate_bg(img))
+
+
+# Module-level init
+
+_log_available()
+del _log_available