Upload train/train_sequence.py with huggingface_hub

train/train_sequence.py

@@ -0,0 +1,271 @@
+"""
+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)