Architecture

Architecture

This page is for people reading or changing the code. It maps the modules and traces the generation pipeline.

Repository layout

ocr-foundry/
├── src/khocr_gen/        # Python package
├── rust/                 # optional khocr-gen-core PyO3 extension
├── configs/              # example YAML recipes
├── docs/                 # in-repository references and archived proposals
├── tests/                # pytest suite
├── typings/              # type stubs used by `ty`
└── .github/workflows/    # CI and publishing workflows

Module map

Module	Responsibility
cli.py	Top-level entrypoint, subcommand dispatch, config-default injection
config.py	GenerationConfig and the generation option surface
config_loader.py	YAML loading, key normalisation, validation, and default injection
rendering.py	ImageRenderer: text rendering, decorations, effects, mixed-font compositing
augmentation.py	Augmentation registry and Python implementations
data_generator.py	Dataset generation, splits, output files, and vocabulary
verify.py	khocr-gen verify sample rendering
_rust_accel.py	Native extension loader and Python fallback shim
corpus.py	Corpus loading, filtering, counting, and character frequencies
fonts.py	Font discovery, pools, style tags, and face cache
parallel.py	Worker resolution, multiprocessing context, and chunk dispatch
normalizer.py	Khmer normalization wrappers and NormalizerConfig
combine.py / combine_cmd.py	Dataset merging
generate.py	khocr-gen generate command wiring
viewer.py	khocr-gen view LMDB preview and extraction
lmdb_pack.py	Image and label packing into LMDB
line_height.py	Height, font-scale, and padding sampling and validation
errors.py	KhocrGenError hierarchy
logging.py	Logging setup
_pil_compat.py	Pillow resampling-constant shims

Configuration is centralised

GenerationConfig (in config.py) is the single source of truth. It owns:

• add_args(parser): registers every CLI flag,
• from_args(args): builds the config from parsed arguments,
• iter_aug_methods() / enabled_aug_methods(): drives the augmentation registry,
• resolve_split_ratios(): normalises the train/val/test ratios,
• to_dict() / from_dict(): serialisation for worker processes.

Sub-configs are dataclasses: AugMethodConfig (with ExtremeResizeConfig overriding the pixel-height semantics), TextDecorationConfig, TextEffectConfig, and NormalizerConfig.

Errors

KhocrGenError
├── InputValidationError    # bad corpus / flags / config
├── GenerationError         # failures during generation
└── FontLoadError           # font discovery or parsing problems

Validation failures are raised as typed errors and surfaced as clean CLI messages rather than tracebacks.

The generation pipeline

corpus file
   │  corpus.py — load, filter by length, dedupe
   ▼
text lines
   │  data_generator.DatasetGenerator._split_lines — seeded shuffle, disjoint train/val/test
   ▼
per-split text lists
   │  (optional) rare-character oversampling — train only
   ▼
for each line × copies:
   ├── font selection        fonts.FontManager (script pool, style variants, glyph coverage)
   ├── clean render          rendering.ImageRenderer (decorations → effects → mixed-font spans)
   ├── height/font/padding   line_height.sample_*   (variable line-height modes)
   ├── resize to target      rendering._resize_to_target (uniform resize, never crop)
   └── one augmentation      augmentation.AUG_METHODS (weighted pick, sampled intensity)
   ▼
image + label
   │  data_generator — images/, labels.txt, metadata.jsonl, generation_errors.jsonl
   ▼
   ├── vocab.json            built by scanning the written labels files
   └── lmdb_pack.pack_lmdb   optional: image-N / label-N / num-samples keys

Key invariants:

1. One augmentation per image. Methods are never stacked; selection is a weighted draw across all methods with prob > 0.
2. Decorations and effects happen on the clean canvas, before augmentation, so an augmentation can never corrupt the glyph shapes that make the label correct.
3. Resize, never crop. Variable line height is achieved by resizing the naturally-sized canvas, so glyphs are never clipped.
4. Splits are disjoint at the text level. No text string appears in two splits.

Parallelism

Generation is chunked across worker processes:

Concern	Module	Behaviour
Worker count	parallel.resolve_worker_count	--workers 0 → auto; 1 when the workload is small (< 2048 samples), the host has one CPU, or --workers 1; otherwise cpu_count // 2. An explicit --workers N is clamped to cpu_count.
Chunk size	parallel.resolve_chunk_size	Derived from sample count and worker count, clamped to [64, 256] (_DEFAULT_RENDER_CHUNK_SIZE / _MAX_RENDER_CHUNK_SIZE). Serial runs use the default chunk size.
Start method	parallel.resolve_mp_context	Linux → forkserver (workers come from a single-threaded server, so a multithreaded parent cannot deadlock on fork); macOS/Windows → spawn. Override with KHOCR_GEN_MP_START_METHOD (forkserver, fork, or spawn).
Timeout	--worker-timeout	Seconds before a worker batch is considered timed out (default 300).

Config objects are passed to workers via GenerationConfig.to_dict() / from_dict(). from_dict silently ignores unknown keys, so workers started with an older config dict do not crash when new keys are added. That is deliberate, and worth preserving when you add fields.

_AUTO_PARALLEL_MIN_SAMPLES = 2048 in both augmentation.py and parallel.py gates the automatic switch to multiprocessing.

Rendering internals

ImageRenderer is the most intricate module. Things worth knowing before changing it:

• Font face caching: FontManager._dynamic_font_cache is a bounded LRU with 1024 entries. File hashes and embedded license records are cached per process in provenance.py; multilingual cmap coverage is cached in text_support.py.
• _RGB_PREFERRED_METHODS: the renderer may hand RGB arrays to augmentation methods; that set is why gradient_illumination is excluded from Rust acceleration.
• Effects are single-select: TEXT_EFFECT_NAMES is ordered, and sampling walks it taking the first probability hit.
• DecorStyle: the sampled per-line decoration state (active flags, color, sub/superscript targets) carried from sampling into drawing.

Rust boundary

_rust_accel.py is the only place Python touches native code. It exposes:

• HAS_RUST_ACCEL: the import-time capability flag,
• split_text_spans: mixed-script span splitting,
• estimate_bg: utilities,
• the accelerated apply_* augmentation kernels.

Everything is optional: augmentation.AUG_METHODS selects native implementations when available, and rendering.py keeps its Python span splitter. Python owns font coverage. Details in Rust Acceleration.

Testing

Tests under tests/ are grouped by subsystem. Start with test_config.py and test_config_loader.py for option changes, test_line_height.py for height sampling, and test_rendering.py for rendering changes.

Coverage is configured in pyproject.toml with branch coverage over src/khocr_gen, omitting _rust_accel.py (it is a thin loader over code tested on the Rust side).

See also

• Development: build, lint, test, release
• Rust Acceleration: the native extension in depth
• Output and Storage: the files this pipeline produces