| """Satr — Xattot's text-line recognizer (PyTorch definition, for fine-tuning). |
| |
| Inference does not need this file: `satr/satr.onnx` runs with onnxruntime only. |
| Architecture: CNN (height 48 -> 1, width / 4) + 2 Transformer blocks (d=192, 4 heads) + linear head, trained with CTC. |
| 3.26M parameters. Classes are copied unchanged from the training notebook, so the weights load with strict=True. |
| |
| import torch, satr_model |
| model, chars = satr_model.load("satr/satr.safetensors", "satr/vocab.json") |
| x = torch.from_numpy(line_gray_uint8[None, None]) # (1, 1, 48, W), white background, height 48 |
| print(satr_model.decode(model(x).log_softmax(-1), chars)) |
| |
| Input: grayscale line image resized to height 48 (keep aspect ratio). CTC blank = index 0, class i = chars[i - 1]. |
| License: CC BY-NC-SA 4.0 (see LICENSE). Fine-tuned weights may be shared under the same license with attribution. |
| """ |
| import json |
|
|
| import torch |
| import torch.nn.functional as F |
| from torch import nn |
|
|
|
|
| class Blok(nn.Module): |
| def __init__(self, ci, co, s): |
| super().__init__() |
| self.c1 = nn.Conv2d(ci, co, 3, s, 1, bias=False); self.b1 = nn.BatchNorm2d(co) |
| self.c2 = nn.Conv2d(co, co, 3, 1, 1, bias=False); self.b2 = nn.BatchNorm2d(co) |
| self.sk = nn.Sequential(nn.Conv2d(ci, co, 1, s, bias=False), nn.BatchNorm2d(co)) if (ci != co or s != 1) else nn.Identity() |
|
|
| def forward(self, x): |
| return F.gelu(self.b2(self.c2(F.gelu(self.b1(self.c1(x))))) + self.sk(x)) |
|
|
|
|
| class TBlok(nn.Module): |
| |
| def __init__(self, d, h=4, ff=3, p=0.1): |
| super().__init__() |
| self.h, self.p = h, p |
| self.n1, self.qkv, self.o = nn.LayerNorm(d), nn.Linear(d, 3 * d), nn.Linear(d, d) |
| self.n2, self.ff = nn.LayerNorm(d), nn.Sequential(nn.Linear(d, d * ff), nn.GELU(), nn.Linear(d * ff, d)) |
| self.dr = nn.Dropout(p) |
|
|
| def forward(self, x): |
| q, k, v = self.qkv(self.n1(x)).unflatten(-1, (3, self.h, -1)).permute(2, 0, 3, 1, 4) |
| a = F.scaled_dot_product_attention(q, k, v, dropout_p=self.p if self.training else 0.0) |
| x = x + self.dr(self.o(a.transpose(1, 2).flatten(2))) |
| return x + self.dr(self.ff(self.n2(x))) |
|
|
|
|
| class Satr(nn.Module): |
| def __init__(self, n_cls, d=192): |
| super().__init__() |
| self.cnn = nn.Sequential( |
| nn.Conv2d(1, 32, 3, 1, 1, bias=False), nn.BatchNorm2d(32), nn.GELU(), |
| Blok(32, 48, (2, 2)), Blok(48, 48, 1), |
| Blok(48, 96, (2, 2)), Blok(96, 96, 1), |
| Blok(96, 144, (2, 1)), Blok(144, 144, 1), |
| Blok(144, d, (2, 1)), Blok(d, d, 1)) |
| self.yig = nn.Conv2d(d, d, (3, 1), bias=False) |
| self.pos = nn.Conv1d(d, d, 9, padding=4, groups=d) |
| self.tr = nn.Sequential(TBlok(d), TBlok(d)) |
| self.out = nn.Sequential(nn.LayerNorm(d), nn.Linear(d, n_cls)) |
|
|
| def forward(self, x): |
| x = (x.float() / 255.0 - 0.5) / 0.5 if x.dtype == torch.uint8 else x |
| f = self.yig(self.cnn(x)).squeeze(2) |
| f = f + self.pos(f) |
| return self.out(self.tr(f.transpose(1, 2))) |
|
|
|
|
| def load(weights="satr/satr.safetensors", vocab="satr/vocab.json", device="cpu"): |
| from safetensors.torch import load_file |
| chars = json.load(open(vocab, encoding="utf-8"))["chars"] |
| model = Satr(len(chars) + 1) |
| model.load_state_dict(load_file(weights, device="cpu"), strict=True) |
| return model.to(device).eval(), chars |
|
|
|
|
| def decode(logp, chars): |
| """Greedy CTC decoding of (N,T,C) scores.""" |
| out = [] |
| for s in logp.argmax(-1).cpu().numpy(): |
| p, t = 0, [] |
| for c in s: |
| if c != p and c != 0: |
| t.append(chars[c - 1]) |
| p = c |
| out.append("".join(t)) |
| return out |
|
|
