wormocr 1.0.0__py3-none-any.whl

wormocr/exceptions.py

@@ -0,0 +1,25 @@
+"""Errors raised by wormOCR. They are intentional and safe to show to users."""
+
+
+class WormOCRError(Exception):
+    """Base class for expected wormOCR failures."""
+
+
+class ConfigError(WormOCRError):
+    """A configuration file or override is invalid."""
+
+
+class CheckpointError(WormOCRError):
+    """A checkpoint is missing, incompatible, or unusable."""
+
+
+class DataError(WormOCRError):
+    """A training or evaluation sample cannot be used."""
+
+
+class RecognitionError(WormOCRError):
+    """An image cannot be recognized."""
+
+
+class ExportError(WormOCRError):
+    """A deployment artifact could not be created or verified."""

wormocr/export/__init__.py

@@ -0,0 +1,5 @@
+"""Deployment export helpers."""
+
+from wormocr.export.onnx_export import export_onnx, onnx_max_abs_diff, quantize_onnx
+
+__all__ = ["export_onnx", "onnx_max_abs_diff", "quantize_onnx"]

wormocr/export/onnx_export.py

@@ -0,0 +1,75 @@
+"""ONNX export and dynamic INT8 quantization."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from wormocr.exceptions import ExportError
+
+
+def export_onnx(model: torch.nn.Module, path: str | Path, opset: int = 17) -> Path:
+    destination = Path(path)
+    destination.parent.mkdir(parents=True, exist_ok=True)
+    model.eval()
+    height = int(model.image_height)
+    dummy = torch.zeros(1, 1, height, 128, dtype=torch.float32)
+    try:
+        torch.onnx.export(
+            model,
+            dummy,
+            str(destination),
+            input_names=["image"],
+            output_names=["log_probs"],
+            dynamic_axes={"image": {0: "batch", 3: "width"}, "log_probs": {0: "batch", 1: "time"}},
+            opset_version=int(opset),
+            dynamo=False,
+        )
+    except TypeError:
+        torch.onnx.export(
+            model,
+            dummy,
+            str(destination),
+            input_names=["image"],
+            output_names=["log_probs"],
+            dynamic_axes={"image": {0: "batch", 3: "width"}, "log_probs": {0: "batch", 1: "time"}},
+            opset_version=int(opset),
+        )
+    if not destination.exists() or destination.stat().st_size == 0:
+        raise ExportError(f"ONNX export did not write {destination}")
+    return destination
+
+
+def quantize_onnx(source: str | Path, destination: str | Path) -> Path:
+    try:
+        from onnxruntime.quantization import QuantType, quantize_dynamic
+    except ImportError as exc:
+        raise ExportError(
+            "INT8 export needs the optional packages onnx and onnxruntime. "
+            "Install them with: pip install \"wormocr[export]\""
+        ) from exc
+    output = Path(destination)
+    output.parent.mkdir(parents=True, exist_ok=True)
+    try:
+        quantize_dynamic(str(source), str(output), weight_type=QuantType.QInt8)
+    except Exception as exc:
+        raise ExportError(f"INT8 quantization failed: {exc}") from exc
+    return output
+
+
+def onnx_max_abs_diff(model: torch.nn.Module, path: str | Path, width: int = 160, batch: int = 2) -> float:
+    try:
+        import onnxruntime as ort
+    except ImportError as exc:
+        raise ExportError("onnxruntime is not installed") from exc
+    model.eval()
+    sample = torch.randn(batch, 1, int(model.image_height), width)
+    with torch.no_grad():
+        reference = model(sample).detach().cpu().numpy()
+    session = ort.InferenceSession(str(path), providers=["CPUExecutionProvider"])
+    output = session.run(None, {"image": sample.numpy().astype(np.float32)})[0]
+    if output.shape != reference.shape:
+        raise ExportError(f"ONNX shape {output.shape} != PyTorch shape {reference.shape}")
+    return float(np.max(np.abs(reference - output)))

wormocr/hub.py

@@ -0,0 +1,157 @@
+"""Released wormOCR weights and the explicit download command.
+
+Weights are published as GitHub release assets, never stored in git, and
+never fetched implicitly: only `wormocr download` / `download()` touches the
+network. Every file is checked against a SHA-256 recorded here.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+import shutil
+import tempfile
+import urllib.request
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+from wormocr.exceptions import CheckpointError
+
+REPOSITORY = "arjun988/wormOCR"
+DEFAULT_MODEL = "wormocr-2M"
+
+
+@dataclass(frozen=True)
+class Release:
+    name: str
+    tag: str
+    files: dict[str, str]
+    description: str
+
+    def url(self, filename: str) -> str:
+        return f"https://github.com/{REPOSITORY}/releases/download/{self.tag}/{filename}"
+
+
+RELEASES: dict[str, Release] = {
+    "wormocr-2M": Release(
+        name="wormocr-2M",
+        tag="wormocr-2M",
+        files={
+            "wormocr-2M.safetensors": "a6bee28adfc15f219fc6ff9d0dccdce45dc92c73559bc5474dda3d7e887376f2",
+            "wormocr-2M.json": "62b7c98be82c34684e29da2d1c97e2a3efac4ce3eae4b4cd0153845f608e24c2",
+            "wormocr-2M.onnx": "18498d163622022c8d48cf3fca7fa3e2c6bbc140db720e293f4d9a2f7b4cb25e",
+        },
+        description="2,174,880 parameters. Printed and handwritten English lines. Replaces the printed-only weights previously published on this tag.",
+    ),
+    "wormocr-9M": Release(
+        name="wormocr-9M",
+        tag="wormocr-9M",
+        files={
+            "wormocr-9M.safetensors": "e10fbc8d96fb230056e9c10a4083ab971306179ada66b7cfc644fbf3e0402183",
+            "wormocr-9M.json": "cc6ca2695be4dde105d61591c5175653c02502bdc33711d59ec633bb470c05a9",
+            "wormocr-9M.onnx": "9eb068021344ce07f6846bbe7915d3d4919c1f5e9cd2de441b6fe71f656e51de",
+        },
+        description="8,776,416 parameters. Same training mix as wormocr-2M, larger CRNN (height 48). Stronger on handwritten notes, IAM, and Bentham.",
+    ),
+}
+
+
+def cache_dir() -> Path:
+    override = os.environ.get("WORMOCR_HOME")
+    if override:
+        return Path(override).expanduser()
+    base = os.environ.get("XDG_CACHE_HOME") or (Path.home() / ".cache")
+    return Path(base) / "wormocr"
+
+
+def get_release(name: str) -> Release:
+    try:
+        return RELEASES[name]
+    except KeyError:
+        available = ", ".join(sorted(RELEASES))
+        raise CheckpointError(f"unknown wormOCR model {name!r}. Released models: {available}") from None
+
+
+def sha256(path: str | Path) -> str:
+    digest = hashlib.sha256()
+    with open(path, "rb") as handle:
+        for block in iter(lambda: handle.read(1 << 20), b""):
+            digest.update(block)
+    return digest.hexdigest()
+
+
+def model_dir(name: str, root: str | Path | None = None) -> Path:
+    return (Path(root) if root else cache_dir()) / name
+
+
+def local_weights(name: str, root: str | Path | None = None) -> Path | None:
+    """Path to the downloaded safetensors file, or None if it is not there."""
+    release = get_release(name)
+    folder = model_dir(name, root)
+    weights = folder / f"{release.name}.safetensors"
+    config = folder / f"{release.name}.json"
+    return weights if weights.exists() and config.exists() else None
+
+
+def download(
+    name: str = DEFAULT_MODEL,
+    root: str | Path | None = None,
+    force: bool = False,
+    include_onnx: bool = True,
+    log=print,
+) -> Path:
+    """Download a released model into the cache and verify every file."""
+    release = get_release(name)
+    folder = model_dir(name, root)
+    folder.mkdir(parents=True, exist_ok=True)
+    for filename, expected in release.files.items():
+        if filename.endswith(".onnx") and not include_onnx:
+            continue
+        target = folder / filename
+        if target.exists() and not force and sha256(target) == expected:
+            log(f"{filename}: already downloaded and verified")
+            continue
+        url = release.url(filename)
+        log(f"{filename}: downloading {url}")
+        handle, temporary = tempfile.mkstemp(dir=folder, suffix=".part")
+        os.close(handle)
+        try:
+            with urllib.request.urlopen(url, timeout=60) as response, open(temporary, "wb") as out:
+                shutil.copyfileobj(response, out)
+            actual = sha256(temporary)
+            if actual != expected:
+                raise CheckpointError(f"{filename}: SHA-256 mismatch (expected {expected}, got {actual})")
+            os.replace(temporary, target)
+        except OSError as exc:
+            raise CheckpointError(f"{filename}: download failed: {exc}") from exc
+        finally:
+            if os.path.exists(temporary):
+                os.remove(temporary)
+        log(f"{filename}: {target.stat().st_size / 1e6:.1f} MB, SHA-256 verified")
+    return folder / f"{release.name}.safetensors"
+
+
+def load_safetensors(path: str | Path, map_location: Any = "cpu") -> dict[str, Any]:
+    """Read released weights (`<name>.safetensors` + `<name>.json`) as a checkpoint payload."""
+    try:
+        from safetensors.torch import load_file
+    except ImportError as exc:
+        raise CheckpointError("loading .safetensors weights needs: pip install safetensors") from exc
+    weights = Path(path)
+    meta_path = weights.with_suffix(".json")
+    if not meta_path.exists():
+        raise CheckpointError(f"{weights} needs its metadata file {meta_path.name} next to it")
+    meta = json.loads(meta_path.read_text(encoding="utf-8"))
+    for key in ("config", "charset"):
+        if key not in meta:
+            raise CheckpointError(f"{meta_path} is missing {key!r}")
+    device = str(map_location) if map_location is not None else "cpu"
+    return {
+        "model_state": load_file(str(weights), device=device),
+        "config": meta["config"],
+        "charset": meta["charset"],
+        "maturity": meta.get("maturity", "untrained"),
+        "name": meta.get("name"),
+    }

wormocr/inference/__init__.py

@@ -0,0 +1,5 @@
+"""Decoding and recognition runtime."""
+
+from wormocr.inference.decode import ctc_beam_search, greedy_decode
+
+__all__ = ["ctc_beam_search", "greedy_decode"]

wormocr/inference/decode.py

@@ -0,0 +1,113 @@
+"""CTC decoders.
+
+Greedy decoding is the default. Beam search is optional and only worth its cost
+when a measurement shows fewer errors.
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+
+
+def _logadd(left: float, right: float) -> float:
+    if left == float("-inf"):
+        return right
+    if right == float("-inf"):
+        return left
+    if left > right:
+        return left + math.log1p(math.exp(right - left))
+    return right + math.log1p(math.exp(left - right))
+
+
+def greedy_decode(
+    log_probs: torch.Tensor,
+    lengths: torch.Tensor | None = None,
+    blank: int = 0,
+) -> list[list[int]]:
+    """Collapse CTC greedy paths. Repeated labels are kept when a blank separates them."""
+    if log_probs.ndim != 3:
+        raise ValueError("log_probs must have shape (batch, time, classes)")
+    chosen = log_probs.argmax(dim=-1).detach().cpu()
+    batch, time, _classes = log_probs.shape
+    decoded: list[list[int]] = []
+    for index in range(batch):
+        limit = time if lengths is None else max(0, min(time, int(lengths[index])))
+        previous = blank
+        sequence: list[int] = []
+        row = chosen[index]
+        for step in range(limit):
+            label = int(row[step])
+            if label != blank and label != previous:
+                sequence.append(label)
+            previous = label
+        decoded.append(sequence)
+    return decoded
+
+
+def greedy_confidence(
+    log_probs: torch.Tensor,
+    length: int,
+    blank: int = 0,
+) -> float:
+    """Mean probability of the labels that greedy decoding actually emitted."""
+    limit = max(0, min(int(length), int(log_probs.shape[0])))
+    if limit == 0:
+        return 0.0
+    probabilities = log_probs[:limit].detach().float().exp()
+    labels = probabilities.argmax(dim=-1)
+    chosen: list[float] = []
+    previous = blank
+    for step, label in enumerate(labels.tolist()):
+        label = int(label)
+        if label != blank and label != previous:
+            chosen.append(float(probabilities[step, label]))
+        previous = label
+    if not chosen:
+        return float(probabilities[:, blank].mean())
+    return float(sum(chosen) / len(chosen))
+
+
+def ctc_beam_search(log_probs: torch.Tensor, beam_width: int = 8, blank: int = 0) -> list[int]:
+    """Prefix beam search for one sequence of shape (time, classes)."""
+    if beam_width < 1:
+        raise ValueError("beam_width must be at least 1")
+    if log_probs.ndim != 2:
+        raise ValueError("beam search expects log_probs of shape (time, classes)")
+    frame_log = log_probs.detach().float().cpu()
+    time, classes = frame_log.shape
+    neg = float("-inf")
+    beam: dict[tuple[int, ...], list[float]] = {(): [0.0, neg]}
+
+    def add(store: dict[tuple[int, ...], list[float]], prefix: tuple[int, ...], blank_p: float, non_p: float) -> None:
+        slot = store.get(prefix)
+        if slot is None:
+            store[prefix] = [blank_p, non_p]
+            return
+        slot[0] = _logadd(slot[0], blank_p)
+        slot[1] = _logadd(slot[1], non_p)
+
+    for step in range(time):
+        frame = frame_log[step]
+        ranked = sorted(beam.items(), key=lambda item: _logadd(item[1][0], item[1][1]), reverse=True)[
+            :beam_width
+        ]
+        nxt: dict[tuple[int, ...], list[float]] = {}
+        for prefix, (prob_blank, prob_label) in ranked:
+            total = _logadd(prob_blank, prob_label)
+            add(nxt, prefix, total + float(frame[blank]), neg)
+            for label in range(classes):
+                if label == blank:
+                    continue
+                emit = float(frame[label])
+                if prefix and label == prefix[-1]:
+                    add(nxt, prefix, neg, prob_label + emit)
+                    add(nxt, prefix + (label,), neg, prob_blank + emit)
+                else:
+                    add(nxt, prefix + (label,), neg, total + emit)
+        beam = dict(
+            sorted(nxt.items(), key=lambda item: _logadd(item[1][0], item[1][1]), reverse=True)[:beam_width]
+        )
+    best = max(beam.items(), key=lambda item: _logadd(item[1][0], item[1][1]))
+    return list(best[0])

wormocr/inference/engine.py

@@ -0,0 +1,122 @@
+"""Load a checkpoint and recognize text lines."""
+
+from __future__ import annotations
+
+import warnings
+from pathlib import Path
+
+import torch
+from PIL import Image
+
+from wormocr.checkpoints import load_checkpoint
+from wormocr.inference.decode import ctc_beam_search, greedy_confidence, greedy_decode
+from wormocr.models.crnn import build_model
+from wormocr.preprocessing import prepare_line
+from wormocr.runtime import resolve_device
+from wormocr.vocabulary import Vocabulary
+
+
+class LineEngine:
+    def __init__(
+        self,
+        model,
+        vocabulary: Vocabulary,
+        device: torch.device,
+        config: dict,
+        maturity: str,
+        beam_width: int = 0,
+        max_width: int = 4096,
+        batch_size: int = 8,
+    ) -> None:
+        self.model = model
+        self.vocabulary = vocabulary
+        self.device = device
+        self.config = config
+        self.maturity = maturity
+        self.beam_width = int(beam_width)
+        self.max_width = int(max_width)
+        self.batch_size = max(1, int(batch_size))
+        self.image_height = int(config["model"]["image_height"])
+        self.model.to(device)
+        self.model.eval()
+
+    @classmethod
+    def load(
+        cls,
+        path: str | Path,
+        device: str | None = None,
+        beam_width: int = 0,
+        batch_size: int = 8,
+    ) -> "LineEngine":
+        torch_device = resolve_device(device)
+        payload = load_checkpoint(path, map_location=torch_device)
+        vocabulary = Vocabulary(payload["charset"])
+        model = build_model(payload["config"], vocabulary.size)
+        model.load_state_dict(payload["model_state"])
+        maturity = str(payload.get("maturity") or "untrained")
+        if maturity in {"untrained", "smoke-tested"}:
+            warnings.warn(
+                f"This checkpoint is {maturity} and is not a quality OCR model.",
+                UserWarning,
+                stacklevel=2,
+            )
+        return cls(
+            model=model,
+            vocabulary=vocabulary,
+            device=torch_device,
+            config=payload["config"],
+            maturity=maturity,
+            beam_width=beam_width,
+            max_width=int(payload["config"].get("data", {}).get("max_width") or 1024) * 2,
+            batch_size=batch_size,
+        )
+
+    def recognize(self, images: list[Image.Image]) -> list[dict[str, float | str]]:
+        if not images:
+            return []
+        prepared: list[tuple[torch.Tensor, float, int]] = []
+        for image in images:
+            tensor, pad_value = prepare_line(
+                image,
+                height=self.image_height,
+                align=2,
+                max_width=self.max_width,
+            )
+            prepared.append((tensor, pad_value, int(tensor.shape[-1])))
+        order = sorted(range(len(prepared)), key=lambda index: prepared[index][2])
+        outputs: list[dict[str, float | str] | None] = [None] * len(prepared)
+        for start in range(0, len(order), self.batch_size):
+            chunk = order[start : start + self.batch_size]
+            widths = [prepared[index][2] for index in chunk]
+            max_width = max(widths)
+            batch = prepared[chunk[0]][0].new_empty(len(chunk), 1, self.image_height, max_width)
+            for row, index in enumerate(chunk):
+                tensor, pad_value, width = prepared[index]
+                batch[row, :, :, :width] = tensor
+                if width < max_width:
+                    batch[row, :, :, width:] = pad_value
+            with torch.inference_mode():
+                log_probs = self.model(batch.to(self.device))
+            length_tensor = self.model.output_lengths(torch.tensor(widths, dtype=torch.long))
+            length_tensor = torch.minimum(length_tensor, torch.tensor(log_probs.shape[1]))
+            if self.beam_width > 0:
+                sequences = [
+                    ctc_beam_search(
+                        log_probs[row, : int(length_tensor[row])].cpu(),
+                        self.beam_width,
+                        self.vocabulary.blank,
+                    )
+                    for row in range(len(chunk))
+                ]
+            else:
+                sequences = greedy_decode(log_probs, length_tensor, self.vocabulary.blank)
+            for row, index in enumerate(chunk):
+                outputs[index] = {
+                    "text": self.vocabulary.decode(sequences[row]),
+                    "confidence": greedy_confidence(
+                        log_probs[row].cpu(),
+                        int(length_tensor[row]),
+                        self.vocabulary.blank,
+                    ),
+                }
+        return [item for item in outputs if item is not None]

wormocr/models/__init__.py

@@ -0,0 +1,5 @@
+"""Recognizer architectures."""
+
+from wormocr.models.crnn import CRNN, build_model, count_parameters
+
+__all__ = ["CRNN", "build_model", "count_parameters"]