wehub-resource-sync
389 lines
15 KiB
Python
389 lines
15 KiB
Python
"""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
|
|
|
|
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}
|
|
|
|
# Full-page OCR regeneration schedule (chandra-style), used ONLY when
|
|
# settings.SURYA_FULLPAGE_REGEN is True. Round 0 is greedy; a page whose output
|
|
# loops / fails to parse is re-requested at escalating temperature (top_p 0.95)
|
|
# before resorting to the (slower) block-mode fallback. Mirrors chandra's
|
|
# retry_temperature = min(0.2*(n+1), 0.8) over MAX_VLLM_RETRIES=6 retries.
|
|
_REGEN_ROUNDS = [(0.0, None)] + [(min(0.2 * (n + 1), 0.8), 0.95) for n in range(6)]
|
|
|
|
|
|
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] = []
|
|
|
|
for page_idx, (img, layout) in enumerate(zip(images, layout_results)):
|
|
for block_idx, box in enumerate(layout.bboxes):
|
|
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))
|
|
|
|
outputs = manager.generate(batch) if batch else []
|
|
|
|
# Index outputs by (page_idx, block_idx)
|
|
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] = []
|
|
for block_idx, box in enumerate(layout.bboxes):
|
|
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()
|
|
results: List[Optional[PageOCRResult]] = [None] * len(images)
|
|
|
|
def _build_page(out, img):
|
|
"""A good full-page output -> PageOCRResult; otherwise None (regen/fallback).
|
|
A genuinely blank page is a valid empty result (not a failure)."""
|
|
w, h = img.size
|
|
page_bbox = [0, 0, float(w), float(h)]
|
|
if out is None or out.error:
|
|
return None
|
|
if not out.raw:
|
|
return (
|
|
PageOCRResult(blocks=[], image_bbox=page_bbox)
|
|
if is_blank_region(img)
|
|
else None
|
|
)
|
|
if _detect_repeat_loop(out.raw):
|
|
return None
|
|
try:
|
|
parsed = parse_full_page_html(out.raw)
|
|
except Exception:
|
|
return None
|
|
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,
|
|
)
|
|
)
|
|
return PageOCRResult(
|
|
blocks=_drop_blank_text_blocks(img, blocks), image_bbox=page_bbox
|
|
)
|
|
|
|
# Progressive-temperature regeneration (chandra-style), gated by
|
|
# settings.SURYA_FULLPAGE_REGEN (default off -> single greedy pass, then
|
|
# block-mode fallback, i.e. surya's prior behavior).
|
|
rounds = _REGEN_ROUNDS if settings.SURYA_FULLPAGE_REGEN else [(0.0, None)]
|
|
pending = list(range(len(images)))
|
|
for round_i, (temp, top_p) in enumerate(rounds):
|
|
if not pending:
|
|
break
|
|
batch = [
|
|
BatchInputItem(
|
|
image=images[i],
|
|
prompt_type=PROMPT_TYPE_HIGH_ACCURACY_BBOX,
|
|
max_tokens=settings.SURYA_MAX_TOKENS_FULL_PAGE,
|
|
temperature=(temp if round_i > 0 else None),
|
|
top_p=top_p,
|
|
metadata={"page_idx": i},
|
|
)
|
|
for i in pending
|
|
]
|
|
out_by_page = {o.metadata["page_idx"]: o for o in manager.generate(batch)}
|
|
still: List[int] = []
|
|
for i in pending:
|
|
r = _build_page(out_by_page.get(i), images[i])
|
|
if r is not None:
|
|
results[i] = r
|
|
else:
|
|
still.append(i)
|
|
if still and round_i < len(rounds) - 1:
|
|
logger.info(
|
|
f"regenerating {len(still)} full-page output(s) at higher temperature"
|
|
)
|
|
pending = still
|
|
needs_fallback: List[int] = pending
|
|
|
|
# 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
|