File size: 2,892 Bytes
4b917c4 | 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 | """32M checkpoint consumer for RegionalPipeline's RTL bank; raw CTC retained."""
import json,math
from pathlib import Path
import numpy as np
import torch
from recognizer import Recognizer
from geometry import crop_line64,pad_batch
from visual_token import RTLVisualChar
class RTLReader:
def __init__(self, checkpoint, vocabulary, device='cpu', frame_budget=8192):
self.device = torch.device(device)
data = json.loads(Path(vocabulary).read_text())
self.alphabet = data['chars']
self.model = Recognizer(classes=len(self.alphabet) + 1, gradient_checkpointing=False)
state = torch.load(checkpoint, map_location='cpu', weights_only=True)
self.model.load_state_dict(state['model'])
self.model.to(self.device).eval()
self.frame_budget = frame_budget
@torch.inference_mode()
def __call__(self, gray, boxes):
from indra_ocr.pipeline import MIN_CONF
crops, receipts = [], []
for box in boxes:
crop, receipt = crop_line64(gray, box)
crops.append(crop)
receipts.append(receipt)
order = sorted(range(len(crops)), key=lambda i: crops[i].shape[1])
parts = []
for i in order:
frames = math.ceil(crops[i].shape[1] / 8) * 2
if parts and len(parts[-1]) < 32 and (len(parts[-1]) + 1) * frames <= self.frame_budget:
parts[-1].append(i)
else:
parts.append([i])
readings = [None] * len(crops)
for ids in parts:
x, widths = pad_batch([crops[i] for i in ids])
x = torch.from_numpy(x).to(self.device)
logits, lengths = self.model(x, torch.from_numpy(widths))
probs, best = logits.float().softmax(-1).max(-1)
best, probs = best.cpu().numpy(), probs.cpu().numpy()
for k, i in enumerate(ids):
n = int(lengths[k])
seq, pm = best[k, :n], probs[k, :n]
keep = (seq != 0) & np.concatenate(([True], seq[1:] != seq[:-1]))
steps = np.nonzero(keep)[0]
token_ids = seq[steps].tolist()
visual = ''.join(self.alphabet[t - 1] for t in token_ids)
conf = float(pm[keep].mean()) if keep.any() else 0.
g = receipts[i]
left, _, right, _ = g['padded_source_extent']
xs = left + (steps * 4 + 3.5) * (right - left) / crops[i].shape[1]
toks = list(zip([RTLVisualChar(self.alphabet[t - 1]) for t in token_ids], xs.tolist()))
receipts[i] = dict(g, raw_visual_text=visual, raw_token_ids=token_ids, confidence=conf, frame_centers=xs.tolist(), confidence_threshold=MIN_CONF, confidence_filtered=conf < MIN_CONF)
readings[i] = (visual, conf, toks) if conf >= MIN_CONF else ('', conf, [])
return readings, receipts
|