Suya OCR API — vLLM-backed, OpenAI-compatible OCR service
FastAPI service wrapping the Surya-OCR-2 model (datalab-to) served through vLLM: legacy /v1/api/ai/* endpoints, an OpenAI-compatible /v1/chat/completions endpoint, a coalescing request batcher, a local OCR CLI, Docker packaging, multilingual example outputs, and quantization/concurrency benchmarks. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Fu Dai
@@ -0,0 +1,141 @@
|
||||
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 denorm_bbox, parse_layout
|
||||
from surya.inference.prompts import LAYOUT_JSON_SCHEMA, PROMPT_TYPE_LAYOUT
|
||||
from surya.inference.schema import BatchInputItem
|
||||
from surya.layout.label import LAYOUT_PRED_RELABEL, TEXT_LABELS
|
||||
from surya.layout.schema import LayoutBox, LayoutResult
|
||||
from surya.logging import get_logger
|
||||
from surya.settings import settings
|
||||
from surya.timing import timing_span
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
|
||||
class LayoutPredictor:
|
||||
"""Run LAYOUT_PROMPT on full pages, parse JSON, return LayoutResult per image."""
|
||||
|
||||
def __init__(self, manager: Optional[SuryaInferenceManager] = None):
|
||||
self.manager = manager # If None, get_default_manager() is used at call time
|
||||
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):
|
||||
# Manager-backed; .to() is a no-op for compatibility with BasePredictor callers.
|
||||
return
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
images: List[Image.Image],
|
||||
target_image_sizes: Optional[List[tuple]] = None,
|
||||
max_tokens: Optional[int] = None,
|
||||
) -> List[LayoutResult]:
|
||||
"""Run layout on a batch of images.
|
||||
|
||||
target_image_sizes: optional list of (width, height) tuples — if
|
||||
provided, bboxes are denormalized to these sizes instead of each
|
||||
input image's size. Useful when layout runs on a low-DPI render but
|
||||
you want bboxes in the OCR image's coordinate space.
|
||||
"""
|
||||
if not images:
|
||||
return []
|
||||
manager = self.manager or get_default_manager()
|
||||
|
||||
max_tokens = max_tokens or settings.SURYA_MAX_TOKENS_LAYOUT
|
||||
guided = LAYOUT_JSON_SCHEMA if settings.SURYA_GUIDED_LAYOUT else None
|
||||
with timing_span("layout_build_batch", image_count=len(images), max_tokens=max_tokens):
|
||||
batch = [
|
||||
BatchInputItem(
|
||||
image=img,
|
||||
prompt_type=PROMPT_TYPE_LAYOUT,
|
||||
max_tokens=max_tokens,
|
||||
guided_json=guided,
|
||||
)
|
||||
for img in images
|
||||
]
|
||||
with timing_span("layout_manager_generate", item_count=len(batch)):
|
||||
outputs = manager.generate(batch)
|
||||
|
||||
if target_image_sizes is not None and len(target_image_sizes) != len(images):
|
||||
raise ValueError("target_image_sizes must match images length")
|
||||
|
||||
with timing_span("layout_parse_outputs", item_count=len(outputs)):
|
||||
results: List[LayoutResult] = []
|
||||
for idx, (img, out) in enumerate(zip(images, outputs)):
|
||||
if target_image_sizes is not None:
|
||||
w, h = target_image_sizes[idx]
|
||||
else:
|
||||
w, h = img.size
|
||||
page_bbox = [0, 0, float(w), float(h)]
|
||||
if out.error or not out.raw:
|
||||
results.append(
|
||||
LayoutResult(
|
||||
bboxes=[], image_bbox=page_bbox, raw=out.raw, error=True
|
||||
)
|
||||
)
|
||||
continue
|
||||
try:
|
||||
parsed = parse_layout(out.raw)
|
||||
except Exception as e:
|
||||
logger.warning(f"Layout parse failed: {e}; raw[:300]={out.raw[:300]!r}")
|
||||
results.append(
|
||||
LayoutResult(
|
||||
bboxes=[], image_bbox=page_bbox, raw=out.raw, error=True
|
||||
)
|
||||
)
|
||||
continue
|
||||
|
||||
confidence = out.mean_token_prob if out.mean_token_prob is not None else 1.0
|
||||
img_w, img_h = img.size
|
||||
boxes: List[LayoutBox] = []
|
||||
dropped_blank = 0
|
||||
for blk in parsed:
|
||||
canon = LAYOUT_PRED_RELABEL.get(blk.label, blk.label)
|
||||
# Drop text-labeled blocks the model hallucinated over an
|
||||
# essentially-blank region (mostly white OR near-uniform
|
||||
# color). Visual blocks (Picture / Figure / Table / etc.)
|
||||
# are allowed to be uniform — that's normal content.
|
||||
if canon in TEXT_LABELS:
|
||||
img_bbox = denorm_bbox(
|
||||
blk.bbox, img_w, img_h, scale=settings.BBOX_SCALE
|
||||
)
|
||||
x0, y0, x1, y1 = (max(0, int(v)) for v in img_bbox)
|
||||
if x1 > x0 and y1 > y0:
|
||||
if is_blank_region(img.crop((x0, y0, x1, y1))):
|
||||
dropped_blank += 1
|
||||
continue
|
||||
pixel_bbox = denorm_bbox(blk.bbox, w, h, scale=settings.BBOX_SCALE)
|
||||
boxes.append(
|
||||
LayoutBox(
|
||||
polygon=list(pixel_bbox),
|
||||
label=canon,
|
||||
raw_label=blk.label,
|
||||
position=len(boxes),
|
||||
count=blk.count,
|
||||
confidence=confidence,
|
||||
)
|
||||
)
|
||||
if dropped_blank:
|
||||
logger.info(
|
||||
f"dropped {dropped_blank} text-labeled layout block(s) over "
|
||||
f"blank/uniform regions"
|
||||
)
|
||||
results.append(
|
||||
LayoutResult(
|
||||
bboxes=boxes, image_bbox=page_bbox, raw=out.raw, error=False
|
||||
)
|
||||
)
|
||||
return results
|
||||
@@ -0,0 +1,44 @@
|
||||
"""Surya2 layout labels emitted by the model + canonicalization to surya's
|
||||
public label vocabulary."""
|
||||
|
||||
# Canonical text-bearing labels — used by blank-region filters to decide
|
||||
# which blocks may be dropped when their underlying image region is empty.
|
||||
# Excludes Picture/Figure/Diagram/Table/Form/Equation/etc., which can legitimately
|
||||
# contain whitespace or solid fills.
|
||||
TEXT_LABELS = frozenset(
|
||||
{
|
||||
"Text",
|
||||
"SectionHeader",
|
||||
"PageHeader",
|
||||
"PageFooter",
|
||||
"Caption",
|
||||
"Footnote",
|
||||
"Code",
|
||||
"Bibliography",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
# Canonicalize raw model labels to public surya label names. Marker and other
|
||||
# downstream consumers depend on these names.
|
||||
LAYOUT_PRED_RELABEL = {
|
||||
"Caption": "Caption",
|
||||
"Footnote": "Footnote",
|
||||
"Equation-Block": "Equation",
|
||||
"List-Group": "ListGroup",
|
||||
"Page-Header": "PageHeader",
|
||||
"Page-Footer": "PageFooter",
|
||||
"Image": "Picture",
|
||||
"Section-Header": "SectionHeader",
|
||||
"Table": "Table",
|
||||
"Text": "Text",
|
||||
"Complex-Block": "Figure",
|
||||
"Code-Block": "Code",
|
||||
"Form": "Form",
|
||||
"Table-Of-Contents": "TableOfContents",
|
||||
"Figure": "Figure",
|
||||
"Chemical-Block": "ChemicalBlock",
|
||||
"Diagram": "Diagram",
|
||||
"Bibliography": "Bibliography",
|
||||
"Blank-Page": "BlankPage",
|
||||
}
|
||||
@@ -0,0 +1,19 @@
|
||||
from typing import List, Optional
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from surya.common.polygon import PolygonBox
|
||||
|
||||
|
||||
class LayoutBox(PolygonBox):
|
||||
label: str # canonicalized via LAYOUT_PRED_RELABEL
|
||||
raw_label: str # original model label, before canonicalization
|
||||
position: int # reading order index
|
||||
count: int = 0 # model's token estimate for OCR output (multiple of 50)
|
||||
|
||||
|
||||
class LayoutResult(BaseModel):
|
||||
bboxes: List[LayoutBox]
|
||||
image_bbox: List[float]
|
||||
raw: Optional[str] = None # raw model output, useful for debugging
|
||||
error: bool = False
|
||||
Reference in New Issue
Block a user