wormocr 1.0.0__py3-none-any.whl

wormocr/checkpoints.py

@@ -0,0 +1,55 @@
+"""Checkpoint format.
+
+Checkpoints are PyTorch pickle files. Load only files you trust.
+They store the model, optimizer, charset, config, and a maturity label.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+import torch
+
+from wormocr import __version__
+from wormocr.exceptions import CheckpointError
+
+FORMAT_VERSION = 1
+
+
+def public_config(config: dict[str, Any]) -> dict[str, Any]:
+    return {key: value for key, value in config.items() if not str(key).startswith("_")}
+
+
+def save_checkpoint(path: str | Path, payload: dict[str, Any]) -> None:
+    destination = Path(path)
+    destination.parent.mkdir(parents=True, exist_ok=True)
+    body = dict(payload)
+    body["wormocr_format"] = FORMAT_VERSION
+    body["wormocr_version"] = __version__
+    torch.save(body, destination)
+
+
+def load_checkpoint(path: str | Path, map_location: str | torch.device = "cpu") -> dict[str, Any]:
+    source = Path(path)
+    if not source.exists():
+        raise CheckpointError(
+            f"checkpoint not found: {source}. Download the released model with "
+            "`wormocr download`, train one with `wormocr train --config configs/tiny.yaml`, "
+            "or pass an explicit --checkpoint."
+        )
+    if source.suffix == ".safetensors":
+        from wormocr.hub import load_safetensors
+
+        return load_safetensors(source, map_location=map_location)
+    try:
+        payload = torch.load(source, map_location=map_location, weights_only=False)
+    except TypeError:
+        payload = torch.load(source, map_location=map_location)
+    except Exception as exc:
+        raise CheckpointError(f"cannot read checkpoint {source}: {exc}") from exc
+    if not isinstance(payload, dict) or "model_state" not in payload:
+        raise CheckpointError(f"{source} is not a wormOCR checkpoint")
+    if "charset" not in payload or "config" not in payload:
+        raise CheckpointError(f"{source} is missing its charset or config")
+    return payload

wormocr/cli.py

@@ -0,0 +1,308 @@
+"""Command-line interface.
+
+Every documented command is implemented here. Checkpoints are local files.
+Only `wormocr download` uses the network.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+from wormocr import __version__
+from wormocr.benchmark.report import run_benchmark
+from wormocr.config import load_config
+from wormocr.data.fonts import discover_fonts
+from wormocr.data.synthetic import generate_dataset
+from wormocr.evaluation.errors import ErrorReport
+from wormocr.evaluation.metrics import cer
+from wormocr.evaluation.robustness import measure_robustness
+from wormocr.exceptions import WormOCRError
+from wormocr.export.onnx_export import export_onnx, onnx_max_abs_diff, quantize_onnx
+from wormocr.hub import DEFAULT_MODEL, RELEASES, download
+from wormocr.inference.engine import LineEngine
+from wormocr.models.crnn import build_model, count_parameters
+from wormocr.runtime import environment_report
+from wormocr.training.loop import build_datasets, evaluate_model, make_loader, train
+from wormocr.vocabulary import Vocabulary
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = build_parser()
+    args = parser.parse_args(argv)
+    try:
+        return int(args.func(args) or 0)
+    except WormOCRError as exc:
+        print(f"error: {exc}", file=sys.stderr)
+        return 1
+    except MemoryError as exc:
+        print(f"error: {exc}", file=sys.stderr)
+        return 1
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(prog="wormocr", description="Small local OCR for printed English text.")
+    parser.add_argument("--version", action="version", version=f"wormocr {__version__}")
+    commands = parser.add_subparsers(dest="command", required=True)
+
+    info = commands.add_parser("info", help="show the runtime, fonts, and model sizes")
+    info.set_defaults(func=command_info)
+
+    train_cmd = commands.add_parser("train", help="train or smoke-test a recognizer")
+    train_cmd.add_argument("--config", required=True)
+    train_cmd.add_argument("--mode", required=True, choices=["smoke", "debug", "baseline", "resume"])
+    train_cmd.add_argument("--checkpoint", help="checkpoint to resume from")
+    train_cmd.add_argument("--device", default=None)
+    train_cmd.add_argument("--set", dest="overrides", action="append", default=[], help="dotted override, such as train.max_steps=20")
+    train_cmd.set_defaults(func=command_train)
+
+    predict = commands.add_parser("predict", help="recognize an image or PDF")
+    predict.add_argument("--image", action="append", required=True)
+    predict.add_argument("--model", default=DEFAULT_MODEL, help=f"released model name (default: {DEFAULT_MODEL})")
+    predict.add_argument("--checkpoint", default=None)
+    predict.add_argument("--device", default="cpu")
+    predict.add_argument("--mode", default="auto", choices=["auto", "line", "page"])
+    predict.add_argument("--beam", type=int, default=0)
+    predict.add_argument("--json", action="store_true")
+    predict.add_argument("--no-orientation", action="store_true")
+    predict.set_defaults(func=command_predict)
+
+    evaluate = commands.add_parser("evaluate", help="measure CER and WER")
+    evaluate.add_argument("--model", required=True, help="checkpoint path")
+    evaluate.add_argument("--labels", default=None)
+    evaluate.add_argument("--image-root", default=None)
+    evaluate.add_argument("--synthetic-samples", type=int, default=100)
+    evaluate.add_argument("--seed", type=int, default=12345)
+    evaluate.add_argument("--device", default="cpu")
+    evaluate.add_argument("--beam", type=int, default=0)
+    evaluate.add_argument("--batch-size", type=int, default=8)
+    evaluate.add_argument("--errors", default=None, help="write an error-analysis markdown file")
+    evaluate.add_argument("--robustness", type=int, default=0, help="samples per corruption; 0 skips it")
+    evaluate.add_argument("--json", action="store_true")
+    evaluate.set_defaults(func=command_evaluate)
+
+    analyze = commands.add_parser("analyze", help="write an error analysis report")
+    analyze.add_argument("--model", required=True)
+    analyze.add_argument("--synthetic-samples", type=int, default=50)
+    analyze.add_argument("--seed", type=int, default=12345)
+    analyze.add_argument("--device", default="cpu")
+    analyze.add_argument("--output", required=True)
+    analyze.set_defaults(func=command_analyze)
+
+    export = commands.add_parser("export", help="export ONNX or INT8 ONNX")
+    export.add_argument("--model", required=True, help="checkpoint path")
+    export.add_argument("--format", required=True, choices=["onnx", "onnx-int8"])
+    export.add_argument("--output", required=True)
+    export.add_argument("--opset", type=int, default=17)
+    export.set_defaults(func=command_export)
+
+    benchmark = commands.add_parser("benchmark", help="measure size, accuracy, latency, and memory")
+    benchmark.add_argument("--model", required=True, help="checkpoint path")
+    benchmark.add_argument("--output", default="artifacts/benchmarks/report.md")
+    benchmark.add_argument("--samples", type=int, default=40)
+    benchmark.add_argument("--repeats", type=int, default=20)
+    benchmark.add_argument("--device", default="cpu")
+    benchmark.add_argument("--seed", type=int, default=12345)
+    benchmark.set_defaults(func=command_benchmark)
+
+    generate = commands.add_parser("generate", help="write a synthetic line dataset and manifest")
+    generate.add_argument("--config", default="configs/tiny.yaml")
+    generate.add_argument("--samples", type=int, default=32)
+    generate.add_argument("--out", required=True)
+    generate.add_argument("--seed", type=int, default=0)
+    generate.add_argument("--augment", action="store_true")
+    generate.set_defaults(func=command_generate)
+
+    serve_cmd = commands.add_parser("serve", help="open a local web UI for images and PDFs")
+    serve_cmd.add_argument("--model", default=DEFAULT_MODEL, help=f"released model name (default: {DEFAULT_MODEL})")
+    serve_cmd.add_argument("--checkpoint", default=None)
+    serve_cmd.add_argument("--device", default="cpu")
+    serve_cmd.add_argument("--host", default="127.0.0.1")
+    serve_cmd.add_argument("--port", type=int, default=8000)
+    serve_cmd.add_argument("--threads", type=int, default=None, help="CPU threads for PyTorch (default: half the logical CPUs, at most 16)")
+    serve_cmd.set_defaults(func=command_serve)
+
+    download_cmd = commands.add_parser("download", help="download released weights from GitHub and verify them")
+    download_cmd.add_argument("name", nargs="?", default=DEFAULT_MODEL, choices=sorted(RELEASES))
+    download_cmd.add_argument("--dir", default=None, help="cache folder (default: $WORMOCR_HOME or ~/.cache/wormocr)")
+    download_cmd.add_argument("--force", action="store_true", help="download again even if verified files exist")
+    download_cmd.add_argument("--no-onnx", action="store_true", help="skip the ONNX file")
+    download_cmd.set_defaults(func=command_download)
+    return parser
+
+
+def command_download(args: argparse.Namespace) -> int:
+    path = download(args.name, root=args.dir, force=args.force, include_onnx=not args.no_onnx)
+    print(f"{args.name} is ready: {path}")
+    return 0
+
+
+def command_serve(args: argparse.Namespace) -> int:
+    from wormocr.web import serve
+
+    serve(
+        model=args.model,
+        checkpoint=args.checkpoint,
+        device=args.device,
+        host=args.host,
+        port=args.port,
+        threads=args.threads,
+    )
+    return 0
+
+
+def command_info(_args: argparse.Namespace) -> int:
+    report = environment_report()
+    try:
+        fonts = discover_fonts("")
+        report["fonts"] = len(fonts)
+    except WormOCRError as exc:
+        report["fonts"] = 0
+        report["font_error"] = str(exc)
+    charset = Vocabulary()
+    for name in ("micro", "tiny", "small"):
+        config_name = f"configs/{name}.yaml"
+        try:
+            config = load_config(config_name)
+        except WormOCRError:
+            continue
+        model = build_model(config, charset.size)
+        report[f"{name}_parameters"] = count_parameters(model)
+    print(json.dumps(report, indent=2))
+    return 0
+
+
+def command_train(args: argparse.Namespace) -> int:
+    config = load_config(args.config, args.overrides, mode=args.mode)
+    if args.device:
+        config["device"] = args.device
+    result = train(config, mode=args.mode, resume_from=args.checkpoint)
+    print(json.dumps({
+        "last_checkpoint": result.last_checkpoint,
+        "best_checkpoint": result.best_checkpoint,
+        "steps": result.steps,
+        "best_cer": result.best_cer,
+        "maturity": result.maturity,
+        "final_loss": result.final_loss,
+    }, indent=2))
+    return 0
+
+
+def command_predict(args: argparse.Namespace) -> int:
+    from wormocr.api import OCR
+
+    ocr = OCR(model=args.model, checkpoint=args.checkpoint, device=args.device, beam_width=args.beam)
+    results = [
+        ocr.predict(image, mode=args.mode, correct_orientation=not args.no_orientation)
+        for image in args.image
+    ]
+    if args.json or len(results) > 1:
+        print(json.dumps([result.to_dict() for result in results], indent=2, ensure_ascii=True))
+    else:
+        print(results[0].text)
+    return 0
+
+
+def command_evaluate(args: argparse.Namespace) -> int:
+    engine = LineEngine.load(args.model, device=args.device, beam_width=args.beam)
+    config = dict(engine.config)
+    if args.labels:
+        config["data"]["labels"] = args.labels
+        config["data"]["image_root"] = args.image_root or str(Path(args.labels).parent)
+        config["data"]["val_labels"] = None
+    else:
+        config["data"]["labels"] = None
+        config["data"]["train_samples"] = args.synthetic_samples
+        config["data"]["augment"] = False
+        config["seed"] = args.seed
+    if args.labels:
+        dataset, _val, vocabulary = build_datasets(config)
+    else:
+        from wormocr.data.dataset import OnlineSynthDataset
+
+        vocabulary = engine.vocabulary
+        dataset = OnlineSynthDataset(config, args.synthetic_samples, args.seed, augment=False)
+    loader = make_loader(dataset, args.batch_size, False, args.seed, 0)
+    totals, pairs = evaluate_model(
+        engine.model, loader, vocabulary, engine.device, collect=bool(args.errors)
+    )
+    payload = totals.to_dict()
+    payload["maturity"] = engine.maturity
+    payload["checkpoint"] = str(args.model)
+    payload["synthetic"] = not bool(args.labels)
+    if args.robustness:
+        payload["robustness"] = measure_robustness(engine, args.robustness, args.seed)
+    if args.errors:
+        report = ErrorReport()
+        for reference, hypothesis in pairs:
+            report.add(reference, hypothesis, cer(reference, hypothesis))
+        Path(args.errors).parent.mkdir(parents=True, exist_ok=True)
+        Path(args.errors).write_text(report.to_markdown(), encoding="utf-8")
+        payload["errors"] = args.errors
+    print(json.dumps(payload, indent=2))
+    return 0
+
+
+def command_analyze(args: argparse.Namespace) -> int:
+    args.errors = args.output
+    args.labels = None
+    args.image_root = None
+    args.beam = 0
+    args.batch_size = 8
+    args.robustness = 0
+    args.json = True
+    return command_evaluate(args)
+
+
+def command_export(args: argparse.Namespace) -> int:
+    engine = LineEngine.load(args.model, device="cpu")
+    output = Path(args.output)
+    if args.format == "onnx":
+        export_onnx(engine.model, output, opset=args.opset)
+        diff = onnx_max_abs_diff(engine.model, output)
+        if diff > 1e-3:
+            raise WormOCRError(f"ONNX outputs differ from PyTorch by {diff:.4g}")
+        print(json.dumps({"output": str(output), "bytes": output.stat().st_size, "max_abs_diff": diff}, indent=2))
+        return 0
+    fp32 = output.with_suffix(".fp32.onnx")
+    export_onnx(engine.model, fp32, opset=args.opset)
+    fp32_diff = onnx_max_abs_diff(engine.model, fp32)
+    quantize_onnx(fp32, output)
+    int8_diff = onnx_max_abs_diff(engine.model, output)
+    print(json.dumps({
+        "output": str(output),
+        "fp32": str(fp32),
+        "fp32_bytes": fp32.stat().st_size,
+        "int8_bytes": output.stat().st_size,
+        "fp32_max_abs_diff": fp32_diff,
+        "int8_max_abs_diff": int8_diff,
+    }, indent=2))
+    return 0
+
+
+def command_benchmark(args: argparse.Namespace) -> int:
+    report = run_benchmark(
+        args.model,
+        args.output,
+        samples=args.samples,
+        seed=args.seed,
+        repeats=args.repeats,
+        device=args.device,
+    )
+    print(f"wrote {args.output}")
+    synthetic = report["synthetic_held_out"]
+    print(
+        f"parameters={report['parameters']} maturity={report['maturity']} "
+        f"CER={synthetic['cer']:.4f} median_model_ms="
+        f"{report['latency']['model_only_batch1_width256']['median_ms']:.2f}"
+    )
+    return 0
+
+
+def command_generate(args: argparse.Namespace) -> int:
+    config = load_config(args.config)
+    manifest = generate_dataset(config, args.samples, args.out, args.seed, augment=args.augment)
+    print(manifest)
+    return 0

wormocr/config.py

@@ -0,0 +1,220 @@
+"""YAML configuration, defaults, and dotted command-line overrides."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+import yaml
+
+from wormocr.exceptions import ConfigError
+
+DEFAULT_CONFIG: dict[str, Any] = {
+    "seed": 1337,
+    "device": "auto",
+    "model": {
+        "name": "tiny",
+        "in_channels": 1,
+        "image_height": 32,
+        "width_stride": 4,
+        "channels": [32, 64, 128, 192],
+        "separable": True,
+        "rnn_hidden": 320,
+        "rnn_layers": 2,
+        "rnn_dropout": 0.1,
+        "proj_size": 256,
+    },
+    "vocab": {"charset": None},
+    "data": {
+        "train_samples": 4000,
+        "val_samples": 400,
+        "val_seed": None,
+        "font_size_min": 16,
+        "font_size_max": 32,
+        "max_chars": 40,
+        "augment": True,
+        "font_dir": None,
+        "labels": None,
+        "val_labels": None,
+        "image_root": None,
+        "val_image_root": None,
+        "max_width": 1024,
+    },
+    "train": {
+        "epochs": 10,
+        "batch_size": 32,
+        "lr": 0.001,
+        "min_lr_ratio": 0.1,
+        "weight_decay": 0.01,
+        "grad_clip": 5.0,
+        "warmup_steps": 100,
+        "max_steps": None,
+        "max_minutes": None,
+        "early_stopping_patience": 5,
+        "checkpoint_dir": "artifacts/checkpoints/tiny",
+        "log_every": 20,
+        "save_every": 5,
+        "amp": False,
+        "num_workers": 0,
+        "torch_threads": None,
+        "ema_decay": 0.0,
+        "eval_every": 0,
+    },
+    "eval": {"batch_size": 32, "beam_width": 0, "synthetic_samples": 200},
+    "export": {"opset": 17},
+}
+
+MODE_PRESETS: dict[str, dict[str, Any]] = {
+    "smoke": {
+        "train": {
+            "epochs": 1,
+            "max_steps": 2,
+            "batch_size": 4,
+            "log_every": 1,
+            "early_stopping_patience": 0,
+            "warmup_steps": 0,
+            "save_every": 0,
+            "max_minutes": 5,
+            "num_workers": 0,
+            "eval_every": 0,
+        },
+        "data": {"train_samples": 8, "val_samples": 4, "augment": False},
+    },
+    "debug": {
+        "train": {
+            "epochs": 1,
+            "max_steps": 40,
+            "log_every": 5,
+            "early_stopping_patience": 0,
+            "warmup_steps": 0,
+            "max_minutes": 15,
+            "num_workers": 0,
+            "eval_every": 0,
+        },
+        "data": {"train_samples": 128, "val_samples": 32},
+    },
+    "baseline": {},
+    "resume": {},
+}
+
+
+def repo_root() -> Path:
+    here = Path(__file__).resolve()
+    for parent in here.parents:
+        if (parent / "pyproject.toml").exists() and (parent / "configs").exists():
+            return parent
+    return Path.cwd()
+
+
+def config_path(name: str) -> Path:
+    candidate = Path(name)
+    if candidate.exists():
+        return candidate
+    rooted = repo_root() / name
+    if rooted.exists():
+        return rooted
+    by_stem = repo_root() / "configs" / f"{name}.yaml"
+    if by_stem.exists():
+        return by_stem
+    raise ConfigError(f"config not found: {name}")
+
+
+def deep_update(base: dict[str, Any], override: dict[str, Any]) -> dict[str, Any]:
+    merged = dict(base)
+    for key, value in override.items():
+        if isinstance(value, dict) and isinstance(merged.get(key), dict):
+            merged[key] = deep_update(merged[key], value)
+        else:
+            merged[key] = value
+    return merged
+
+
+def parse_override_value(raw: str) -> Any:
+    try:
+        return yaml.safe_load(raw)
+    except yaml.YAMLError:
+        return raw
+
+
+def set_dotted(config: dict[str, Any], dotted: str, value: Any) -> None:
+    if not dotted or dotted.startswith(".") or dotted.endswith("."):
+        raise ConfigError(f"invalid override key: {dotted!r}")
+    cursor = config
+    keys = dotted.split(".")
+    for key in keys[:-1]:
+        child = cursor.get(key)
+        if not isinstance(child, dict):
+            child = {}
+            cursor[key] = child
+        cursor = child
+    cursor[keys[-1]] = value
+
+
+def apply_overrides(config: dict[str, Any], overrides: list[str] | None) -> dict[str, Any]:
+    updated = config
+    for item in overrides or []:
+        if "=" not in item:
+            raise ConfigError(f"override must look like key=value, got {item!r}")
+        key, raw = item.split("=", 1)
+        set_dotted(updated, key.strip(), parse_override_value(raw))
+    return updated
+
+
+def apply_mode(config: dict[str, Any], mode: str, labels_set: bool = False) -> dict[str, Any]:
+    if mode not in MODE_PRESETS:
+        known = ", ".join(sorted(MODE_PRESETS))
+        raise ConfigError(f"unknown train mode {mode!r}. Expected one of: {known}")
+    preset = MODE_PRESETS[mode]
+    if mode == "smoke" and labels_set:
+        preset = {"train": preset.get("train", {})}
+    return deep_update(config, preset)
+
+
+def validate_config(config: dict[str, Any]) -> None:
+    channels = config["model"]["channels"]
+    if not isinstance(channels, list) or len(channels) != 4:
+        raise ConfigError("model.channels must be a list of four positive integers")
+    if any(int(channel) < 4 for channel in channels):
+        raise ConfigError("each model channel count must be at least 4")
+    height = int(config["model"]["image_height"])
+    if height % 16 != 0 or height < 16:
+        raise ConfigError("model.image_height must be a multiple of 16 and at least 16")
+    if int(config["model"]["rnn_hidden"]) < 8:
+        raise ConfigError("model.rnn_hidden must be at least 8")
+    if int(config["model"]["rnn_layers"]) < 1:
+        raise ConfigError("model.rnn_layers must be at least 1")
+    if int(config["train"]["batch_size"]) < 1:
+        raise ConfigError("train.batch_size must be at least 1")
+    if int(config["data"]["train_samples"]) < 2:
+        raise ConfigError("data.train_samples must be at least 2")
+    if float(config["train"]["lr"]) <= 0:
+        raise ConfigError("train.lr must be positive")
+
+
+def load_config(
+    path: str | Path,
+    overrides: list[str] | None = None,
+    mode: str | None = None,
+) -> dict[str, Any]:
+    source = config_path(str(path))
+    try:
+        loaded = yaml.safe_load(source.read_text(encoding="utf-8")) or {}
+    except OSError as exc:
+        raise ConfigError(f"cannot read config {source}") from exc
+    except yaml.YAMLError as exc:
+        raise ConfigError(f"invalid YAML in {source}: {exc}") from exc
+    if not isinstance(loaded, dict):
+        raise ConfigError(f"config root must be a mapping: {source}")
+    config = deep_update(DEFAULT_CONFIG, loaded)
+    config = apply_overrides(config, overrides)
+    if mode:
+        config = apply_mode(config, mode, labels_set=bool(config["data"].get("labels")))
+        config = apply_overrides(config, overrides)
+    validate_config(config)
+    config["_config_path"] = str(source)
+    return config
+
+
+def dump_config(config: dict[str, Any]) -> str:
+    public = {key: value for key, value in config.items() if not str(key).startswith("_")}
+    return yaml.safe_dump(public, sort_keys=False)

wormocr/data/__init__.py

@@ -0,0 +1,5 @@
+"""Synthetic and file-backed OCR datasets."""
+
+from wormocr.data.dataset import LabelDataset, OnlineSynthDataset, collate_lines
+
+__all__ = ["LabelDataset", "OnlineSynthDataset", "collate_lines"]