"""RecognitionPredictor: per-block OCR via BLOCK_PROMPT. Given page images and corresponding LayoutResult (or any list of LayoutBox), crops each block, runs BLOCK_PROMPT, returns PageOCRResult per page. """ from __future__ import annotations from typing import List, Optional from PIL import Image from surya.common.blank import is_blank_region from surya.inference import SuryaInferenceManager, get_default_manager from surya.inference.parsers import clean_block_html, parse_full_page_html from surya.inference.prompts import ( PROMPT_TYPE_BLOCK, PROMPT_TYPE_HIGH_ACCURACY_BBOX, SKIP_OCR_LABELS, ) from surya.inference.schema import BatchInputItem from surya.inference.util import image_token_budget from surya.layout.label import LAYOUT_PRED_RELABEL, TEXT_LABELS from surya.layout.schema import LayoutResult from surya.logging import get_logger from surya.recognition.schema import ( BlockOCRResult, PageOCRResult, ) from surya.settings import settings from surya.timing import timing_span logger = get_logger() # Surya's canonical labels we shouldn't OCR (mirrors model-emitted SKIP_OCR_LABELS # after canonicalization). SKIP_CANON_LABELS = {LAYOUT_PRED_RELABEL.get(lbl, lbl) for lbl in SKIP_OCR_LABELS} def _crop_block(image: Image.Image, polygon, pad: int = 4) -> Image.Image: xs = [p[0] for p in polygon] ys = [p[1] for p in polygon] x0 = max(0, int(min(xs)) - pad) y0 = max(0, int(min(ys)) - pad) x1 = min(image.size[0], int(max(xs)) + pad) y1 = min(image.size[1], int(max(ys)) + pad) if x1 <= x0 or y1 <= y0: return image.crop((0, 0, 1, 1)) return image.crop((x0, y0, x1, y1)) def _drop_blank_text_blocks( image: Image.Image, blocks: List[BlockOCRResult], ) -> List[BlockOCRResult]: """Drop text-labeled blocks whose source page region is essentially blank. Full-page OCR can emit text divs for regions that are visually empty (margins, gutter space) — the model hallucinates a paragraph where there is none. We crop the region, count near-white pixels, and drop the block when the fraction exceeds ``blank_pixel_fraction``. Only text-like labels (see ``TEXT_LABELS``) are eligible: tables, forms, equations, and visual blocks may legitimately contain large whitespace and are left untouched. """ kept: List[BlockOCRResult] = [] dropped = 0 for blk in blocks: if blk.label not in TEXT_LABELS or blk.skipped or blk.error: kept.append(blk) continue crop = _crop_block(image, blk.polygon) if not is_blank_region(crop): kept.append(blk) continue dropped += 1 if dropped: logger.info(f"dropped {dropped} blank text block(s) from full-page OCR") return kept def _detect_repeat_loop( text: str, base_max_repeats: int = 4, window_size: int = 500, scaling_factor: float = 3.0, ) -> bool: """True iff the tail of ``text`` ends in a repeating sequence. Ported from chandra's detect_repeat_token. For each candidate length 1..window_size/2, takes that many trailing chars and counts consecutive identical preceding blocks. Shorter loops need many repeats to count; longer ones only need a few. Catches the typical decoder failure mode where a page output gets stuck emitting the same div / phrase until it hits max_tokens. """ if not text: return False for seq_len in range(1, window_size // 2 + 1): candidate = text[-seq_len:] max_repeats = int(base_max_repeats * (1 + scaling_factor / seq_len)) repeats = 0 pos = len(text) - seq_len while pos >= 0 and text[pos : pos + seq_len] == candidate: repeats += 1 pos -= seq_len if repeats > max_repeats: return True return False class RecognitionPredictor: """Per-block OCR. Construct with a SuryaInferenceManager (or rely on default).""" def __init__(self, manager: Optional[SuryaInferenceManager] = None): self.manager = manager self._disable_tqdm = settings.DISABLE_TQDM @property def disable_tqdm(self) -> bool: return self._disable_tqdm @disable_tqdm.setter def disable_tqdm(self, value: bool) -> None: self._disable_tqdm = bool(value) def to(self, *args, **kwargs): return def __call__( self, images: List[Image.Image], layout_results: Optional[List[LayoutResult]] = None, *, full_page: Optional[bool] = None, ) -> List[PageOCRResult]: """Run OCR on each page. Mode resolution: - ``full_page=None`` (default): block mode if ``layout_results`` is given, else full-page mode. This is the most-do-what-I-mean form. - ``full_page=True``: full-page OCR (single HIGH_ACCURACY_BBOX_PROMPT request per page). ``layout_results`` is ignored — a warning is logged if it was supplied. - ``full_page=False``: block mode (per-layout-block OCR request). ``layout_results`` is required. Full-page is the more accurate path; block mode is for callers that specifically need per-block crops (e.g. for downstream merging with text-line detection). """ if not images: return [] if full_page is None: full_page = layout_results is None if full_page: if layout_results is not None: logger.info( "RecognitionPredictor called with full_page=True and " "layout_results; layout will be used as fallback if the " "full-page output devolves into a repetition loop." ) return self._full_page_ocr(images, fallback_layout=layout_results) if layout_results is None: raise ValueError("layout_results required when full_page=False") if len(images) != len(layout_results): raise ValueError( f"images and layout_results must be same length " f"({len(images)} vs {len(layout_results)})" ) manager = self.manager or get_default_manager() # Build a flat batch across all pages for max concurrency batch: List[BatchInputItem] = [] block_index_map: List[tuple[int, int]] = [] # (page_idx, block_idx) skipped_flags: List[bool] = [] with timing_span("recognition_build_block_batch", page_count=len(images)): for page_idx, (img, layout) in enumerate(zip(images, layout_results)): page_boxes = sorted(layout.bboxes, key=lambda b: (b.position, b.bbox[1], b.bbox[0])) page_boxes = page_boxes[: settings.SURYA_MAX_BLOCKS_PER_PAGE] if len(layout.bboxes) > len(page_boxes): logger.info( f"capped OCR blocks for page {page_idx}: " f"{len(layout.bboxes)} -> {len(page_boxes)}" ) for block_idx, box in enumerate(page_boxes): skip = box.label in SKIP_CANON_LABELS skipped_flags.append(skip) if skip: continue crop = _crop_block(img, box.polygon) max_tokens = image_token_budget( box.count, ceiling=settings.SURYA_MAX_TOKENS_BLOCK_CEILING ) batch.append( BatchInputItem( image=crop, prompt_type=PROMPT_TYPE_BLOCK, max_tokens=max_tokens, metadata={"page_idx": page_idx, "block_idx": block_idx}, ) ) block_index_map.append((page_idx, block_idx)) with timing_span( "recognition_manager_generate", item_count=len(batch), skipped_count=sum(1 for flag in skipped_flags if flag), max_tokens_sum=sum(item.max_tokens or 0 for item in batch), ): outputs = manager.generate(batch) if batch else [] # Index outputs by (page_idx, block_idx) with timing_span( "recognition_assemble_pages", output_count=len(outputs), token_count=sum(out.token_count or 0 for out in outputs), ): out_by_key = {} for out in outputs: key = (out.metadata["page_idx"], out.metadata["block_idx"]) out_by_key[key] = out # Assemble PageOCRResult per page results: List[PageOCRResult] = [] for page_idx, (img, layout) in enumerate(zip(images, layout_results)): w, h = img.size blocks: List[BlockOCRResult] = [] page_boxes = sorted(layout.bboxes, key=lambda b: (b.position, b.bbox[1], b.bbox[0])) page_boxes = page_boxes[: settings.SURYA_MAX_BLOCKS_PER_PAGE] for block_idx, box in enumerate(page_boxes): skip = box.label in SKIP_CANON_LABELS if skip: blocks.append( BlockOCRResult( polygon=box.polygon, label=box.label, raw_label=box.raw_label, reading_order=box.position, html="", skipped=True, confidence=1.0, ) ) continue out = out_by_key.get((page_idx, block_idx)) if out is None or out.error: blocks.append( BlockOCRResult( polygon=box.polygon, label=box.label, raw_label=box.raw_label, reading_order=box.position, html="", skipped=False, error=True, confidence=0.0, ) ) continue html = clean_block_html(out.raw) conf = out.mean_token_prob if out.mean_token_prob is not None else 1.0 blocks.append( BlockOCRResult( polygon=box.polygon, label=box.label, raw_label=box.raw_label, reading_order=box.position, html=html, skipped=False, error=False, confidence=conf, raw_logprobs=out.logprobs, ) ) results.append( PageOCRResult(blocks=blocks, image_bbox=[0, 0, float(w), float(h)]) ) return results def _full_page_ocr( self, images: List[Image.Image], fallback_layout: Optional[List[LayoutResult]] = None, ) -> List[PageOCRResult]: """One HIGH_ACCURACY_BBOX_PROMPT request per page; parses divs into blocks. On per-page failure (parse error, empty output, or a detected repetition loop in the decoder output), falls back to layout + block-mode OCR for that page only. ``fallback_layout``, if given, provides per-page LayoutResults to use on fallback; otherwise the LayoutPredictor is invoked lazily for just the affected pages. """ manager = self.manager or get_default_manager() with timing_span("recognition_build_full_page_batch", page_count=len(images)): batch = [ BatchInputItem( image=img, prompt_type=PROMPT_TYPE_HIGH_ACCURACY_BBOX, max_tokens=settings.SURYA_MAX_TOKENS_FULL_PAGE, metadata={"page_idx": i}, ) for i, img in enumerate(images) ] with timing_span( "recognition_full_page_generate", item_count=len(batch), max_tokens=settings.SURYA_MAX_TOKENS_FULL_PAGE, ): outputs = manager.generate(batch) out_by_page = {o.metadata["page_idx"]: o for o in outputs} results: List[Optional[PageOCRResult]] = [None] * len(images) needs_fallback: List[int] = [] for page_idx, img in enumerate(images): w, h = img.size page_bbox = [0, 0, float(w), float(h)] out = out_by_page.get(page_idx) if out is None or out.error: # Hard failure (request lost / server error). Always fallback. needs_fallback.append(page_idx) continue if not out.raw: # Empty model output. If the page is genuinely blank, the # model is correct — return an empty result. Only fall back # when the page has content the model failed to emit. if is_blank_region(img): results[page_idx] = PageOCRResult(blocks=[], image_bbox=page_bbox) else: logger.info( f"empty full-page output for non-blank page {page_idx}; " f"falling back to layout + block OCR" ) needs_fallback.append(page_idx) continue if _detect_repeat_loop(out.raw): logger.info( f"full-page output for page {page_idx} appears to loop; " f"falling back to layout + block OCR" ) needs_fallback.append(page_idx) continue try: parsed = parse_full_page_html(out.raw) except Exception as e: logger.warning( f"Full-page parse failed for page {page_idx}: {e}; " f"falling back to layout + block OCR" ) needs_fallback.append(page_idx) continue confidence = out.mean_token_prob if out.mean_token_prob is not None else 1.0 blocks: List[BlockOCRResult] = [] for idx, item in enumerate(parsed): x0 = item.bbox[0] / settings.BBOX_SCALE * w y0 = item.bbox[1] / settings.BBOX_SCALE * h x1 = item.bbox[2] / settings.BBOX_SCALE * w y1 = item.bbox[3] / settings.BBOX_SCALE * h polygon = [[x0, y0], [x1, y0], [x1, y1], [x0, y1]] canon = LAYOUT_PRED_RELABEL.get(item.label, item.label) skipped = canon in SKIP_CANON_LABELS blocks.append( BlockOCRResult( polygon=polygon, label=canon, raw_label=item.label, reading_order=idx, html="" if skipped else item.html, skipped=skipped, error=False, confidence=confidence, ) ) blocks = _drop_blank_text_blocks(img, blocks) results[page_idx] = PageOCRResult(blocks=blocks, image_bbox=page_bbox) # Block-mode fallback for any pages whose full-page output failed or looped. if needs_fallback: fb_images = [images[i] for i in needs_fallback] if fallback_layout is not None: fb_layouts = [fallback_layout[i] for i in needs_fallback] else: # Lazy import to avoid the surya.layout ↔ surya.recognition cycle. from surya.layout import LayoutPredictor logger.info( f"running layout for {len(fb_images)} page(s) requiring " f"block-mode fallback" ) fb_layouts = LayoutPredictor(self.manager)(fb_images) fb_results = self.__call__(fb_images, fb_layouts, full_page=False) for fb_idx, page_idx in enumerate(needs_fallback): results[page_idx] = fb_results[fb_idx] # Backfill any still-None pages with empty results (defensive — shouldn't happen). out_results: List[PageOCRResult] = [] for page_idx, img in enumerate(images): r = results[page_idx] if r is None: w, h = img.size r = PageOCRResult(blocks=[], image_bbox=[0, 0, float(w), float(h)]) out_results.append(r) return out_results