"""Local handwriting OCR using HuggingFace VisionEncoderDecoder models (offline, free). PaddleOCR's printed-text recogniser is weak on handwriting. This module pairs PaddleOCR's *detector* (to find text-line boxes) with a TrOCR-family handwriting *recogniser* (a local HuggingFace transformers model) to read each line. The recogniser is selected by language + an optional 'large' quality flag: en (default) -> microsoft/trocr-base-handwritten (English, fast-ish) en + large -> microsoft/trocr-large-handwritten (English, slower, more accurate) hi (Hindi) -> sabaridsnfuji/Hindi_Offline_Handwritten_OCR (Hindi, EXPERIMENTAL) ne (Nepali / generic Devanagari) -> aayushpuri01/TrOCR-Devanagari It is OPTIONAL and heavy: it needs ``transformers`` + ``torch`` installed and a one-time model download (hundreds of MB). It is slow on CPU and opt-in per document. If the deps/model are missing it raises a clear, actionable error instead of crashing, and unknown languages fall back to the English model. Backward compatibility: ``lang='en'`` + ``large=False`` (the defaults, and the module-level ``hw_engine`` singleton) behave byte-for-byte like the original single-model implementation. """ import os import threading import cv2 import numpy as np # ---------------------------------------------------------------------------- # Model registry. Each entry describes one recogniser. # # id : HuggingFace repo for the VisionEncoderDecoderModel weights. # processor : how to build the TrOCRProcessor: # "bundled" -> TrOCRProcessor.from_pretrained(id) (id ships the # preprocessor_config + tokenizer) # ("manual", feat_repo, tok_repo) -> build from a ViT feature # extractor repo + a decoder tokenizer repo, because # the weights repo ships ONLY config + safetensors. # max_new_tokens : decode cap (the Hindi model is trained for <=64 chars/crop). # experimental : True for models whose quality on this task is unvalidated. # ---------------------------------------------------------------------------- _MODELS = { "en": { "id": "microsoft/trocr-base-handwritten", "processor": "bundled", "max_new_tokens": 96, "experimental": False, }, "en_large": { "id": "microsoft/trocr-large-handwritten", "processor": "bundled", "max_new_tokens": 96, "experimental": False, }, "hi": { # Only genuinely Hindi-trained option found. Repo ships weights only, so # the processor must be assembled from the encoder + decoder base repos. # Quality is modest and UNVALIDATED here — treat as experimental. "id": "sabaridsnfuji/Hindi_Offline_Handwritten_OCR", "processor": ("manual", "google/vit-base-patch16-224-in21k", "surajp/RoBERTa-hindi-guj-san"), "max_new_tokens": 64, # model is trained for <=64 chars per crop "experimental": True, }, "ne": { # Nepali Devanagari; bundled processor, clean drop-in. Usable as a # generic Devanagari fallback (reads many Hindi glyphs, Nepali-biased). "id": "aayushpuri01/TrOCR-Devanagari", "processor": "bundled", "max_new_tokens": 64, "experimental": True, }, } # Default fallback when a requested language has no handwriting model. _DEFAULT_KEY = "en" # Backwards-compat constant some callers/tests may import. MODEL_NAME = _MODELS[_DEFAULT_KEY]["id"] def _resolve_key(lang: str, large: bool) -> str: """Map a (lang, large) request to a registry key, falling back to English. Devanagari-family languages share script; 'hi'/'ne'/'devanagari' all route to a Devanagari model. 'large' currently only applies to English (no large Devanagari model exists), so it is ignored for non-English. """ lang = (lang or "en").strip().lower() if lang in ("en", "english"): return "en_large" if large else "en" if lang in ("hi", "hin", "hindi", "mr", "marathi", "sa"): return "hi" if lang in ("ne", "nep", "nepali", "devanagari"): return "ne" return _DEFAULT_KEY def is_available() -> bool: """True if the optional handwriting dependencies are importable.""" try: import torch # noqa: F401 import transformers # noqa: F401 return True except Exception: return False def _dist(a, b) -> float: return float(np.hypot(a[0] - b[0], a[1] - b[1])) def _crop_quad(img: np.ndarray, box) -> np.ndarray: """Perspective-warp a 4-point text box to an upright rectangle crop.""" pts = np.array(box, dtype="float32") w = int(max(_dist(pts[0], pts[1]), _dist(pts[2], pts[3]))) h = int(max(_dist(pts[1], pts[2]), _dist(pts[3], pts[0]))) w = max(w, 1) h = max(h, 1) dst = np.array([[0, 0], [w, 0], [w, h], [0, h]], dtype="float32") matrix = cv2.getPerspectiveTransform(pts, dst) return cv2.warpPerspective(img, matrix, (w, h)) class HandwritingEngine: """Lazy-loaded TrOCR-family wrapper for ONE registry entry. Thread-safe; one model instance reused. Construct via the module-level ``get_engine(lang, large)`` factory so identical models are shared. """ def __init__(self, key: str = _DEFAULT_KEY) -> None: self._key = key if key in _MODELS else _DEFAULT_KEY self._cfg = _MODELS[self._key] self._lock = threading.Lock() self._model = None self._processor = None self._loaded = False self._error = None @property def model_id(self) -> str: return self._cfg["id"] @property def experimental(self) -> bool: return bool(self._cfg.get("experimental")) def _build_processor(self): from transformers import TrOCRProcessor proc = self._cfg["processor"] if proc == "bundled": return TrOCRProcessor.from_pretrained(self._cfg["id"]) # ("manual", feature_extractor_repo, tokenizer_repo) _, feat_repo, tok_repo = proc from transformers import AutoTokenizer # AutoImageProcessor is the modern name; in transformers 4.40.2 # AutoFeatureExtractor still resolves it fine as a fallback. try: from transformers import AutoImageProcessor as _FeatLoader except Exception: # very old fallback from transformers import AutoFeatureExtractor as _FeatLoader feature_extractor = _FeatLoader.from_pretrained(feat_repo) tokenizer = AutoTokenizer.from_pretrained(tok_repo) # TrOCRProcessor accepts an image processor as `image_processor=`; older # transformers used `feature_extractor=`. 4.40.2 accepts both kwargs. try: return TrOCRProcessor(image_processor=feature_extractor, tokenizer=tokenizer) except TypeError: return TrOCRProcessor(feature_extractor=feature_extractor, tokenizer=tokenizer) def _ensure(self) -> None: if self._loaded: if self._error: raise RuntimeError(self._error) return with self._lock: if self._loaded: if self._error: raise RuntimeError(self._error) return try: import torch from transformers import VisionEncoderDecoderModel torch.set_num_threads(max(1, os.cpu_count() or 1)) self._processor = self._build_processor() self._model = VisionEncoderDecoderModel.from_pretrained( self._cfg["id"] ) self._model.eval() self._loaded = True except (ImportError, ModuleNotFoundError) as exc: # Missing deps are DETERMINISTIC — they won't fix themselves at # runtime, so cache the failure permanently (fail fast). self._error = ( "Handwriting model '{}' unavailable ({}). Enable it with: " "pip install transformers torch (first use downloads the " "model, a few hundred MB; non-English models also fetch a " "tokenizer/feature-extractor repo).".format( self._cfg["id"], exc ) ) self._loaded = True raise RuntimeError(self._error) except Exception as exc: # TRANSIENT failure (network drop / partial HF download / HF # outage / disk full). Do NOT poison the engine: leave _loaded # False so the next request retries from_pretrained (which can # resume the partial download). Otherwise the engine would be # dead for the whole process until a server restart. raise RuntimeError( "Handwriting model '{}' could not be loaded ({}). This is " "often a transient network/download issue — try again; a " "later request will retry the download.".format( self._cfg["id"], exc ) ) from exc def _recognize(self, crop_bgr: np.ndarray) -> str: import torch from PIL import Image rgb = cv2.cvtColor(crop_bgr, cv2.COLOR_BGR2RGB) pil = Image.fromarray(rgb) pixel_values = self._processor( images=pil, return_tensors="pt" ).pixel_values with torch.no_grad(): generated = self._model.generate( pixel_values, max_new_tokens=self._cfg["max_new_tokens"] ) return self._processor.batch_decode( generated, skip_special_tokens=True )[0] def ocr_text(self, img_bgr: np.ndarray, boxes: list) -> str: """Recognize handwriting in ``img_bgr`` given detected line ``boxes``. Boxes are ordered into reading order (top-to-bottom, left-to-right), each cropped and passed through the recogniser, joined by newlines. """ self._ensure() if not boxes: return "" items = [] for box in boxes: # Skip any malformed box (not a 4-point quad) so a single bad # detection can't abort the whole page's handwriting OCR. if not box or len(box) != 4: continue try: ys = [p[1] for p in box] xs = [p[0] for p in box] except (TypeError, IndexError): continue items.append({ "box": box, "cy": sum(ys) / len(ys), "cx": sum(xs) / len(xs), "h": max(ys) - min(ys), }) if not items: return "" heights = [it["h"] for it in items if it["h"] > 0] tol = max((float(np.median(heights)) if heights else 12.0) * 0.6, 8.0) items.sort(key=lambda it: (it["cy"], it["cx"])) rows = [[items[0]]] cy = items[0]["cy"] for it in items[1:]: if abs(it["cy"] - cy) <= tol: rows[-1].append(it) else: rows.append([it]) cy = it["cy"] lines = [] for row in rows: for it in sorted(row, key=lambda x: x["cx"]): try: crop = _crop_quad(img_bgr, it["box"]) if crop.size == 0: continue text = self._recognize(crop).strip() except Exception: text = "" if text: lines.append(text) return "\n".join(lines) # ---------------------------------------------------------------------------- # Factory: one cached engine per registry key (so 'en', 'en_large' and a # Devanagari model can coexist without reloading). Thread-safe. # ---------------------------------------------------------------------------- _engines: dict = {} _engines_lock = threading.Lock() def get_engine(lang: str = "en", large: bool = False) -> "HandwritingEngine": """Return a cached handwriting engine for the requested (lang, large).""" key = _resolve_key(lang, large) eng = _engines.get(key) if eng is not None: return eng with _engines_lock: eng = _engines.get(key) if eng is None: eng = HandwritingEngine(key) _engines[key] = eng return eng # Backwards-compatible module-level singleton (English base), so existing # `from pipeline.handwriting import hw_engine` callers keep working unchanged. hw_engine = get_engine("en", large=False)