"""NetraOCR processor: Qwen2VL image processing + the NetraOCR tokenizer + dual-vision placeholder expansion. Each image placeholder `<|image_pad|>` in the chat-templated prompt is expanded to `frames * (t*h*w / merge²)` tokens (frames = 2 in dual-vision mode). `mm_token_type_ids` marks them, which the model needs for its 3D rotary positions. """ from __future__ import annotations import json import os import torch from transformers import AutoImageProcessor, PreTrainedTokenizerFast class NetraProcessor: def __init__(self, tokenizer, image_processor, frames: int = 2): self.tokenizer = tokenizer self.image_processor = image_processor self.frames = frames self.image_token_id = tokenizer.convert_tokens_to_ids("<|image_pad|>") @classmethod def from_pretrained(cls, path: str, frames: int = 2) -> "NetraProcessor": # Load tokenizer.json verbatim. AutoTokenizer may rebuild a model-specific class and drop or # alter the appended Indic tokens (the same failure mode found with the Sarvam tokenizer). cfg = json.load(open(os.path.join(path, "tokenizer_config.json"), encoding="utf-8")) tok = PreTrainedTokenizerFast( tokenizer_file=os.path.join(path, "tokenizer.json"), eos_token=cfg.get("eos_token", "<|im_end|>"), pad_token=cfg.get("pad_token", "<|endoftext|>"), ) tok.padding_side = "left" tmpl = os.path.join(path, "chat_template.jinja") if os.path.exists(tmpl): tok.chat_template = open(tmpl, encoding="utf-8").read() return cls(tok, AutoImageProcessor.from_pretrained(path), frames) def prompt(self, instruction: str, n_images: int = 1) -> str: content = [{"type": "image"}] * n_images + [{"type": "text", "text": instruction}] return self.tokenizer.apply_chat_template([{"role": "user", "content": content}], add_generation_prompt=True, tokenize=False, enable_thinking=False) def __call__(self, texts: list[str], images: list, images_kwargs: dict | None = None, targets: list[str] | None = None) -> dict: img = self.image_processor(images=images, return_tensors="pt", **(images_kwargs or {})) grid = img["image_grid_thw"] merge = self.image_processor.merge_size counts = (grid.prod(-1) // merge**2 * self.frames).tolist() rows, label_rows, k = [], [], 0 for i, text in enumerate(texts): ids = self.tokenizer(text, add_special_tokens=False)["input_ids"] expanded = [] for t in ids: if t == self.image_token_id: expanded.extend([self.image_token_id] * counts[k]) k += 1 else: expanded.append(t) labels = [-100] * len(expanded) if targets is not None: tgt = self.tokenizer(targets[i], add_special_tokens=False)["input_ids"] + [self.tokenizer.eos_token_id] expanded = expanded + tgt labels = labels + tgt rows.append(expanded) label_rows.append(labels) if k != len(counts): raise ValueError(f"{len(counts)} images but {k} <|image_pad|> placeholders in the prompts") width = max(len(r) for r in rows) pad = self.tokenizer.pad_token_id input_ids = torch.full((len(rows), width), pad, dtype=torch.long) attention = torch.zeros((len(rows), width), dtype=torch.long) labels = torch.full((len(rows), width), -100, dtype=torch.long) for i, (r, lab) in enumerate(zip(rows, label_rows)): # left padding input_ids[i, width - len(r):] = torch.tensor(r) attention[i, width - len(r):] = 1 labels[i, width - len(r):] = torch.tensor(lab) out = { "input_ids": input_ids, "attention_mask": attention, "mm_token_type_ids": (input_ids == self.image_token_id).int() * attention.int(), "pixel_values": img["pixel_values"], "image_grid_thw": grid, } if targets is not None: out["labels"] = labels return out def decode(self, ids) -> str: return self.tokenizer.decode(ids, skip_special_tokens=True) def split_page_regions(img, max_regions: int = 10, min_band_px: int = 140) -> list: """Split a dense full-page document/ledger image into horizontal text-band crops using ink projection valleys. Each band is Tight-cropped horizontally around ink and ready for `area_clamp`, giving small handwritten or multi-column ledger lines 4-8x more pixels per character than single-pass whole-page downsampling. """ import numpy as np from PIL import Image rgb = img.convert("RGB") w, h = rgb.size if h < min_band_px * 2: return [rgb] gray = np.asarray(rgb.convert("L"), dtype=np.float32) bg = float(np.percentile(gray, 85)) ink = (bg - gray) > max(18.0, 0.12 * bg) row_ink = ink.mean(axis=1) # Smooth row ink profile with a 9-px box filter to locate inter-line valleys kernel = np.ones(9, dtype=np.float32) / 9.0 smooth = np.convolve(row_ink, kernel, mode="same") active = np.where(smooth > 0.008)[0] if len(active) < 20: return [rgb] y_top, y_bot = max(0, int(active[0]) - 12), min(h, int(active[-1]) + 12) span = y_bot - y_top if span < min_band_px * 2: return [rgb.crop((0, y_top, w, y_bot))] target_bands = min(max_regions, max(2, int(round(span / max(min_band_px, span / max_regions))))) step = span / target_bands cuts = [y_top] for b_idx in range(1, target_bands): nominal = int(y_top + b_idx * step) win_lo = max(cuts[-1] + min_band_px // 2, nominal - int(step * 0.35)) win_hi = min(y_bot - min_band_px // 2, nominal + int(step * 0.35)) if win_hi > win_lo: valley = win_lo + int(np.argmin(smooth[win_lo:win_hi])) else: valley = nominal cuts.append(valley) cuts.append(y_bot) crops = [] for y0, y1 in zip(cuts[:-1], cuts[1:]): pad0, pad1 = max(0, y0 - 8), min(h, y1 + 8) band_ink = ink[pad0:pad1] if band_ink.mean() < 0.004: continue col_active = np.where(band_ink.mean(axis=0) > 0.003)[0] if len(col_active) >= 10: x0 = max(0, int(col_active[0]) - 12) x1 = min(w, int(col_active[-1]) + 12) else: x0, x1 = 0, w if (pad1 - pad0) >= 24 and (x1 - x0) >= 48: crops.append(rgb.crop((x0, pad0, x1, pad1))) return crops if crops else [rgb]