wormocr 1.0.0__py3-none-any.whl

wormocr/models/crnn.py

@@ -0,0 +1,236 @@
+"""CNN (optionally depthwise-separable) + BiGRU + CTC log-probabilities.
+
+Width is downsampled by `width_stride` (2 or 4). Height 32 is reduced to a
+short feature column and flattened. Stride 4 halves the sequence the GRU must
+read and makes CTC alignment converge much sooner; older checkpoints that do
+not store the field load as stride 2.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+from torch import nn
+
+from wormocr.exceptions import ConfigError
+
+
+@dataclass(frozen=True)
+class Spec:
+    kind: str
+    out_channels: int
+    kernel: tuple[int, int]
+    stride: tuple[int, int]
+    padding: tuple[int, int]
+
+
+def build_specs(channels: list[int], width_stride: int = 2) -> list[Spec]:
+    if len(channels) != 4:
+        raise ConfigError("the recognizer expects four channel stages")
+    if int(width_stride) not in (2, 4):
+        raise ConfigError("model.width_stride must be 2 or 4")
+    first, second, third, fourth = (int(value) for value in channels)
+    second_pool = (2, 2) if int(width_stride) == 4 else (2, 1)
+    return [
+        Spec("conv", first, (3, 3), (1, 1), (1, 1)),
+        Spec("pool", 0, (2, 2), (2, 2), (0, 0)),
+        Spec("conv", second, (3, 3), (1, 1), (1, 1)),
+        Spec("pool", 0, second_pool, second_pool, (0, 0)),
+        Spec("conv", third, (3, 3), (1, 1), (1, 1)),
+        Spec("pool", 0, (2, 1), (2, 1), (0, 0)),
+        Spec("conv", fourth, (3, 3), (1, 1), (1, 1)),
+        Spec("pool", 0, (2, 1), (2, 1), (0, 0)),
+    ]
+
+
+def _spatial(size: int, kernel: int, stride: int, padding: int) -> int:
+    return (int(size) + 2 * padding - kernel) // stride + 1
+
+
+def output_width(width: int, specs: list[Spec]) -> int:
+    size = int(width)
+    for spec in specs:
+        size = _spatial(size, spec.kernel[1], spec.stride[1], spec.padding[1])
+        if size <= 0:
+            return 0
+    return size
+
+
+class ConvBlock(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel: tuple[int, int],
+        stride: tuple[int, int],
+        padding: tuple[int, int],
+        separable: bool,
+    ) -> None:
+        super().__init__()
+        if separable:
+            self.conv: nn.Module = nn.Sequential(
+                nn.Conv2d(
+                    in_channels,
+                    in_channels,
+                    kernel,
+                    stride=stride,
+                    padding=padding,
+                    groups=in_channels,
+                    bias=False,
+                ),
+                nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False),
+            )
+        else:
+            self.conv = nn.Conv2d(
+                in_channels,
+                out_channels,
+                kernel,
+                stride=stride,
+                padding=padding,
+                bias=False,
+            )
+        self.bn = nn.BatchNorm2d(out_channels)
+        self.act = nn.ReLU(inplace=False)
+
+    def forward(self, images: torch.Tensor) -> torch.Tensor:
+        return self.act(self.bn(self.conv(images)))
+
+
+class Encoder(nn.Module):
+    def __init__(self, in_channels: int, specs: list[Spec], separable: bool) -> None:
+        super().__init__()
+        self.specs = specs
+        blocks: list[nn.Module] = []
+        channels = in_channels
+        for spec in specs:
+            if spec.kind == "conv":
+                blocks.append(
+                    ConvBlock(
+                        channels,
+                        spec.out_channels,
+                        spec.kernel,
+                        spec.stride,
+                        spec.padding,
+                        separable=separable and channels > 1,
+                    )
+                )
+                channels = spec.out_channels
+            elif spec.kind == "pool":
+                blocks.append(
+                    nn.MaxPool2d(
+                        kernel_size=spec.kernel,
+                        stride=spec.stride,
+                        padding=spec.padding,
+                    )
+                )
+            else:
+                raise ConfigError(f"unknown encoder spec {spec.kind}")
+        self.blocks = nn.ModuleList(blocks)
+
+    def forward(self, images: torch.Tensor) -> torch.Tensor:
+        features = images
+        for block in self.blocks:
+            features = block(features)
+        return features
+
+    def output_width(self, width: int) -> int:
+        return output_width(width, self.specs)
+
+
+class CRNN(nn.Module):
+    """Variable-width line recognizer. Forward returns log-probabilities (B, T, C)."""
+
+    def __init__(
+        self,
+        num_classes: int,
+        in_channels: int = 1,
+        image_height: int = 32,
+        channels: list[int] | None = None,
+        separable: bool = True,
+        rnn_hidden: int = 256,
+        rnn_layers: int = 2,
+        rnn_dropout: float = 0.1,
+        proj_size: int = 256,
+        width_stride: int = 2,
+    ) -> None:
+        super().__init__()
+        if num_classes < 2:
+            raise ConfigError("num_classes must include the blank and at least one character")
+        stage_channels = list(channels or [32, 64, 128, 192])
+        self.image_height = int(image_height)
+        self.num_classes = int(num_classes)
+        self.width_stride = int(width_stride)
+        self.encoder = Encoder(in_channels, build_specs(stage_channels, self.width_stride), separable)
+        with torch.no_grad():
+            probe = torch.zeros(1, in_channels, self.image_height, 32)
+            encoded = self.encoder(probe)
+        if encoded.shape[-2] < 1 or encoded.shape[-1] < 1:
+            raise ConfigError(
+                f"encoder collapsed a {self.image_height}px-tall input; "
+                "choose a taller image_height that is a multiple of 16"
+            )
+        feature_height = int(encoded.shape[-2])
+        feature_channels = int(encoded.shape[1])
+        feature_size = feature_channels * feature_height
+        self.feature_size = feature_size
+        self.proj = nn.Sequential(
+            nn.Linear(feature_size, proj_size),
+            nn.ReLU(inplace=False),
+        )
+        dropout = float(rnn_dropout) if rnn_layers > 1 else 0.0
+        self.rnn = nn.GRU(
+            input_size=proj_size,
+            hidden_size=int(rnn_hidden),
+            num_layers=int(rnn_layers),
+            batch_first=True,
+            bidirectional=True,
+            dropout=dropout if rnn_layers > 1 else 0.0,
+        )
+        self.classifier = nn.Linear(int(rnn_hidden) * 2, self.num_classes)
+        nn.init.zeros_(self.classifier.bias)
+        # Class 0 is the CTC blank. A negative bias stops a fresh model from
+        # taking the blank at every frame before it has aligned a character.
+        self.classifier.bias.data[0] = -2.0
+
+    def output_lengths(self, widths: torch.Tensor) -> torch.Tensor:
+        lengths = [self.encoder.output_width(int(width)) for width in widths.detach().cpu().tolist()]
+        return torch.tensor(lengths, dtype=torch.long)
+
+    def forward(self, images: torch.Tensor) -> torch.Tensor:
+        if images.ndim != 4:
+            raise ValueError("images must have shape (batch, channels, height, width)")
+        if images.shape[2] != self.image_height:
+            raise ValueError(
+                f"expected image height {self.image_height}, got {images.shape[2]}"
+            )
+        features = self.encoder(images)
+        batch, channels, height, time = features.shape
+        sequence = features.permute(0, 3, 1, 2).contiguous().view(batch, time, channels * height)
+        sequence = self.proj(sequence)
+        sequence, _hidden = self.rnn(sequence)
+        return self.classifier(sequence).float().log_softmax(dim=-1)
+
+
+def build_model(config: dict, num_classes: int) -> CRNN:
+    model_cfg = config["model"]
+    return CRNN(
+        num_classes=num_classes,
+        in_channels=int(model_cfg.get("in_channels", 1)),
+        image_height=int(model_cfg["image_height"]),
+        channels=[int(value) for value in model_cfg["channels"]],
+        separable=bool(model_cfg["separable"]),
+        rnn_hidden=int(model_cfg["rnn_hidden"]),
+        rnn_layers=int(model_cfg["rnn_layers"]),
+        rnn_dropout=float(model_cfg.get("rnn_dropout", 0.0)),
+        proj_size=int(model_cfg.get("proj_size", model_cfg["rnn_hidden"])),
+        width_stride=int(model_cfg.get("width_stride", 2)),
+    )
+
+
+def count_parameters(module: nn.Module) -> int:
+    return sum(parameter.numel() for parameter in module.parameters())
+
+
+def state_dict_num_bytes(module: nn.Module) -> int:
+    return sum(parameter.numel() * parameter.element_size() for parameter in module.parameters())

wormocr/preprocessing.py

@@ -0,0 +1,135 @@
+"""Image loading and line normalization.
+
+The recognizer consumes grayscale tensors of a fixed height and variable width,
+normalized to [-1, 1]. Padding on the right uses the image border so a short
+line in a batch is not given a false edge.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+import torch
+from PIL import Image, UnidentifiedImageError
+
+from wormocr.exceptions import RecognitionError
+
+
+def to_grayscale(image: Image.Image) -> Image.Image:
+    if image.mode == "RGBA":
+        background = Image.new("RGBA", image.size, (255, 255, 255, 255))
+        image = Image.alpha_composite(background, image)
+        return image.convert("L")
+    if image.mode != "L":
+        return image.convert("L")
+    return image
+
+
+def border_color(image: Image.Image) -> int:
+    array = np.asarray(image.convert("L"))
+    if array.size == 0:
+        return 255
+    border = np.concatenate(
+        [array[:, 0], array[:, -1], array[0, :], array[-1, :]]
+    )
+    return int(np.median(border))
+
+
+def resize_to_height(image: Image.Image, height: int) -> Image.Image:
+    if height < 8:
+        raise ValueError("image height must be at least 8")
+    width, source_height = image.size
+    if width < 1 or source_height < 1:
+        raise ValueError("image is empty")
+    if source_height == height:
+        return image
+    new_width = max(1, int(round(width * (height / float(source_height)))))
+    return image.resize((new_width, height), Image.Resampling.BILINEAR)
+
+
+def _normalize_pixel(value: float) -> float:
+    return (value / 255.0 - 0.5) / 0.5
+
+
+def prepare_line(
+    image: Image.Image,
+    height: int = 32,
+    align: int = 2,
+    max_width: int = 2048,
+    min_width: int = 8,
+) -> tuple[torch.Tensor, float]:
+    """Return a float tensor of shape (1, height, width) and its pad value."""
+    if align < 1:
+        raise ValueError("align must be positive")
+    gray = to_grayscale(image)
+    if gray.width < 1 or gray.height < 1:
+        raise ValueError("image is empty")
+    resized = resize_to_height(gray, height)
+    if resized.width > max_width:
+        resized = resized.resize((max_width, height), Image.Resampling.BILINEAR)
+    color = border_color(resized)
+    if resized.width < min_width:
+        canvas = Image.new("L", (min_width, height), color=color)
+        canvas.paste(resized, (0, 0))
+        resized = canvas
+    if resized.width % align != 0:
+        padded_width = resized.width + (align - resized.width % align)
+        canvas = Image.new("L", (padded_width, height), color=color)
+        canvas.paste(resized, (0, 0))
+        resized = canvas
+        color = border_color(resized)
+    array = np.asarray(resized, dtype=np.float32)
+    tensor = torch.from_numpy(array).unsqueeze(0)
+    tensor = (tensor / 255.0 - 0.5) / 0.5
+    return tensor, float(_normalize_pixel(color))
+
+
+def image_from_array(array: np.ndarray) -> Image.Image:
+    if array.ndim not in (2, 3):
+        raise RecognitionError("expected an HxW or HxWx3 image array")
+    values = array
+    if np.issubdtype(values.dtype, np.floating):
+        values = np.clip(values * 255.0, 0, 255).astype(np.uint8)
+    else:
+        values = values.astype(np.uint8)
+    if values.ndim == 2:
+        return Image.fromarray(values, mode="L")
+    channels = values.shape[2]
+    if channels == 3:
+        return Image.fromarray(values, mode="RGB")
+    if channels == 4:
+        return Image.fromarray(values, mode="RGBA")
+    raise RecognitionError(f"unsupported channel count: {channels}")
+
+
+def load_image(source: str | Path) -> Image.Image:
+    path = Path(source)
+    if not path.exists():
+        raise RecognitionError(f"image not found: {path}")
+    if path.suffix.lower() == ".pdf":
+        raise RecognitionError("use the document loader for PDF files")
+    try:
+        with Image.open(path) as handle:
+            return handle.convert("RGBA") if handle.mode == "RGBA" else handle.convert("RGB")
+    except UnidentifiedImageError as exc:
+        raise RecognitionError(f"cannot read image: {path}") from exc
+    except OSError as exc:
+        raise RecognitionError(f"cannot read image: {path}") from exc
+
+
+def coerce_image(source: str | Path | Image.Image | np.ndarray) -> Image.Image:
+    if isinstance(source, Image.Image):
+        return source
+    if isinstance(source, np.ndarray):
+        return image_from_array(source)
+    if isinstance(source, (str, Path)):
+        return load_image(source)
+    raise RecognitionError(
+        "image must be a file path, a PIL image, or a numpy array"
+    )
+
+
+def looks_like_line(image: Image.Image, max_height: int = 96) -> bool:
+    width, height = image.size
+    return height <= max_height and width >= height

wormocr/runtime.py

@@ -0,0 +1,56 @@
+"""Device selection and a small environment report."""
+
+from __future__ import annotations
+
+import os
+import platform
+import shutil
+
+import torch
+
+from wormocr.exceptions import WormOCRError
+
+
+def resolve_device(name: str | None) -> torch.device:
+    requested = (name or "auto").strip().lower()
+    if requested == "auto":
+        return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    if requested.startswith("cuda") and not torch.cuda.is_available():
+        raise WormOCRError(
+            "CUDA was requested, but this PyTorch installation cannot see a GPU. "
+            f"torch {torch.__version__} is the build in use. Run with --device cpu."
+        )
+    try:
+        return torch.device(requested)
+    except (TypeError, RuntimeError) as exc:
+        raise WormOCRError(f"unknown device {name!r}") from exc
+
+
+def environment_report() -> dict:
+    report: dict = {
+        "python": platform.python_version(),
+        "platform": platform.platform(),
+        "cpu": platform.processor() or os.environ.get("PROCESSOR_IDENTIFIER", "unknown"),
+        "cpu_count": os.cpu_count(),
+        "torch": torch.__version__,
+        "cuda_available": bool(torch.cuda.is_available()),
+        "cuda_build": torch.version.cuda,
+        "torch_threads": int(torch.get_num_threads()),
+        "mkldnn": bool(torch.backends.mkldnn.is_available()),
+    }
+    if torch.cuda.is_available():
+        report["gpu"] = torch.cuda.get_device_name(0)
+    try:
+        import onnxruntime
+
+        report["onnxruntime"] = onnxruntime.__version__
+    except ImportError:
+        report["onnxruntime"] = None
+    try:
+        import onnx
+
+        report["onnx"] = onnx.__version__
+    except ImportError:
+        report["onnx"] = None
+    report["tesseract"] = shutil.which("tesseract")
+    return report

wormocr/training/__init__.py

@@ -0,0 +1,5 @@
+"""Training entry points."""
+
+from wormocr.training.loop import TrainResult, evaluate_model, train
+
+__all__ = ["TrainResult", "evaluate_model", "train"]