| """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": |
| |
| |
| 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)): |
| 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) |
|
|
| |
| 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] |
|
|
|
|
