83 lines
3.5 KiB
Python
83 lines
3.5 KiB
Python
"""
|
|
Model architecture - this is the corrected version of the flowchart's "AI" box.
|
|
|
|
On the diagram, CNN and BiLSTM were drawn as two independent branches off "AI".
|
|
For text recognition that's not how it works: they're sequential stages of one
|
|
pipeline (a CRNN), plus a CTC layer the diagram was missing entirely:
|
|
|
|
image -> CNN (visual features, replaces "Līnijas atpazišana" /
|
|
"Pareiza novietošana" - i.e. stroke patterns + spatial layout)
|
|
-> reshape rows->sequence
|
|
-> BiLSTM x2 (context over the sequence, replaces "Vārdu atpazīšana" /
|
|
"Pieturzīmes" / "Cipari" - i.e. word/punctuation/digit recognition
|
|
needs context from neighboring characters, which only the LSTM
|
|
stage gives you)
|
|
-> Linear classifier over the alphabet
|
|
-> CTC loss/decoding (aligns variable-length predictions to text
|
|
without needing per-character bounding boxes)
|
|
|
|
CTC is what makes "Teksta zonas atpazīšana????" mostly unnecessary for the
|
|
input side too: you don't need to segment characters up front, just feed
|
|
whole line images in.
|
|
"""
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
|
|
class CNNBackbone(nn.Module):
|
|
"""Reduces a [B,1,32,W] line image to a [B,512,1,W'] feature map."""
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.net = nn.Sequential(
|
|
nn.Conv2d(1, 64, 3, 1, 1), nn.ReLU(inplace=True), nn.MaxPool2d(2, 2), # 32x W -> 16 x W/2
|
|
nn.Conv2d(64, 128, 3, 1, 1), nn.ReLU(inplace=True), nn.MaxPool2d(2, 2), # -> 8 x W/4
|
|
nn.Conv2d(128, 256, 3, 1, 1), nn.ReLU(inplace=True),
|
|
nn.Conv2d(256, 256, 3, 1, 1), nn.ReLU(inplace=True), nn.MaxPool2d((2, 1), (2, 1)), # -> 4 x W/4
|
|
nn.Conv2d(256, 512, 3, 1, 1), nn.BatchNorm2d(512), nn.ReLU(inplace=True),
|
|
nn.Conv2d(512, 512, 3, 1, 1), nn.BatchNorm2d(512), nn.ReLU(inplace=True), nn.MaxPool2d((2, 1), (2, 1)), # -> 2 x W/4
|
|
nn.Conv2d(512, 512, 2, 1, 0), nn.ReLU(inplace=True), # -> 1 x (W/4 - 1)
|
|
)
|
|
|
|
def forward(self, x):
|
|
return self.net(x)
|
|
|
|
|
|
class BiLSTMHead(nn.Module):
|
|
def __init__(self, in_dim: int, hidden: int, num_classes: int, num_layers: int = 2):
|
|
super().__init__()
|
|
self.lstm = nn.LSTM(
|
|
in_dim, hidden, num_layers=num_layers, bidirectional=True, batch_first=False
|
|
)
|
|
self.fc = nn.Linear(hidden * 2, num_classes)
|
|
|
|
def forward(self, x):
|
|
# x: [T, B, in_dim]
|
|
out, _ = self.lstm(x)
|
|
return self.fc(out) # [T, B, num_classes]
|
|
|
|
|
|
class CRNN(nn.Module):
|
|
def __init__(self, num_classes: int, lstm_hidden: int = 256):
|
|
super().__init__()
|
|
self.cnn = CNNBackbone()
|
|
self.rnn = BiLSTMHead(in_dim=512, hidden=lstm_hidden, num_classes=num_classes)
|
|
|
|
def forward(self, images: torch.Tensor) -> torch.Tensor:
|
|
"""
|
|
images: [B, 1, 32, W]
|
|
returns log-probs: [T, B, num_classes], T = sequence length after CNN downsampling
|
|
"""
|
|
feats = self.cnn(images) # [B, 512, 1, W']
|
|
feats = feats.squeeze(2) # [B, 512, W']
|
|
feats = feats.permute(2, 0, 1) # [W'(=T), B, 512]
|
|
logits = self.rnn(feats) # [T, B, num_classes]
|
|
return logits.log_softmax(dim=2)
|
|
|
|
def output_length(self, input_width: int) -> int:
|
|
"""Sequence length T the CNN produces for a given input image width (for CTC input_lengths)."""
|
|
w = input_width // 2 // 2 # two stride-2 pools
|
|
w = w - 1 # final kernel=2,stride=1,pad=0 conv
|
|
return max(w, 1)
|