File size: 6,708 Bytes
e3f99fb | 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 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 | """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]
|