File size: 4,407 Bytes
824cafa | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 | """
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))
|