Upload src/recognize.py with huggingface_hub

src/recognize.py

@@ -0,0 +1,126 @@
+"""
+OCR recognition: preprocessed ROI image -> 8-character passport string.
+
+Model: CRNN (CNN + biGRU + 8 position heads).
+Input:  grayscale image of any size (resized internally to 48x256).
+Output: (string, confidence)  e.g. ("НР430098", 0.97)
+"""
+
+from pathlib import Path
+import cv2
+import numpy as np
+import torch
+import torch.nn as nn
+
+
+# -- Constants -----------------------------------------------------------------
+
+DIGITS         = list("0123456789")
+SERIES_LETTERS = list("АВЕКМНОРСТИЮ")
+ALL_CHARS      = DIGITS + SERIES_LETTERS   # 22 classes
+
+N_CHARS = 8
+SEQ_H   = 48
+SEQ_W   = 256
+
+CONFIDENCE_THRESHOLD = 0.50   # below this -> flagged as unreadable
+
+
+# -- Model architecture (must match train/train_sequence.py) -------------------
+
+class SequenceCNN(nn.Module):
+    RNN_SLOTS  = 24
+    CHAR_SLOTS = [1, 4, 7, 10, 13, 16, 19, 22]
+
+    def __init__(self, n_classes: int = len(ALL_CHARS)):
+        super().__init__()
+        self.features = nn.Sequential(
+            nn.Conv2d(1, 32, 3, padding=1), nn.BatchNorm2d(32), nn.ReLU(),
+            nn.MaxPool2d(2),
+            nn.Conv2d(32, 64, 3, padding=1), nn.BatchNorm2d(64), nn.ReLU(),
+            nn.MaxPool2d(2),
+            nn.Conv2d(64, 128, 3, padding=1), nn.BatchNorm2d(128), nn.ReLU(),
+            nn.MaxPool2d(2),
+            nn.Conv2d(128, 128, 3, padding=1), nn.BatchNorm2d(128), nn.ReLU(),
+        )
+        self.pool  = nn.AdaptiveAvgPool2d((1, self.RNN_SLOTS))
+        self.rnn   = nn.GRU(128, 128, num_layers=1, batch_first=True,
+                            bidirectional=True)
+        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)
+        feat = self.pool(feat).squeeze(2)
+        feat = feat.permute(0, 2, 1)
+        rnn_out, _ = self.rnn(feat)
+        return torch.stack(
+            [self.heads[i](rnn_out[:, self.CHAR_SLOTS[i], :])
+             for i in range(N_CHARS)],
+            dim=1)
+
+
+# -- Load ----------------------------------------------------------------------
+
+def load_recognizer(model_path: str | Path):
+    """Load CRNN model from checkpoint. Returns (model, meta_dict)."""
+    ck = torch.load(str(model_path), map_location="cpu", weights_only=True)
+    n_classes = len(ck["all_chars"])
+    model = SequenceCNN(n_classes=n_classes)
+    model.load_state_dict(ck["model_state"])
+    for m in model.modules():
+        if hasattr(m, "training"):
+            m.training = False
+    return model, ck
+
+
+# -- Inference -----------------------------------------------------------------
+
+def recognize(image: np.ndarray, model: SequenceCNN,
+              checkpoint: dict) -> tuple[str, float]:
+    """
+    Recognize 8-character passport string from a preprocessed ROI image.
+
+    Args:
+        image:      grayscale numpy array (any size, will be resized).
+        model:      loaded SequenceCNN.
+        checkpoint: dict returned by load_recognizer (contains idx2char etc.).
+
+    Returns:
+        (prediction, confidence)
+        confidence < CONFIDENCE_THRESHOLD means the image is likely unreadable.
+    """
+    all_chars  = checkpoint["all_chars"]
+    idx2char   = {int(k): v for k, v in checkpoint["idx2char"].items()}
+    seq_h      = checkpoint.get("seq_h", SEQ_H)
+    seq_w      = checkpoint.get("seq_w", SEQ_W)
+
+    letter_set = set(SERIES_LETTERS)
+    digit_set  = set(DIGITS)
+    letter_idx = [i for i, c in enumerate(all_chars) if c in letter_set]
+    digit_idx  = [i for i, c in enumerate(all_chars) if c in digit_set]
+
+    img = cv2.resize(image, (seq_w, seq_h), interpolation=cv2.INTER_AREA)
+    arr = img.astype(np.float32) / 255.0
+    x   = torch.tensor(arr).unsqueeze(0).unsqueeze(0)   # (1, 1, H, W)
+
+    with torch.no_grad():
+        logits = model(x)[0]   # (8, n_classes)
+
+    result, confidences = [], []
+    for pos in range(N_CHARS):
+        mask = torch.full((len(all_chars),), float("-inf"))
+        for idx in (letter_idx if pos < 2 else digit_idx):
+            mask[idx] = 0.0
+        masked = logits[pos] + mask
+        probs  = torch.softmax(masked, dim=0)
+        best   = probs.argmax().item()
+        result.append(idx2char[best])
+        confidences.append(probs[best].item())
+
+    return "".join(result), float(np.mean(confidences))