Skip to content

Rust Acceleration

LazyGreed edited this page Nov 11, 2026 · 3 revisions

Rust Acceleration

khocr-gen ships an optional native extension, _rust_accel, that speeds up the hottest paths in the generation pipeline: pixel-level augmentation and font glyph checking.

It is built with PyO3/maturin from the rust/ crate (package name khocr-gen-core) and installed as a regular editable dependency of the Python project. There is no separate build step.

Every accelerated code path has a pure-Python fallback, so the project works with or without a Rust toolchain. Acceleration is detected once at import time:

from khocr_gen import _rust_accel as ra
print(ra.HAS_RUST_ACCEL)   # True when the native extension loaded

Installation

uv sync                           # install runtime dependencies (pure-Python)
uv sync --extra dev               # development dependencies (pytest, ruff, ty)
uv sync --extra dev --extra rust  # also compiles and installs khocr-gen-core from rust/

Building khocr-gen-core requires a Rust toolchain (cargo/rustc, edition 2021) on PATH. Running uv sync --extra dev alone installs development tools without requiring a Rust toolchain; all accelerated operations fall back to pure-Python implementations. Installing Rust (https://rustup.rs) enables building with --extra rust.

Check whether acceleration is active:

uv run python -c "from khocr_gen import _rust_accel as ra; print(ra.HAS_RUST_ACCEL)"

What's accelerated

Augmentation — 21 of 27 methods

khocr_gen.augmentation.AUG_METHODS transparently swaps in the native implementation for each method below when HAS_RUST_ACCEL is true; otherwise the pure-Python/OpenCV implementation is used.

Accelerated: sauvola, geo_warp, vertical_crop, blur, jpeg_compression, rotation, salt_pepper, background_texture, lowdpi, oversample, perspective, elastic, random_crop, online_blur, online_noise, hsv, reverse, brightness_contrast, pixelation, morphological, anisotropic_dilation.

Not accelerated:

Method Why
distortion, albu_noise Thin wrappers around Albumentations' own compiled pipeline — no native win to be had
gradient_illumination The Rust implementation is grayscale-only, but the renderer may pass RGB arrays through _RGB_PREFERRED_METHODS, so it's excluded until an RGB path is added
extreme_resize No native implementation yet
low_contrast_caption No native implementation yet
luminous_halo Colored bloom is Python-only; a grayscale native path would drop the hue

Font coverage

Python manages font registration, cached file hashes and embedded license records. Multilingual coverage uses fontTools cmap inspection and fails closed on unreadable faces; Pillow/RAQM shapes and renders text. Legacy mode retains cached raster/tofu checks. Font handling works with or without the Rust extension.

Script-span splitting (split_text_spans)

ImageRenderer._render_mixed_font uses _rust_accel.split_text_spans to split mixed Khmer/English text into contiguous per-script spans, replacing the equivalent pure-Python ImageRenderer._split_text_into_spans (kept as the fallback).

Utilities

estimate_bg (background-colour estimation) has a native implementation wired into augmentation.py's helper calls.

Module layout

rust/
├── Cargo.toml           # crate: khocr-gen-core, cdylib, PyO3 abi3-py312
├── pyproject.toml       # maturin build config (module-name = "_rust_accel")
└── src/
    ├── lib.rs           # PyO3 bindings + #[pymodule] registration (the only file Python sees)
    ├── augmentation.rs  # 20 accelerated augmentation methods + RGB variants
    └── utils.rs         # background estimation, text span splitting and kernel helpers

The project-root pyproject.toml declares khocr-gen-core as a path dependency via [tool.uv.sources], pointing at rust/ with editable = true, under both the rust and dev optional-dependency extras. That is what makes uv sync build and keep it installed automatically.

Iterating on the Rust crate

For fast edit/build/test cycles on rust/src/*.rs without going through uv sync each time, build directly against the project's existing .venv:

VIRTUAL_ENV=/path/to/ocr-foundry/.venv uv run --with maturin maturin develop --release --manifest-path rust/Cargo.toml

Do not cd rust && uv run --with maturin maturin develop. rust/pyproject.toml declares its own (maturin) project, so a bare uv run from inside rust/ treats it as a separate project root and creates a stray rust/.venv. The extension gets installed there instead of the real project environment, and it silently looks like the build "isn't taking effect". Always pin VIRTUAL_ENV to the top-level .venv (or run from the repo root with --manifest-path) when building manually.

Also note: uv sync re-resolves and re-installs declared dependencies on every invocation. If you hand-build with maturin develop for a quick iteration, a later uv sync may reinstall its own cached build over yours. Run uv cache clean khocr-gen-core if uv sync seems to reinstall a stale binary (for example a build missing a class or function you just added).

Verifying acceleration is active

uv run python -c "from khocr_gen import _rust_accel as ra; print(ra.HAS_RUST_ACCEL)"

True means the extension is loaded and the registry swapped in the native methods. Generation output is equivalent either way, because the fallbacks implement the same behaviour. Pixel-identical output between Rust and Python paths is not guaranteed for methods that draw from their own thread RNG. See the note in Augmentation about verify reproducibility.

There is also a timing harness at the repo root:

uv run python benchmark_rust_vs_python_aug.py

CI

.github/workflows/ci.yml runs two jobs:

  • python — ruff check, ty check (non-blocking), and pytest --cov on a Python 3.12/3.13 matrix, after uv sync --extra dev builds the extension via maturin.
  • rust — cargo fmt --check (non-blocking), cargo clippy -D warnings (non-blocking), and cargo build --release (the real gate) in rust/.

See also

Clone this wiki locally