#!/usr/bin/env python3
# Copyright 2026 Muse Mesh Private Limited. Licensed under the Apache License, Version 2.0 (LICENSE-CODE).
"""Sansar OCR 2B: transcribe printed Sanskrit (Devanagari) pages with transformers.

The page is prepared exactly as in training: near-white outer margins are trimmed, then the image is resized
(aspect kept) so that width x height <= 1.8 MP and both sides are multiples of 32. One fixed prompt, thinking off,
greedy decoding.

    pip install "transformers>=5.18" torch torchvision pillow numpy accelerate
    python ocr_page.py page1.jpg page2.png ...            # prints each page's text
    python ocr_page.py --model MuseMesh/sansar-ocr-2b --batch 4 --out_dir out/ pages/*.jpg
"""
import argparse
import math
import sys
from pathlib import Path

REPO = "MuseMesh/sansar-ocr-2b"
PROMPT = ("Transcribe all the text printed on this page exactly as printed, in reading order, line by line. "
          "Output only the text.")
MAX_PIXELS = 1_800_000
MAX_NEW_TOKENS = 3072


def trim_margins(im, pad=24, thr_delta=60):
    """Crop near-white outer margins (a dark scanner border counts as content and stays)."""
    import numpy as np
    g = im.convert("L")
    small = g.resize((max(1, g.width // 4), max(1, g.height // 4)))
    a = np.asarray(small).astype(int)
    paper = np.percentile(a, 90)
    ink = a < paper - thr_delta
    rows = np.where(ink.mean(1) > 0.003)[0]
    cols = np.where(ink.mean(0) > 0.003)[0]
    if len(rows) < 2 or len(cols) < 2:
        return im
    y0, y1, x0, x1 = rows[0] * 4, (rows[-1] + 1) * 4, cols[0] * 4, (cols[-1] + 1) * 4
    if (x1 - x0) * (y1 - y0) < 0.15 * im.width * im.height:
        return im
    return im.crop((max(0, x0 - pad), max(0, y0 - pad), min(im.width, x1 + pad), min(im.height, y1 + pad)))


def fit_pixels(im, max_pixels=MAX_PIXELS, min_side=64, mult=32):
    from PIL import Image
    w, h = im.size
    s = min(1.0, math.sqrt(max_pixels / (w * h)))
    nw, nh = max(min_side, int(w * s) // mult * mult), max(min_side, int(h * s) // mult * mult)
    if (nw, nh) != (w, h):
        im = im.resize((nw, nh), Image.LANCZOS if s < 1 else Image.BICUBIC)
    return im


def prepare_page(image, max_pixels=MAX_PIXELS):
    """A path or PIL image -> the RGB image the model was trained on."""
    from PIL import Image
    im = Image.open(image) if not isinstance(image, Image.Image) else image
    im.load()
    return fit_pixels(trim_margins(im.convert("RGB")), max_pixels)


def load(model=REPO, revision=None, device=None):
    import torch
    from transformers import AutoModelForImageTextToText, AutoProcessor
    device = device or ("cuda" if torch.cuda.is_available() else "cpu")
    dtype = torch.bfloat16 if device == "cuda" else torch.float32
    proc = AutoProcessor.from_pretrained(model, revision=revision)
    proc.tokenizer.padding_side = "left"
    net = AutoModelForImageTextToText.from_pretrained(model, revision=revision, dtype=dtype).to(device).eval()
    return net, proc


def prompt_text(proc):
    msgs = [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": PROMPT}]}]
    return proc.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True, enable_thinking=False)


def transcribe(pages, net, proc, batch=4, max_new_tokens=MAX_NEW_TOKENS, max_pixels=MAX_PIXELS):
    """pages: paths or PIL images -> list of page texts (lines separated by \\n)."""
    import torch
    text = prompt_text(proc)
    out = []
    for i in range(0, len(pages), batch):
        ims = [prepare_page(p, max_pixels) for p in pages[i:i + batch]]
        x = proc(text=[text] * len(ims), images=ims, return_tensors="pt", padding=True).to(net.device)
        with torch.no_grad():
            g = net.generate(**x, max_new_tokens=max_new_tokens, do_sample=False)
        for ids in g[:, x["input_ids"].shape[1]:]:
            out.append(proc.tokenizer.decode(ids, skip_special_tokens=True).strip())
    return out


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("pages", nargs="+")
    ap.add_argument("--model", default=REPO)
    ap.add_argument("--revision", default=None)
    ap.add_argument("--batch", type=int, default=4)
    ap.add_argument("--max_new_tokens", type=int, default=MAX_NEW_TOKENS)
    ap.add_argument("--out_dir", default=None, help="write <page stem>.txt here instead of printing")
    a = ap.parse_args()
    net, proc = load(a.model, a.revision)
    texts = transcribe(a.pages, net, proc, batch=a.batch, max_new_tokens=a.max_new_tokens)
    for p, t in zip(a.pages, texts):
        if a.out_dir:
            Path(a.out_dir).mkdir(parents=True, exist_ok=True)
            (Path(a.out_dir) / f"{Path(p).stem}.txt").write_text(t + "\n", encoding="utf-8")
        else:
            sys.stdout.write(f"=== {p}\n{t}\n")


if __name__ == "__main__":
    main()