Views
No views yet
| 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% |
parseq_finetuned_v2.pth for best results on real Malayalam scene text.| 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% |
[PAD], [BOS], [EOS], [UNK])pip install torch torchvision huggingface_hub Pillow1import math
2import torch
3import torch.nn as nn
4
5class PatchEmbed(nn.Module):
6 def __init__(self, img_h=32, img_w=128, patch_h=4, patch_w=8, in_chans=3, embed_dim=384):
7 super().__init__()
8 self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=(patch_h, patch_w), stride=(patch_h, patch_w))
9 self.norm = nn.LayerNorm(embed_dim)
10 def forward(self, x):
11 return self.norm(self.proj(x).flatten(2).transpose(1, 2))
12
13class SinusoidalPE(nn.Module):
14 def __init__(self, embed_dim, max_len=512, dropout=0.1):
15 super().__init__()
16 self.dropout = nn.Dropout(dropout)
17 pe = torch.zeros(max_len, embed_dim)
18 pos = torch.arange(0, max_len).unsqueeze(1).float()
19 div = torch.exp(torch.arange(0, embed_dim, 2).float() * (-math.log(10000.0) / embed_dim))
20 pe[:, 0::2] = torch.sin(pos * div)
21 pe[:, 1::2] = torch.cos(pos * div)
22 self.register_buffer('pe', pe.unsqueeze(0))
23 def forward(self, x):
24 return self.dropout(x + self.pe[:, :x.size(1)])
25
26class ViTEncoder(nn.Module):
27 def __init__(self, img_h=32, img_w=128, patch_h=4, patch_w=8, in_chans=3,
28 embed_dim=384, depth=6, num_heads=6, mlp_ratio=4.0, dropout=0.1):
29 super().__init__()
30 self.patch_embed = PatchEmbed(img_h, img_w, patch_h, patch_w, in_chans, embed_dim)
31 self.pos_enc = SinusoidalPE(embed_dim, max_len=512, dropout=dropout)
32 encoder_layer = nn.TransformerEncoderLayer(d_model=embed_dim, nhead=num_heads,
33 dim_feedforward=int(embed_dim*mlp_ratio), dropout=dropout,
34 activation='gelu', batch_first=True, norm_first=True)
35 self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=depth)
36 self.norm = nn.LayerNorm(embed_dim)
37 def forward(self, x):
38 return self.norm(self.encoder(self.pos_enc(self.patch_embed(x))))
39
40class TransformerDecoder(nn.Module):
41 def __init__(self, vocab_size, embed_dim=384, depth=6, num_heads=6,
42 mlp_ratio=4.0, dropout=0.1, max_label_len=26, pad_idx=0):
43 super().__init__()
44 self.pad_idx = pad_idx
45 self.token_embed = nn.Embedding(vocab_size, embed_dim, padding_idx=pad_idx)
46 self.pos_enc = SinusoidalPE(embed_dim, max_len=max_label_len+2, dropout=dropout)
47 decoder_layer = nn.TransformerDecoderLayer(d_model=embed_dim, nhead=num_heads,
48 dim_feedforward=int(embed_dim*mlp_ratio), dropout=dropout,
49 activation='gelu', batch_first=True, norm_first=True)
50 self.decoder = nn.TransformerDecoder(decoder_layer, num_layers=depth)
51 self.norm = nn.LayerNorm(embed_dim)
52 self.head = nn.Linear(embed_dim, vocab_size)
53 def _causal_mask(self, sz, device):
54 return torch.triu(torch.ones(sz, sz, device=device), diagonal=1).bool()
55 def forward(self, tgt_inp, memory):
56 B, T = tgt_inp.shape
57 x = self.pos_enc(self.token_embed(tgt_inp))
58 x = self.decoder(tgt=x, memory=memory,
59 tgt_mask=self._causal_mask(T, tgt_inp.device),
60 tgt_key_padding_mask=(tgt_inp == self.pad_idx))
61 return self.head(self.norm(x))
62
63class PARSeqOCR(nn.Module):
64 def __init__(self, vocab_size, img_h=32, img_w=128, patch_h=4, patch_w=8,
65 embed_dim=384, enc_depth=6, dec_depth=6, num_heads=6,
66 mlp_ratio=4.0, dropout=0.1, max_label_len=25, pad_idx=0):
67 super().__init__()
68 self.max_label_len = max_label_len
69 self.pad_idx = pad_idx
70 self.encoder = ViTEncoder(img_h, img_w, patch_h, patch_w, 3, embed_dim,
71 enc_depth, num_heads, mlp_ratio, dropout)
72 self.decoder = TransformerDecoder(vocab_size, embed_dim, dec_depth, num_heads,
73 mlp_ratio, dropout, max_label_len, pad_idx)
74
75 def forward(self, images, tgt_inp):
76 return self.decoder(tgt_inp, self.encoder(images))
77
78 @torch.no_grad()
79 def greedy_decode(self, images, bos_idx, eos_idx, max_len=None):
80 """Fast greedy decoding — good for batches."""
81 self.eval()
82 max_len = max_len or self.max_label_len
83 B, device = images.size(0), images.device
84 memory = self.encoder(images)
85 generated = torch.full((B, 1), bos_idx, dtype=torch.long, device=device)
86 finished = torch.zeros(B, dtype=torch.bool, device=device)
87 for _ in range(max_len):
88 next_token = self.decoder(generated, memory)[:, -1, :].argmax(-1)
89 next_token = torch.where(finished, torch.full_like(next_token, self.pad_idx), next_token)
90 generated = torch.cat([generated, next_token.unsqueeze(1)], dim=1)
91 finished = finished | (next_token == eos_idx)
92 if finished.all(): break
93 preds = []
94 for seq in generated.tolist():
95 seq = seq[1:]
96 if eos_idx in seq: seq = seq[:seq.index(eos_idx)]
97 preds.append(seq)
98 return preds
99
100 @torch.no_grad()
101 def beam_decode(self, images, bos_idx, eos_idx, beam_size=5, max_len=None):
102 """Beam search decoding — slightly more accurate, slower."""
103 self.eval()
104 max_len = max_len or self.max_label_len
105 device = images.device
106 B = images.size(0)
107 memory = self.encoder(images)
108 all_preds = []
109 for b in range(B):
110 mem = memory[b:b+1]
111 beams = [(0.0, [bos_idx])]
112 completed = []
113 for _ in range(max_len):
114 new_beams = []
115 for score, tokens in beams:
116 if tokens[-1] == eos_idx:
117 completed.append((score, tokens))
118 continue
119 seq = torch.tensor([tokens], dtype=torch.long, device=device)
120 logits = self.decoder(seq, mem)
121 log_prob = torch.log_softmax(logits[0, -1, :], dim=-1)
122 topk_scores, topk_ids = log_prob.topk(beam_size)
123 for s, t in zip(topk_scores.tolist(), topk_ids.tolist()):
124 new_beams.append((score + s, tokens + [t]))
125 new_beams.sort(key=lambda x: x[0], reverse=True)
126 beams = new_beams[:beam_size]
127 if all(t[-1] == eos_idx for _, t in beams):
128 completed.extend(beams)
129 break
130 completed.extend(beams)
131 completed.sort(key=lambda x: x[0] / max(len(x[1]), 1), reverse=True)
132 best = completed[0][1][1:]
133 if eos_idx in best: best = best[:best.index(eos_idx)]
134 all_preds.append(best)
135 return all_preds1import json, torch
2import torchvision.transforms as T
3from PIL import Image
4from huggingface_hub import hf_hub_download
5
6REPO = 'magles/malayalam-ocr-parseq'
7DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
8
9# Download files
10ckpt_path = hf_hub_download(REPO, 'parseq_finetuned_v2.pth')
11char2idx_path = hf_hub_download(REPO, 'char2idx.json')
12idx2char_path = hf_hub_download(REPO, 'idx2char.json')
13
14# Load vocab
15with open(char2idx_path, encoding='utf-8') as f:
16 char2idx = json.load(f)
17with open(idx2char_path, encoding='utf-8') as f:
18 idx2char = {int(k): v for k, v in json.load(f).items()}
19
20# Load model
21ckpt = torch.load(ckpt_path, map_location=DEVICE, weights_only=False)
22model = PARSeqOCR(vocab_size=len(char2idx))
23model.load_state_dict(ckpt['model'])
24model = model.to(DEVICE)
25model.eval()
26print(f"Model loaded — vocab: {len(char2idx)}")
27
28# Preprocess
29transform = T.Compose([
30 T.Resize((32, 128)),
31 T.ToTensor(),
32 T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
33])
34
35bos_idx = char2idx['[BOS]']
36eos_idx = char2idx['[EOS]']
37
38def predict(image_path, use_beam=True, beam_size=5):
39 img = Image.open(image_path).convert('RGB')
40 tensor = transform(img).unsqueeze(0).to(DEVICE)
41 if use_beam:
42 indices = model.beam_decode(tensor, bos_idx, eos_idx, beam_size=beam_size)[0]
43 else:
44 indices = model.greedy_decode(tensor, bos_idx, eos_idx)[0]
45 return ''.join(idx2char.get(i, '') for i in indices)
46
47# Single image
48print(predict('your_image.jpg'))
49
50# Batch (greedy is faster for batches)
51def predict_batch(image_paths):
52 imgs = torch.stack([transform(Image.open(p).convert('RGB'))
53 for p in image_paths]).to(DEVICE)
54 all_seqs = model.greedy_decode(imgs, bos_idx, eos_idx)
55 return [''.join(idx2char.get(i, '') for i in seq) for seq in all_seqs]1# Greedy — fast, good for batches, ~84% word accuracy on v1 / 91% on v2
2indices = model.greedy_decode(tensor, bos_idx, eos_idx)[0]
3
4# Beam search — slightly more accurate, slower (processes one image at a time internally)
5# beam_size=5 is the sweet spot — larger values give no further improvement
6indices = model.beam_decode(tensor, bos_idx, eos_idx, beam_size=5)[0]ൾ = U+0D7E vs ള് = U+0D33 + U+0D4D + U+200D).
Normalize before comparing predictions:1CHILLU_MAP = {
2 '\u0d7a': '\u0d23\u0d4d',
3 '\u0d7b': '\u0d28\u0d4d',
4 '\u0d7c': '\u0d30\u0d4d',
5 '\u0d7d': '\u0d32\u0d4d',
6 '\u0d7e': '\u0d33\u0d4d',
7 '\u0d7f': '\u0d15\u0d4d',
8}
9
10def normalize_malayalam(text):
11 text = text.strip().replace('\u200c', '').replace('\u200d', '')
12 for chillu, base in CHILLU_MAP.items():
13 text = text.replace(chillu, base)
14 return text
15
16# Compare
17normalize_malayalam(pred) == normalize_malayalam(gt)parseq_best.pth)parseq_finetuned_best.pth)finetune_ split)val_ split)parseq_finetuned_v2.pth) ← recommended1@inproceedings{bautista2022parseq,
2 title={Scene Text Recognition with Permuted Autoregressive Sequence Models},
3 author={Bautista, Darwin and Atienza, Rowel},
4 booktitle={European Conference on Computer Vision (ECCV)},
5 year={2022}
6}