File size: 5,063 Bytes
aa861d5 e186a72 aa861d5 e186a72 aa861d5 e186a72 aa861d5 5347b34 aa861d5 5347b34 aa861d5 5347b34 aa861d5 e186a72 aa861d5 5347b34 aa861d5 e186a72 aa861d5 e186a72 5347b34 e186a72 5347b34 e186a72 | 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 128 129 130 131 132 133 134 135 136 137 138 139 140 141 | from __future__ import annotations
import math
import pickle
from functools import lru_cache
from pathlib import Path
from typing import Any
import numpy as np
import torch
import torch.nn as nn
from PIL import Image
IMAGE_HEIGHT = 64
IMAGE_WIDTH = 2048
HORIZONTAL_STRIDE = 4
RNN_HIDDEN = 256
RNN_LAYERS = 2
RNN_DROPOUT = 0.05
CHECKPOINT = Path(__file__).with_name("best-model.pt")
LANGUAGE_MODEL = Path(__file__).with_name("char-trigram-lm.pkl")
BEAM_CONFIG = {
"beam_width": 10,
"token_topk": 12,
"lm_weight": 0.4,
"token_bonus": 1.5,
}
class TemporalFeatureDropout(nn.Module):
def __init__(self, probability: float = 0.0):
super().__init__()
self.p = float(probability)
def forward(self, x: torch.Tensor) -> torch.Tensor:
if not self.training or self.p == 0.0:
return x
keep = x.new_empty((x.shape[0], x.shape[1], 1)).bernoulli_(1.0 - self.p)
return x * keep / (1.0 - self.p)
class KoshurCRNN(nn.Module):
def __init__(self, n_classes: int, temporal_dropout_p: float = 0.0):
super().__init__()
self.cnn = nn.Sequential(
nn.Conv2d(1, 48, 3, padding=1), nn.BatchNorm2d(48), nn.ReLU(), nn.MaxPool2d((2, 2)),
nn.Conv2d(48, 96, 3, padding=1), nn.BatchNorm2d(96), nn.ReLU(), nn.MaxPool2d((2, 2)),
nn.Conv2d(96, 160, 3, padding=1), nn.BatchNorm2d(160), nn.ReLU(), nn.MaxPool2d((2, 1)),
nn.Conv2d(160, 192, 3, padding=1), nn.BatchNorm2d(192), nn.ReLU(),
)
self.temporal_dropout = TemporalFeatureDropout(temporal_dropout_p)
self.rnn = nn.GRU(
192 * (IMAGE_HEIGHT // 8), RNN_HIDDEN, num_layers=RNN_LAYERS,
bidirectional=True, batch_first=True, dropout=RNN_DROPOUT,
)
self.head = nn.Linear(RNN_HIDDEN * 2, n_classes)
def forward(self, x: torch.Tensor) -> torch.Tensor:
z = self.cnn(x)
z = z.permute(0, 3, 1, 2).contiguous().flatten(2)
z = self.temporal_dropout(z)
z, _ = self.rnn(z)
return self.head(z).log_softmax(-1).permute(1, 0, 2)
@lru_cache(maxsize=1)
def load_model() -> tuple[KoshurCRNN, dict[int, str]]:
checkpoint = torch.load(CHECKPOINT, map_location="cpu", weights_only=False)
chars = list(checkpoint["chars"])
architecture = checkpoint.get("architecture") or {}
temporal_dropout_p = float(architecture.get("temporal_dropout_p", 0.0))
model = KoshurCRNN(len(chars) + 1, temporal_dropout_p=temporal_dropout_p)
model.load_state_dict(checkpoint["model_state_dict"], strict=True)
model.eval()
return model, {i + 1: char for i, char in enumerate(chars)}
@lru_cache(maxsize=1)
def load_language_model():
# Trusted release artifact containing only fitted character n-gram counts.
with LANGUAGE_MODEL.open("rb") as handle:
return pickle.load(handle)
def preprocess(image: Image.Image) -> tuple[torch.Tensor, int, int]:
if image is None:
raise ValueError("Upload a cropped image containing one Kashmiri text line.")
normal = image.convert("L")
scale = IMAGE_HEIGHT / max(1, normal.height)
resized_width = max(1, min(IMAGE_WIDTH, int(normal.width * scale)))
normal = normal.resize((resized_width, IMAGE_HEIGHT), Image.Resampling.BICUBIC)
canvas = Image.new("L", (IMAGE_WIDTH, IMAGE_HEIGHT), 255)
canvas.paste(normal, (0, 0))
pixels = 1.0 - np.asarray(canvas, dtype=np.float32) / 255.0
input_length = max(
1,
min(IMAGE_WIDTH // HORIZONTAL_STRIDE, int(math.ceil(resized_width / HORIZONTAL_STRIDE))),
)
return torch.from_numpy(pixels).unsqueeze(0).unsqueeze(0), input_length, resized_width
def ctc_decode(ids: list[int], itos: dict[int, str]) -> str:
output: list[str] = []
previous = None
for token in ids:
if token != 0 and token != previous:
output.append(itos.get(token, ""))
previous = token
return "".join(output)[::-1]
def recognize_line(image: Image.Image, *, decoder: str = "beam") -> tuple[str, dict[str, Any]]:
if decoder not in {"beam", "greedy"}:
raise ValueError("decoder must be 'beam' or 'greedy'")
model, itos = load_model()
tensor, input_length, resized_width = preprocess(image)
with torch.inference_mode():
logits = model(tensor)
matrix = logits[:input_length, 0]
diagnostics: dict[str, Any] = {
"input_width": image.width,
"input_height": image.height,
"resized_width": resized_width,
"ctc_frames": input_length,
"decoder": decoder,
}
if decoder == "greedy":
text = ctc_decode(matrix.argmax(-1).tolist(), itos)
else:
from ctc_prefix_beam import prefix_beam_search
ids = prefix_beam_search(
matrix.cpu().numpy(),
lm=load_language_model(),
**BEAM_CONFIG,
)
# Prefix-beam output is already CTC-collapsed in increasing-x order.
text = "".join(itos[token] for token in ids)[::-1]
diagnostics.update(BEAM_CONFIG)
return text, diagnostics
|