Handwriting recognition

pip install torch==2.5.1 opencv-python==4.10.0.84 numpy==1.26.4

import cv2
import numpy as np
import torch
import torch.nn as nn

FONT = {
    "A": ["..#..", ".#.#.", "#...#", "#...#", "#####", "#...#", "#...#"],
    "C": [".###.", "#...#", "#....", "#....", "#....", "#...#", ".###."],
    "E": ["#####", "#....", "#....", "####.", "#....", "#....", "#####"],
    "L": ["#....", "#....", "#....", "#....", "#....", "#....", "#####"],
    "N": ["#...#", "##..#", "#.#.#", "#.#.#", "#..##", "#...#", "#...#"],
    "R": ["####.", "#...#", "#...#", "####.", "#.#..", "#..#.", "#...#"],
    "T": ["#####", "..#..", "..#..", "..#..", "..#..", "..#..", "..#.."],
    "1": ["..#..", ".##..", "..#..", "..#..", "..#..", "..#..", ".###."],
    "9": [".###.", "#...#", "#...#", ".####", "....#", "#...#", ".###."],
}
ALPHABET = "".join(sorted(FONT))
GLYPH = {c: np.array([[x == "#" for x in r] for r in FONT[c]], np.float32) for c in FONT}
IMG_H, IMG_W, SCALE = 24, 96, 2


def draw(word, rng, messy):
    """messy=0 is neat print. Higher values slant, overlap and wobble the letters."""
    canvas = np.zeros((IMG_H, IMG_W), np.float32)
    x = 4 + int(rng.integers(0, 4))
    for ch in word:
        g = np.kron(GLYPH[ch], np.ones((SCALE, SCALE), np.float32))
        h, w = g.shape
        if messy:
            thick = int(rng.integers(0, 2))                     # pen pressure varies
            if thick:
                g = cv2.dilate(g, np.ones((2, 2), np.float32))
        top = (IMG_H - h) // 2 + int(rng.integers(-messy, messy + 1))
        top = max(0, min(IMG_H - h, top))
        if x + w > IMG_W:
            break
        canvas[top:top + h, x:x + w] = np.maximum(canvas[top:top + h, x:x + w], g)
        gap = 2 - int(rng.integers(0, messy + 1))               # letters start to touch
        x += w + max(gap, -2)
    if messy:
        slant = float(rng.uniform(-0.35, 0.35)) * (messy / 3.0)  # a slanted hand
        M = np.array([[1, slant, -slant * IMG_H / 2], [0, 1, 0]], np.float32)
        canvas = cv2.warpAffine(canvas, M, (IMG_W, IMG_H))
    return np.clip(canvas, 0, 1)


def batch(n, rng, messy, lo=3, hi=6):
    words = ["".join(rng.choice(list(ALPHABET), rng.integers(lo, hi + 1))) for _ in range(n)]
    imgs = torch.tensor(np.stack([draw(w, rng, messy) for w in words]))[:, None]
    targets = torch.tensor([ALPHABET.index(c) + 1 for w in words for c in w])
    lengths = torch.tensor([len(w) for w in words])
    return words, imgs, targets, lengths


class CRNN(nn.Module):
    def __init__(self, n_classes):
        super().__init__()
        self.cnn = nn.Sequential(
            nn.Conv2d(1, 16, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
            nn.Conv2d(16, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d((2, 1)),
            nn.Conv2d(32, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d((2, 1)),
        )
        self.rnn = nn.LSTM(32 * 3, 64, bidirectional=True)
        self.head = nn.Linear(128, n_classes)

    def forward(self, x):
        f = self.cnn(x).permute(3, 0, 1, 2).flatten(2)
        return self.head(self.rnn(f)[0])


def decode(logits):
    out = []
    for row in logits.argmax(2).T.tolist():
        chars, prev = [], -1
        for k in row:
            if k != prev and k != 0:
                chars.append(ALPHABET[k - 1])
            prev = k
        out.append("".join(chars))
    return out


def train(messy, steps=900):
    torch.manual_seed(0)
    rng = np.random.default_rng(1)
    model = CRNN(len(ALPHABET) + 1)
    loss_fn = nn.CTCLoss(blank=0, zero_infinity=True)
    opt = torch.optim.Adam(model.parameters(), lr=3e-3)
    for step in range(steps):
        _, imgs, targets, lengths = batch(16, rng, messy)
        logits = model(imgs)
        lens = torch.full((16,), logits.shape[0], dtype=torch.long)
        loss = loss_fn(logits.log_softmax(2), targets, lens, lengths)
        opt.zero_grad()
        loss.backward()
        opt.step()
    model.eval()
    return model, loss.item()


def evaluate(model, messy, n=200):
    rng = np.random.default_rng(99)
    words, imgs, _, _ = batch(n, rng, messy)
    with torch.no_grad():
        read = decode(model(imgs))
    exact = sum(a == b for a, b in zip(words, read)) / n
    return exact, list(zip(words, read))[:4]


print("the same word, printed neatly and written messily:")
for label, messy in [("print", 0), ("messy", 3)]:
    img = draw("RENT", np.random.default_rng(7), messy)
    print(f"  {label}:")
    for row in img[3:22:2]:
        print("    " + "".join("#" if v > 0.5 else "." for v in row[:62]))

print("\ntraining a reader on NEAT PRINT only ...")
printed_model, loss_p = train(0)
print(f"  final training loss {loss_p:.3f}")
for messy in [0, 1, 2, 3]:
    acc, sample = evaluate(printed_model, messy)
    print(f"  tested at messiness {messy}: exact-word accuracy {acc:.3f}   e.g. {sample[0]}")

print("\ntraining the SAME network on messy writing instead ...")
messy_model, loss_m = train(3)
print(f"  final training loss {loss_m:.3f}")
for messy in [0, 3]:
    acc, sample = evaluate(messy_model, messy)
    print(f"  tested at messiness {messy}: exact-word accuracy {acc:.3f}   e.g. {sample[0]}")