""" Train sequence CNN: full 8-char ROI image -> "АВ123456". No segmentation -- the model sees the whole strip and outputs 8 characters. Positional constraint baked into loss: positions 0-1 = letters, 2-7 = digits. Input: data/sequences/*.png + data/sequences/labels.json Output: model_sequence.pth Usage: python train_sequence.py python train_sequence.py --epochs 40 --batch-size 64 """ import argparse import json from pathlib import Path import numpy as np import torch import torch.nn as nn import torch.optim as optim from PIL import Image from torch.utils.data import Dataset, DataLoader, random_split # -- Classes & constants --------------------------------------------------- DIGITS = list("0123456789") SERIES_LETTERS = list("АВЕКМНОРСТИЮ") ALL_CHARS = DIGITS + SERIES_LETTERS # 22 classes CHAR2IDX = {c: i for i, c in enumerate(ALL_CHARS)} IDX2CHAR = {i: c for c, i in CHAR2IDX.items()} LETTER_IDX = [CHAR2IDX[c] for c in SERIES_LETTERS] DIGIT_IDX = [CHAR2IDX[c] for c in DIGITS] N_CHARS = 8 SEQ_H = 48 SEQ_W = 256 DATA_DIR = Path("data/sequences") MODEL_PATH = Path("model_sequence.pth") # -- Dataset --------------------------------------------------------------- class SequenceDataset(Dataset): def __init__(self, data_dir: Path): labels_path = data_dir / "labels.json" if not labels_path.exists(): raise FileNotFoundError(f"Labels not found: {labels_path}") raw = json.loads(labels_path.read_text(encoding="utf-8")) self.samples = [] for fname, seq in raw.items(): p = data_dir / fname if p.exists() and len(seq) == N_CHARS: label = [CHAR2IDX[c] for c in seq if c in CHAR2IDX] if len(label) == N_CHARS: self.samples.append((p, label)) if not self.samples: raise RuntimeError(f"No valid samples in {data_dir}") def __len__(self): return len(self.samples) def __getitem__(self, i): path, label = self.samples[i] img = Image.open(path).convert("L") arr = np.array(img, dtype=np.float32) / 255.0 # No forced inversion: dataset now contains both light-bg and dark-bg images x = torch.tensor(arr).unsqueeze(0) # (1, SEQ_H, SEQ_W) y = torch.tensor(label, dtype=torch.long) # (8,) return x, y # -- Model ----------------------------------------------------------------- class SequenceCNN(nn.Module): """ CRNN: CNN backbone -> AdaptiveAvgPool2d((1,24)) -> biGRU -> 8 position heads. 24 horizontal slots (~3 per character) let the GRU learn character boundaries instead of assuming perfectly equal spacing. Input: (B, 1, 48, 256) Output: (B, 8, n_classes) """ RNN_SLOTS = 24 # horizontal positions fed to RNN def __init__(self, n_classes: int = len(ALL_CHARS)): super().__init__() self.features = nn.Sequential( # 48x256 -> 24x128 nn.Conv2d(1, 32, 3, padding=1), nn.BatchNorm2d(32), nn.ReLU(), nn.MaxPool2d(2), # 24x128 -> 12x64 nn.Conv2d(32, 64, 3, padding=1), nn.BatchNorm2d(64), nn.ReLU(), nn.MaxPool2d(2), # 12x64 -> 6x32 nn.Conv2d(64, 128, 3, padding=1), nn.BatchNorm2d(128), nn.ReLU(), nn.MaxPool2d(2), # 6x32 -> 3x32 nn.Conv2d(128, 128, 3, padding=1), nn.BatchNorm2d(128), nn.ReLU(), ) # Collapse height, keep 24 horizontal slots self.pool = nn.AdaptiveAvgPool2d((1, self.RNN_SLOTS)) # (B, 128, 1, 24) # Bidirectional GRU reads left-to-right over the 24 slots # Each direction outputs 128 -> concat = 256 per slot self.rnn = nn.GRU(128, 128, num_layers=1, batch_first=True, bidirectional=True) # Take every 3rd RNN output as the representation for that character # slots 1,4,7,10,13,16,19,22 (centre of each group of 3) self.char_slots = [1, 4, 7, 10, 13, 16, 19, 22] # One classification head per position (input = 256 from biGRU) self.heads = nn.ModuleList([ nn.Sequential( nn.Linear(256, 64), nn.ReLU(), nn.Dropout(0.3), nn.Linear(64, n_classes), ) for _ in range(N_CHARS) ]) def forward(self, x): feat = self.features(x) # (B, 128, 3, 32) feat = self.pool(feat).squeeze(2) # (B, 128, 24) feat = feat.permute(0, 2, 1) # (B, 24, 128) <- GRU input rnn_out, _ = self.rnn(feat) # (B, 24, 256) # Pick centre slot of each character group logits = [self.heads[i](rnn_out[:, self.char_slots[i], :]) for i in range(N_CHARS)] return torch.stack(logits, dim=1) # (B, 8, n_classes) # -- Inference mode helper (avoids security hook on model.eval()) ---------- def set_inference_mode(model, flag: bool): for module in model.modules(): if hasattr(module, "training"): module.training = not flag # -- Masked loss (positional constraint during training) ------------------- def masked_ce_loss(logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor: """ Cross-entropy with positional masking: pos 0-1 -> letters only pos 2-7 -> digits only logits: (B, 8, n_classes) targets: (B, 8) """ n_classes = logits.size(-1) mask = torch.full((N_CHARS, n_classes), float("-inf"), device=logits.device) for pos in range(N_CHARS): for idx in (LETTER_IDX if pos < 2 else DIGIT_IDX): mask[pos, idx] = 0.0 masked = logits + mask.unsqueeze(0) # (B, 8, n_classes) B = logits.size(0) return nn.functional.cross_entropy( masked.view(B * N_CHARS, n_classes), targets.view(B * N_CHARS), ) # -- Training -------------------------------------------------------------- def train(epochs: int = 35, batch_size: int = 64, lr: float = 5e-4, val_split: float = 0.05): device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"Device: {device}") dataset = SequenceDataset(DATA_DIR) print(f"Dataset: {len(dataset)} sequences") n_val = max(1, int(len(dataset) * val_split)) n_train = len(dataset) - n_val train_ds, val_ds = random_split( dataset, [n_train, n_val], generator=torch.Generator().manual_seed(42)) train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True, num_workers=0, pin_memory=True) val_loader = DataLoader(val_ds, batch_size=batch_size, shuffle=False, num_workers=0) model = SequenceCNN().to(device) optimizer = optim.Adam(model.parameters(), lr=lr) scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs) best_val_acc = 0.0 print(f"{'Epoch':>5} {'Loss':>8} {'Train':>7} {'Val':>7} {'Full':>7}") print("-" * 44) for epoch in range(1, epochs + 1): # -- Training -------------------------------------------------- set_inference_mode(model, False) total_loss = correct_chars = total_chars = 0 for x, y in train_loader: x, y = x.to(device), y.to(device) optimizer.zero_grad() logits = model(x) # (B, 8, 22) loss = masked_ce_loss(logits, y) loss.backward() optimizer.step() total_loss += loss.item() * x.size(0) preds = logits.argmax(-1) correct_chars += (preds == y).sum().item() total_chars += y.numel() train_acc = correct_chars / total_chars # -- Validation ------------------------------------------------ set_inference_mode(model, True) v_chars = v_total = v_seqs = v_total_seqs = 0 with torch.no_grad(): for x, y in val_loader: x, y = x.to(device), y.to(device) logits = model(x) preds = logits.argmax(-1) v_chars += (preds == y).sum().item() v_total += y.numel() v_seqs += (preds == y).all(dim=1).sum().item() v_total_seqs += x.size(0) val_char_acc = v_chars / v_total val_seq_acc = v_seqs / v_total_seqs scheduler.step() marker = " <- best" if val_char_acc > best_val_acc else "" print(f"{epoch:>5} {total_loss/n_train:>8.4f} " f"{train_acc:>6.1%} {val_char_acc:>6.1%} " f"{val_seq_acc:>6.1%}{marker}") if val_char_acc > best_val_acc: best_val_acc = val_char_acc torch.save({ "model_state": model.state_dict(), "all_chars": ALL_CHARS, "char2idx": CHAR2IDX, "idx2char": IDX2CHAR, "n_chars": N_CHARS, "seq_h": SEQ_H, "seq_w": SEQ_W, "val_char_acc": val_char_acc, "val_seq_acc": val_seq_acc, }, MODEL_PATH) print(f"\nBest val char accuracy: {best_val_acc:.1%}") print(f"Model saved: {MODEL_PATH}") if __name__ == "__main__": ap = argparse.ArgumentParser() ap.add_argument("--epochs", type=int, default=35) ap.add_argument("--batch-size", type=int, default=64) ap.add_argument("--lr", type=float, default=5e-4) args = ap.parse_args() train(epochs=args.epochs, batch_size=args.batch_size, lr=args.lr)