| """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 |
|
|
