| --- |
| language: ml |
| tags: |
| - ocr |
| - scene-text |
| - malayalam |
| - image-to-text |
| - pytorch |
| license: mit |
| --- |
| |
| # Malayalam Scene Text OCR — PARSeq-style |
|
|
| A PARSeq-inspired OCR model (ViT encoder + Transformer decoder) for Malayalam scene text recognition. |
|
|
| ## Available Checkpoints |
|
|
| | Checkpoint | Description | Word Acc (real val) | |
| |---|---|---| |
| | `parseq_best.pth` | Pretrained on 950k synthetic images | 97.64% (synthetic) | |
| | `parseq_finetuned_best.pth` | Finetuned v1 (100 epochs, basic) | 84.15% | |
| | `parseq_finetuned_v2.pth` | **Finetuned v2 (150 epochs, label smoothing + cosine restart)** | **91.46%** | |
|
|
| **Use `parseq_finetuned_v2.pth` for best results on real Malayalam scene text.** |
|
|
| --- |
|
|
| ## Benchmark Results |
|
|
| Evaluated on 82 real Malayalam scene text images (val_ split, IndicVignesh dataset). |
| All scores normalized for Malayalam Unicode variants (Chillu characters, ZWJ/ZWNJ). |
|
|
| | Model | Word Acc | Char Acc | |
| |---|---|---| |
| | GPT-5.4 | 34.15% | 63.17% | |
| | Claude Sonnet 4.5 | 35.37% | 56.18% | |
| | Claude Sonnet 4.6 | 84.15% | 93.57% | |
| | Gemini 3 Flash Preview | 85.37% | 94.76% | |
| | Claude Opus 4.6 | 86.59% | 95.41% | |
| | **PARSeq v2 (Ours)** | **91.46%** | **97.18%** | |
|
|
| Our 25M parameter specialized model **beats all frontier VLMs** including Claude Opus 4.6, runs locally at zero inference cost. |
|
|
| --- |
|
|
| ## Model Architecture |
|
|
| - **Encoder**: Vision Transformer (ViT) — patch size 4×8 on 32×128 images → 128 patches |
| - **Decoder**: Autoregressive Transformer decoder with causal masking |
| - **Parameters**: ~25M |
| - **Vocab**: 99 tokens (95 Malayalam characters + `[PAD]`, `[BOS]`, `[EOS]`, `[UNK]`) |
| - **Input**: 32×128 RGB images (auto-resized) |
|
|
| --- |
|
|
| ## Quick Start |
|
|
| ### Install |
|
|
| ```bash |
| pip install torch torchvision huggingface_hub Pillow |
| ``` |
|
|
| ### Define model class |
|
|
| ```python |
| import math |
| import torch |
| import torch.nn as nn |
| |
| class PatchEmbed(nn.Module): |
| def __init__(self, img_h=32, img_w=128, patch_h=4, patch_w=8, in_chans=3, embed_dim=384): |
| super().__init__() |
| self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=(patch_h, patch_w), stride=(patch_h, patch_w)) |
| self.norm = nn.LayerNorm(embed_dim) |
| def forward(self, x): |
| return self.norm(self.proj(x).flatten(2).transpose(1, 2)) |
| |
| class SinusoidalPE(nn.Module): |
| def __init__(self, embed_dim, max_len=512, dropout=0.1): |
| super().__init__() |
| self.dropout = nn.Dropout(dropout) |
| pe = torch.zeros(max_len, embed_dim) |
| pos = torch.arange(0, max_len).unsqueeze(1).float() |
| div = torch.exp(torch.arange(0, embed_dim, 2).float() * (-math.log(10000.0) / embed_dim)) |
| pe[:, 0::2] = torch.sin(pos * div) |
| pe[:, 1::2] = torch.cos(pos * div) |
| self.register_buffer('pe', pe.unsqueeze(0)) |
| def forward(self, x): |
| return self.dropout(x + self.pe[:, :x.size(1)]) |
| |
| class ViTEncoder(nn.Module): |
| def __init__(self, img_h=32, img_w=128, patch_h=4, patch_w=8, in_chans=3, |
| embed_dim=384, depth=6, num_heads=6, mlp_ratio=4.0, dropout=0.1): |
| super().__init__() |
| self.patch_embed = PatchEmbed(img_h, img_w, patch_h, patch_w, in_chans, embed_dim) |
| self.pos_enc = SinusoidalPE(embed_dim, max_len=512, dropout=dropout) |
| encoder_layer = nn.TransformerEncoderLayer(d_model=embed_dim, nhead=num_heads, |
| dim_feedforward=int(embed_dim*mlp_ratio), dropout=dropout, |
| activation='gelu', batch_first=True, norm_first=True) |
| self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=depth) |
| self.norm = nn.LayerNorm(embed_dim) |
| def forward(self, x): |
| return self.norm(self.encoder(self.pos_enc(self.patch_embed(x)))) |
| |
| class TransformerDecoder(nn.Module): |
| def __init__(self, vocab_size, embed_dim=384, depth=6, num_heads=6, |
| mlp_ratio=4.0, dropout=0.1, max_label_len=26, pad_idx=0): |
| super().__init__() |
| self.pad_idx = pad_idx |
| self.token_embed = nn.Embedding(vocab_size, embed_dim, padding_idx=pad_idx) |
| self.pos_enc = SinusoidalPE(embed_dim, max_len=max_label_len+2, dropout=dropout) |
| decoder_layer = nn.TransformerDecoderLayer(d_model=embed_dim, nhead=num_heads, |
| dim_feedforward=int(embed_dim*mlp_ratio), dropout=dropout, |
| activation='gelu', batch_first=True, norm_first=True) |
| self.decoder = nn.TransformerDecoder(decoder_layer, num_layers=depth) |
| self.norm = nn.LayerNorm(embed_dim) |
| self.head = nn.Linear(embed_dim, vocab_size) |
| def _causal_mask(self, sz, device): |
| return torch.triu(torch.ones(sz, sz, device=device), diagonal=1).bool() |
| def forward(self, tgt_inp, memory): |
| B, T = tgt_inp.shape |
| x = self.pos_enc(self.token_embed(tgt_inp)) |
| x = self.decoder(tgt=x, memory=memory, |
| tgt_mask=self._causal_mask(T, tgt_inp.device), |
| tgt_key_padding_mask=(tgt_inp == self.pad_idx)) |
| return self.head(self.norm(x)) |
| |
| class PARSeqOCR(nn.Module): |
| def __init__(self, vocab_size, img_h=32, img_w=128, patch_h=4, patch_w=8, |
| embed_dim=384, enc_depth=6, dec_depth=6, num_heads=6, |
| mlp_ratio=4.0, dropout=0.1, max_label_len=25, pad_idx=0): |
| super().__init__() |
| self.max_label_len = max_label_len |
| self.pad_idx = pad_idx |
| self.encoder = ViTEncoder(img_h, img_w, patch_h, patch_w, 3, embed_dim, |
| enc_depth, num_heads, mlp_ratio, dropout) |
| self.decoder = TransformerDecoder(vocab_size, embed_dim, dec_depth, num_heads, |
| mlp_ratio, dropout, max_label_len, pad_idx) |
| |
| def forward(self, images, tgt_inp): |
| return self.decoder(tgt_inp, self.encoder(images)) |
| |
| @torch.no_grad() |
| def greedy_decode(self, images, bos_idx, eos_idx, max_len=None): |
| """Fast greedy decoding — good for batches.""" |
| self.eval() |
| max_len = max_len or self.max_label_len |
| B, device = images.size(0), images.device |
| memory = self.encoder(images) |
| generated = torch.full((B, 1), bos_idx, dtype=torch.long, device=device) |
| finished = torch.zeros(B, dtype=torch.bool, device=device) |
| for _ in range(max_len): |
| next_token = self.decoder(generated, memory)[:, -1, :].argmax(-1) |
| next_token = torch.where(finished, torch.full_like(next_token, self.pad_idx), next_token) |
| generated = torch.cat([generated, next_token.unsqueeze(1)], dim=1) |
| finished = finished | (next_token == eos_idx) |
| if finished.all(): break |
| preds = [] |
| for seq in generated.tolist(): |
| seq = seq[1:] |
| if eos_idx in seq: seq = seq[:seq.index(eos_idx)] |
| preds.append(seq) |
| return preds |
| |
| @torch.no_grad() |
| def beam_decode(self, images, bos_idx, eos_idx, beam_size=5, max_len=None): |
| """Beam search decoding — slightly more accurate, slower.""" |
| self.eval() |
| max_len = max_len or self.max_label_len |
| device = images.device |
| B = images.size(0) |
| memory = self.encoder(images) |
| all_preds = [] |
| for b in range(B): |
| mem = memory[b:b+1] |
| beams = [(0.0, [bos_idx])] |
| completed = [] |
| for _ in range(max_len): |
| new_beams = [] |
| for score, tokens in beams: |
| if tokens[-1] == eos_idx: |
| completed.append((score, tokens)) |
| continue |
| seq = torch.tensor([tokens], dtype=torch.long, device=device) |
| logits = self.decoder(seq, mem) |
| log_prob = torch.log_softmax(logits[0, -1, :], dim=-1) |
| topk_scores, topk_ids = log_prob.topk(beam_size) |
| for s, t in zip(topk_scores.tolist(), topk_ids.tolist()): |
| new_beams.append((score + s, tokens + [t])) |
| new_beams.sort(key=lambda x: x[0], reverse=True) |
| beams = new_beams[:beam_size] |
| if all(t[-1] == eos_idx for _, t in beams): |
| completed.extend(beams) |
| break |
| completed.extend(beams) |
| completed.sort(key=lambda x: x[0] / max(len(x[1]), 1), reverse=True) |
| best = completed[0][1][1:] |
| if eos_idx in best: best = best[:best.index(eos_idx)] |
| all_preds.append(best) |
| return all_preds |
| ``` |
|
|
| ### Load model and run inference |
|
|
| ```python |
| import json, torch |
| import torchvision.transforms as T |
| from PIL import Image |
| from huggingface_hub import hf_hub_download |
| |
| REPO = 'magles/malayalam-ocr-parseq' |
| DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu' |
| |
| # Download files |
| ckpt_path = hf_hub_download(REPO, 'parseq_finetuned_v2.pth') |
| char2idx_path = hf_hub_download(REPO, 'char2idx.json') |
| idx2char_path = hf_hub_download(REPO, 'idx2char.json') |
| |
| # Load vocab |
| with open(char2idx_path, encoding='utf-8') as f: |
| char2idx = json.load(f) |
| with open(idx2char_path, encoding='utf-8') as f: |
| idx2char = {int(k): v for k, v in json.load(f).items()} |
| |
| # Load model |
| ckpt = torch.load(ckpt_path, map_location=DEVICE, weights_only=False) |
| model = PARSeqOCR(vocab_size=len(char2idx)) |
| model.load_state_dict(ckpt['model']) |
| model = model.to(DEVICE) |
| model.eval() |
| print(f"Model loaded — vocab: {len(char2idx)}") |
| |
| # Preprocess |
| transform = T.Compose([ |
| T.Resize((32, 128)), |
| T.ToTensor(), |
| T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), |
| ]) |
| |
| bos_idx = char2idx['[BOS]'] |
| eos_idx = char2idx['[EOS]'] |
| |
| def predict(image_path, use_beam=True, beam_size=5): |
| img = Image.open(image_path).convert('RGB') |
| tensor = transform(img).unsqueeze(0).to(DEVICE) |
| if use_beam: |
| indices = model.beam_decode(tensor, bos_idx, eos_idx, beam_size=beam_size)[0] |
| else: |
| indices = model.greedy_decode(tensor, bos_idx, eos_idx)[0] |
| return ''.join(idx2char.get(i, '') for i in indices) |
| |
| # Single image |
| print(predict('your_image.jpg')) |
| |
| # Batch (greedy is faster for batches) |
| def predict_batch(image_paths): |
| imgs = torch.stack([transform(Image.open(p).convert('RGB')) |
| for p in image_paths]).to(DEVICE) |
| all_seqs = model.greedy_decode(imgs, bos_idx, eos_idx) |
| return [''.join(idx2char.get(i, '') for i in seq) for seq in all_seqs] |
| ``` |
|
|
| ### Greedy vs Beam Search |
|
|
| ```python |
| # Greedy — fast, good for batches, ~84% word accuracy on v1 / 91% on v2 |
| indices = model.greedy_decode(tensor, bos_idx, eos_idx)[0] |
| |
| # Beam search — slightly more accurate, slower (processes one image at a time internally) |
| # beam_size=5 is the sweet spot — larger values give no further improvement |
| indices = model.beam_decode(tensor, bos_idx, eos_idx, beam_size=5)[0] |
| ``` |
|
|
| ### Malayalam Unicode Normalization |
|
|
| Some ground truth labels use different Unicode encodings for the same visual character |
| (e.g. Chillu characters: `ൾ` = `U+0D7E` vs `ള്` = `U+0D33 + U+0D4D + U+200D`). |
| Normalize before comparing predictions: |
|
|
| ```python |
| CHILLU_MAP = { |
| '\u0d7a': '\u0d23\u0d4d', |
| '\u0d7b': '\u0d28\u0d4d', |
| '\u0d7c': '\u0d30\u0d4d', |
| '\u0d7d': '\u0d32\u0d4d', |
| '\u0d7e': '\u0d33\u0d4d', |
| '\u0d7f': '\u0d15\u0d4d', |
| } |
| |
| def normalize_malayalam(text): |
| text = text.strip().replace('\u200c', '').replace('\u200d', '') |
| for chillu, base in CHILLU_MAP.items(): |
| text = text.replace(chillu, base) |
| return text |
| |
| # Compare |
| normalize_malayalam(pred) == normalize_malayalam(gt) |
| ``` |
|
|
| --- |
|
|
| ## Training Details |
|
|
| ### Pretraining (`parseq_best.pth`) |
| - **Dataset**: 950,000 synthetic Malayalam scene text images |
| - **Epochs**: 20 | **Batch**: 512 | **LR**: 1e-3 (OneCycleLR) | **AMP**: fp16 |
| - **Hardware**: RTX 5090 | **Time**: ~3 hours |
| |
| ### Finetuning v1 (`parseq_finetuned_best.pth`) |
| - **Dataset**: 915 real images (IndicVignesh `finetune_` split) |
| - **Epochs**: 100 | **Batch**: 32 | **LR**: 1e-4 |
| - **Val**: 82 real images (`val_` split) |
|
|
| ### Finetuning v2 (`parseq_finetuned_v2.pth`) ← recommended |
| - **Dataset**: 915 real images (same as v1) |
| - **Epochs**: 150 | **Batch**: 32 | **LR**: 1e-4 |
| - **Label smoothing**: 0.1 — prevents overconfidence, improves generalization |
| - **Scheduler**: CosineAnnealingWarmRestarts (T_0=50) — 3 LR restart cycles |
| - **Encoder frozen** for first 10 epochs, then unfrozen |
| - **Best epoch**: 121 |
| |
| --- |
| |
| ## Citation |
| |
| ```bibtex |
| @inproceedings{bautista2022parseq, |
| title={Scene Text Recognition with Permuted Autoregressive Sequence Models}, |
| author={Bautista, Darwin and Atienza, Rowel}, |
| booktitle={European Conference on Computer Vision (ECCV)}, |
| year={2022} |
| } |
| ``` |
| |
