| """Text-line recognizer: CNN feature extractor -> Transformer encoder -> CTC head. |
| |
| Input is a grayscale line crop of height 32 and variable width. The CNN collapses height to 1 and |
| downsamples width by 4, so each output step covers 4 input columns. Position comes from a depthwise |
| convolution (relative, translation-equivariant) rather than an absolute table, so lines longer than |
| any seen in training are handled the same way as short ones. |
| """ |
|
|
| import torch |
| from torch import nn |
|
|
| from ..charset import NUM_CLASSES |
| from .common import ConvBNAct, TransformerBlock, padding_bias |
|
|
| LINE_HEIGHT = 32 |
| WIDTH_STRIDE = 4 |
| X_STRETCH = 1.25 |
|
|
|
|
| def _separable(cin, cout): |
| """Depthwise 3x3 then pointwise 1x1: about a ninth of a full 3x3 convolution's arithmetic.""" |
| return nn.Sequential(ConvBNAct(cin, cin, groups=cin), ConvBNAct(cin, cout, k=1)) |
|
|
|
|
| class Recognizer(nn.Module): |
| """cnn="full" (v1-v4) or "separable": the same shape of front-end with depthwise-separable convolutions, |
| measured 2.5x (4 threads) to 3x (1 thread) faster per page on CPU, where the full CNN is ~70% of the time.""" |
|
|
| def __init__(self, num_classes=NUM_CLASSES, d=192, layers=5, heads=6, ff=672, dropout=0.1, cnn="full"): |
| super().__init__() |
| conv = ConvBNAct if cnn == "full" else _separable |
| self.cnn = nn.Sequential( |
| ConvBNAct(1, 32), nn.MaxPool2d(2, 2), |
| conv(32, 64), nn.MaxPool2d(2, 2), |
| conv(64, 128), conv(128, 128), nn.MaxPool2d((2, 1), (2, 1)), |
| conv(128, 192), conv(192, 192), nn.MaxPool2d((2, 1), (2, 1)), |
| ConvBNAct(192, 256, k=(2, 3), p=(0, 1)), |
| ) |
| self.proj = nn.Linear(256, d) |
| self.pos = nn.Conv1d(d, d, 7, padding=3, groups=d) |
| self.blocks = nn.ModuleList(TransformerBlock(d, heads, ff, dropout) for _ in range(layers)) |
| self.ln = nn.LayerNorm(d) |
| self.head = nn.Linear(d, num_classes) |
|
|
| def forward(self, images, step_mask=None): |
| """images: (B, 1, 32, W) in [0, 1], ink dark. step_mask: optional (B, W/4) bool, True = padding. |
| |
| Returns logits (B, T, C) with T = W / 4. |
| """ |
| f = self.cnn(images - 0.5).squeeze(2).transpose(1, 2) |
| x = self.proj(f) |
| if step_mask is not None: |
| x = x.masked_fill(step_mask[..., None], 0.0) |
| x = x + self.pos(x.transpose(1, 2)).transpose(1, 2) |
| bias = None if step_mask is None else padding_bias(step_mask, x.dtype) |
| for blk in self.blocks: |
| x = blk(x, bias) |
| return self.head(self.ln(x)) |
|
|
