from __future__ import annotations import torch from torch import nn class AHCDSmallCNN(nn.Module): """Compact character classifier used in the AHCD notebook.""" def __init__(self, num_classes: int = 28) -> None: super().__init__() self.features = nn.Sequential( nn.Conv2d(1, 32, 3, padding=1), nn.ReLU(inplace=True), nn.MaxPool2d(2), nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(inplace=True), nn.MaxPool2d(2), nn.Conv2d(64, 96, 3, padding=1), nn.ReLU(inplace=True), nn.AdaptiveAvgPool2d((4, 4)), ) self.classifier = nn.Sequential( nn.Flatten(), nn.Linear(96 * 4 * 4, 128), nn.ReLU(inplace=True), nn.Dropout(0.2), nn.Linear(128, num_classes), ) def forward(self, images: torch.Tensor) -> torch.Tensor: return self.classifier(self.features(images)) class SimpleCRNN(nn.Module): """Small CTC recogniser for line-level handwriting experiments.""" def __init__(self, num_classes: int, hidden_size: int = 128) -> None: super().__init__() self.cnn = nn.Sequential( nn.Conv2d(1, 32, 3, padding=1), nn.ReLU(inplace=True), nn.MaxPool2d((2, 2)), nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(inplace=True), nn.MaxPool2d((2, 2)), nn.Conv2d(64, 128, 3, padding=1), nn.ReLU(inplace=True), nn.AdaptiveAvgPool2d((1, None)), ) self.rnn = nn.LSTM( input_size=128, hidden_size=hidden_size, num_layers=2, bidirectional=True, batch_first=False, dropout=0.1, ) self.classifier = nn.Linear(hidden_size * 2, num_classes) def forward(self, images: torch.Tensor) -> torch.Tensor: features = self.cnn(images).squeeze(2) sequence = features.permute(2, 0, 1).contiguous() sequence, _ = self.rnn(sequence) return self.classifier(sequence) def best_available_device() -> torch.device: if torch.cuda.is_available(): return torch.device("cuda") if torch.backends.mps.is_available(): return torch.device("mps") return torch.device("cpu") # --------------------------------------------------------------------------- # Final model: attention-pooling CNN-Transformer-CTC backbone. # This is the visual recogniser behind the headline result # (test CER 0.1331 with the full calibrated-decode + reranking pipeline; # test CER ~0.170 with plain greedy CTC decoding on this backbone alone). # Architecture copied verbatim from the project source so checkpoints load exactly. # --------------------------------------------------------------------------- import math class SinusoidalPositionalEncoding(nn.Module): def __init__(self, d_model: int, max_len: int = 8192, dropout: float = 0.1) -> None: super().__init__() self.dropout = nn.Dropout(dropout) positions = torch.arange(max_len, dtype=torch.float32).unsqueeze(1) div_terms = torch.exp(torch.arange(0, d_model, 2, dtype=torch.float32) * (-math.log(10000.0) / d_model)) encoding = torch.zeros(max_len, d_model, dtype=torch.float32) encoding[:, 0::2] = torch.sin(positions * div_terms) encoding[:, 1::2] = torch.cos(positions * div_terms) self.register_buffer("encoding", encoding.unsqueeze(0), persistent=False) def forward(self, sequence: torch.Tensor) -> torch.Tensor: if sequence.size(1) > self.encoding.size(1): raise ValueError( f"Sequence length {sequence.size(1)} exceeds max positional length {self.encoding.size(1)}" ) sequence = sequence + self.encoding[:, : sequence.size(1), :].to(dtype=sequence.dtype) return self.dropout(sequence) class HeightAttentionPool(nn.Module): def __init__(self, channels: int) -> None: super().__init__() hidden = max(16, channels // 4) output = nn.Conv2d(hidden, 1, kernel_size=1) nn.init.zeros_(output.weight) nn.init.zeros_(output.bias) self.score = nn.Sequential( nn.Conv2d(channels, hidden, kernel_size=1), nn.GELU(), output, ) def forward(self, features: torch.Tensor) -> torch.Tensor: weights = torch.softmax(self.score(features), dim=2) return (features * weights).sum(dim=2) class CNNTransformerCTC(nn.Module): def __init__( self, num_classes: int, d_model: int = 256, num_encoder_layers: int = 4, num_heads: int = 8, dim_feedforward: int = 1024, dropout: float = 0.1, height_pooling: str = "mean", ) -> None: super().__init__() if height_pooling not in {"mean", "attention", "mean_max"}: raise ValueError("height_pooling must be one of: mean, attention, mean_max") self.height_pooling = height_pooling self.frontend = nn.Sequential( nn.Conv2d(1, 64, kernel_size=3, padding=1), nn.BatchNorm2d(64), nn.GELU(), nn.MaxPool2d(kernel_size=2, stride=2), nn.Conv2d(64, 128, kernel_size=3, padding=1), nn.BatchNorm2d(128), nn.GELU(), nn.MaxPool2d(kernel_size=2, stride=2), nn.Conv2d(128, 256, kernel_size=3, padding=1), nn.BatchNorm2d(256), nn.GELU(), nn.Conv2d(256, 256, kernel_size=3, padding=1), nn.BatchNorm2d(256), nn.GELU(), ) if height_pooling == "attention": self.height_attention = HeightAttentionPool(256) pooled_channels = 256 elif height_pooling == "mean_max": pooled_channels = 512 else: pooled_channels = 256 self.projection = nn.Linear(pooled_channels, d_model) self.position_encoding = SinusoidalPositionalEncoding(d_model=d_model, dropout=dropout) encoder_layer = nn.TransformerEncoderLayer( d_model=d_model, nhead=num_heads, dim_feedforward=dim_feedforward, dropout=dropout, activation="gelu", batch_first=True, norm_first=True, ) self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_encoder_layers) self.norm = nn.LayerNorm(d_model) self.classifier = nn.Linear(d_model, num_classes) @staticmethod def output_lengths(input_widths: torch.Tensor) -> torch.Tensor: lengths = torch.floor_divide(input_widths, 2) lengths = torch.floor_divide(lengths, 2) return torch.clamp(lengths, min=1).to(dtype=torch.long) def forward(self, images: torch.Tensor, input_widths: torch.Tensor | None = None) -> torch.Tensor: features = self.frontend(images) if self.height_pooling == "attention": pooled = self.height_attention(features) elif self.height_pooling == "mean_max": pooled = torch.cat([features.mean(dim=2), features.amax(dim=2)], dim=1) else: pooled = features.mean(dim=2) sequence = pooled.permute(0, 2, 1).contiguous() sequence = self.projection(sequence) sequence = self.position_encoding(sequence) padding_mask = None if input_widths is not None: output_lengths = self.output_lengths(input_widths).to(device=sequence.device) max_length = sequence.size(1) positions = torch.arange(max_length, device=sequence.device).unsqueeze(0) padding_mask = positions >= output_lengths.unsqueeze(1) encoded = self.encoder(sequence, src_key_padding_mask=padding_mask) encoded = self.norm(encoded) logits = self.classifier(encoded) return logits.permute(1, 0, 2).contiguous() def load_ctc_backbone(checkpoint_path, device: torch.device | None = None): """Load the published CNN-Transformer-CTC backbone and its charset. Returns (model, charset, model_config). The model is set to eval() mode. """ from .ocr_helpers import Charset device = device or best_available_device() checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False) charset = Charset(checkpoint["charset"]) config = checkpoint.get("model_config", {}) model = CNNTransformerCTC( num_classes=charset.vocab_size, d_model=config.get("d_model", 256), num_encoder_layers=config.get("transformer_layers", 4), num_heads=config.get("transformer_heads", 8), dim_feedforward=config.get("dim_feedforward", 1024), dropout=config.get("dropout", 0.1), height_pooling=config.get("height_pooling", "attention"), ) model.load_state_dict(checkpoint["model_state"]) model.eval().to(device) return model, charset, config @torch.no_grad() def recognise_line(model, charset, image, image_height: int = 96, device: torch.device | None = None) -> str: """Run greedy CTC OCR on a single PIL line image and return decoded text.""" from PIL import Image from .ocr_helpers import ImageConfig, preprocess_image, greedy_ctc_decode if isinstance(image, (str, bytes)) or hasattr(image, "__fspath__"): image = Image.open(image) device = device or next(model.parameters()).device array = preprocess_image(image, ImageConfig(target_height=image_height, grayscale=True, normalize=True)) tensor = torch.tensor(array, dtype=torch.float32)[None, None, :, :].to(device) logits = model(tensor) # (T, 1, num_classes) best = logits.argmax(dim=-1).squeeze(1) # (T,) ids = greedy_ctc_decode(best.cpu().tolist(), blank_id=charset.blank_id) return charset.decode(ids)