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,3 @@
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,64 @@
|
||||
"""Pixel-content heuristics for detecting blank or near-uniform image regions.
|
||||
|
||||
Used by both the layout predictor (drop hallucinated layout blocks over empty
|
||||
space) and the recognition predictor (drop hallucinated text blocks from
|
||||
full-page OCR, decide whether an empty full-page output is a correct blank-page
|
||||
read or a failure).
|
||||
|
||||
Two signals, combined:
|
||||
* near-white fraction — most pixels have every RGB channel above a threshold
|
||||
* pixel-value standard deviation — the region is essentially one color
|
||||
(catches uniform-color fills that the white check misses)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
|
||||
# Per-channel value at/above which a pixel is considered "near-white".
|
||||
# Tolerates the small noise typical of PDF renders at 96 DPI.
|
||||
BLANK_WHITE_THRESHOLD = 245
|
||||
# Fraction of pixels that must be near-white for a region to count as blank.
|
||||
BLANK_PIXEL_FRACTION = 0.99
|
||||
# Pixel-value std below which a region is "essentially one color" regardless
|
||||
# of what that color is (catches solid-fill rectangles, dark banners, etc.).
|
||||
UNIFORM_COLOR_STD = 8.0
|
||||
|
||||
|
||||
def near_white_fraction(
|
||||
image: Image.Image, white_threshold: int = BLANK_WHITE_THRESHOLD
|
||||
) -> float:
|
||||
"""Fraction of pixels where every RGB channel ≥ ``white_threshold``."""
|
||||
arr = np.asarray(image.convert("RGB"))
|
||||
if arr.size == 0:
|
||||
return 0.0
|
||||
return float(np.all(arr >= white_threshold, axis=-1).mean())
|
||||
|
||||
|
||||
def is_blank_region(
|
||||
image: Image.Image,
|
||||
*,
|
||||
white_threshold: int = BLANK_WHITE_THRESHOLD,
|
||||
blank_pixel_fraction: float = BLANK_PIXEL_FRACTION,
|
||||
uniform_color_std: float = UNIFORM_COLOR_STD,
|
||||
) -> bool:
|
||||
"""True iff the image is essentially blank — either mostly near-white or
|
||||
near-uniform color. Use this on a per-block crop or a whole page.
|
||||
|
||||
Returns False for empty (0-pixel) crops so callers don't accidentally
|
||||
treat a degenerate bbox as blank.
|
||||
"""
|
||||
arr = np.asarray(image.convert("RGB"))
|
||||
if arr.size == 0:
|
||||
return False
|
||||
if np.all(arr >= white_threshold, axis=-1).mean() > blank_pixel_fraction:
|
||||
return True
|
||||
# Per-channel std — a uniform solid color (e.g., red banner with RGB=(200,50,50))
|
||||
# has each channel constant across pixels, but mixing channels inflates the
|
||||
# aggregate std. Check each channel independently.
|
||||
per_channel_std = arr.reshape(-1, arr.shape[-1]).std(axis=0)
|
||||
if float(per_channel_std.max()) < uniform_color_std:
|
||||
return True
|
||||
return False
|
||||
@@ -0,0 +1,25 @@
|
||||
from typing import Optional, Any
|
||||
|
||||
import torch
|
||||
|
||||
from surya.settings import settings
|
||||
|
||||
|
||||
class ModelLoader:
|
||||
def __init__(self, checkpoint: Optional[str] = None):
|
||||
self.checkpoint = checkpoint
|
||||
|
||||
def model(
|
||||
self,
|
||||
device: torch.device | str | None = settings.TORCH_DEVICE_MODEL,
|
||||
dtype: Optional[torch.dtype | str] = settings.MODEL_DTYPE,
|
||||
attention_implementation: Optional[str] = None,
|
||||
) -> Any:
|
||||
raise NotImplementedError()
|
||||
|
||||
def processor(
|
||||
self,
|
||||
device: torch.device | str | None = settings.TORCH_DEVICE_MODEL,
|
||||
dtype: Optional[torch.dtype | str] = settings.MODEL_DTYPE,
|
||||
) -> Any:
|
||||
raise NotImplementedError()
|
||||
@@ -0,0 +1,122 @@
|
||||
import copy
|
||||
from typing import List, Optional
|
||||
|
||||
import numpy as np
|
||||
from pydantic import BaseModel, field_validator, computed_field
|
||||
import numbers
|
||||
|
||||
|
||||
class PolygonBox(BaseModel):
|
||||
polygon: List[List[float]]
|
||||
confidence: Optional[float] = None
|
||||
|
||||
@field_validator("polygon", mode="before")
|
||||
@classmethod
|
||||
def convert_bbox_to_polygon(cls, value):
|
||||
if isinstance(value, (list, tuple)) and len(value) == 4:
|
||||
if all(isinstance(x, numbers.Number) for x in value):
|
||||
value = [float(v) for v in value]
|
||||
x_min, y_min, x_max, y_max = value
|
||||
polygon = [
|
||||
[x_min, y_min],
|
||||
[x_max, y_min],
|
||||
[x_max, y_max],
|
||||
[x_min, y_max],
|
||||
]
|
||||
return polygon
|
||||
elif all(
|
||||
isinstance(point, (list, tuple)) and len(point) == 2 for point in value
|
||||
):
|
||||
value = [[float(v) for v in point] for point in value]
|
||||
return value
|
||||
elif isinstance(value, np.ndarray):
|
||||
if value.shape == (4, 2):
|
||||
return value.tolist()
|
||||
|
||||
raise ValueError(
|
||||
f"Input must be either a bbox [x_min, y_min, x_max, y_max] or a polygon with 4 corners [(x,y), (x,y), (x,y), (x,y)]. All values must be numeric. You passed {value} of type {type(value)}. The first value is of type {type(value[0])}."
|
||||
)
|
||||
|
||||
@property
|
||||
def height(self):
|
||||
return self.bbox[3] - self.bbox[1]
|
||||
|
||||
@property
|
||||
def width(self):
|
||||
return self.bbox[2] - self.bbox[0]
|
||||
|
||||
@property
|
||||
def area(self):
|
||||
return self.width * self.height
|
||||
|
||||
@computed_field
|
||||
@property
|
||||
def bbox(self) -> List[float]:
|
||||
x_coords = [point[0] for point in self.polygon]
|
||||
y_coords = [point[1] for point in self.polygon]
|
||||
return [min(x_coords), min(y_coords), max(x_coords), max(y_coords)]
|
||||
|
||||
def rescale(self, processor_size, image_size):
|
||||
# Point is in x, y format
|
||||
page_width, page_height = processor_size
|
||||
|
||||
img_width, img_height = image_size
|
||||
width_scaler = img_width / page_width
|
||||
height_scaler = img_height / page_height
|
||||
|
||||
for corner in self.polygon:
|
||||
corner[0] = int(corner[0] * width_scaler)
|
||||
corner[1] = int(corner[1] * height_scaler)
|
||||
|
||||
def round(self, divisor):
|
||||
for corner in self.polygon:
|
||||
corner[0] = int(corner[0] / divisor) * divisor
|
||||
corner[1] = int(corner[1] / divisor) * divisor
|
||||
|
||||
def fit_to_bounds(self, bounds):
|
||||
new_corners = copy.deepcopy(self.polygon)
|
||||
for corner in new_corners:
|
||||
corner[0] = max(min(corner[0], bounds[2]), bounds[0])
|
||||
corner[1] = max(min(corner[1], bounds[3]), bounds[1])
|
||||
self.polygon = new_corners
|
||||
|
||||
def expand(self, x_margin: float, y_margin: float):
|
||||
new_polygon = []
|
||||
x_margin = x_margin * self.width
|
||||
y_margin = y_margin * self.height
|
||||
for idx, poly in enumerate(self.polygon):
|
||||
if idx == 0:
|
||||
new_polygon.append([int(poly[0] - x_margin), int(poly[1] - y_margin)])
|
||||
elif idx == 1:
|
||||
new_polygon.append([int(poly[0] + x_margin), int(poly[1] - y_margin)])
|
||||
elif idx == 2:
|
||||
new_polygon.append([int(poly[0] + x_margin), int(poly[1] + y_margin)])
|
||||
elif idx == 3:
|
||||
new_polygon.append([int(poly[0] - x_margin), int(poly[1] + y_margin)])
|
||||
self.polygon = new_polygon
|
||||
|
||||
def intersection_area(self, other, x_margin=0, y_margin=0):
|
||||
x_overlap = self.x_overlap(other, x_margin)
|
||||
y_overlap = self.y_overlap(other, y_margin)
|
||||
return x_overlap * y_overlap
|
||||
|
||||
def x_overlap(self, other, x_margin=0):
|
||||
return max(
|
||||
0,
|
||||
min(self.bbox[2] + x_margin, other.bbox[2] + x_margin)
|
||||
- max(self.bbox[0] - x_margin, other.bbox[0] - x_margin),
|
||||
)
|
||||
|
||||
def y_overlap(self, other, y_margin=0):
|
||||
return max(
|
||||
0,
|
||||
min(self.bbox[3] + y_margin, other.bbox[3] + y_margin)
|
||||
- max(self.bbox[1] - y_margin, other.bbox[1] - y_margin),
|
||||
)
|
||||
|
||||
@property
|
||||
def center(self):
|
||||
return [(self.bbox[0] + self.bbox[2]) / 2, (self.bbox[1] + self.bbox[3]) / 2]
|
||||
|
||||
def __hash__(self):
|
||||
return hash(tuple(self.bbox))
|
||||
@@ -0,0 +1,57 @@
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
|
||||
from surya.common.load import ModelLoader
|
||||
from surya.settings import settings
|
||||
|
||||
|
||||
class BasePredictor:
|
||||
model_loader_cls = ModelLoader
|
||||
batch_size: Optional[int] = None
|
||||
default_batch_sizes = {"cpu": 1, "mps": 1, "cuda": 1}
|
||||
torch_dtype = settings.MODEL_DTYPE
|
||||
|
||||
@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 __init__(
|
||||
self,
|
||||
checkpoint: Optional[str] = None,
|
||||
device: torch.device | str | None = settings.TORCH_DEVICE_MODEL,
|
||||
dtype: Optional[torch.dtype | str] = None,
|
||||
attention_implementation: Optional[str] = None,
|
||||
):
|
||||
if dtype is None:
|
||||
dtype = self.torch_dtype
|
||||
|
||||
loader = self.model_loader_cls(checkpoint)
|
||||
self.model = loader.model(device, dtype, attention_implementation)
|
||||
self.processor = loader.processor()
|
||||
self._disable_tqdm = settings.DISABLE_TQDM
|
||||
|
||||
def to(self, device_dtype: torch.device | str | None = None):
|
||||
if hasattr(self, "model") and self.model:
|
||||
self.model.to(device_dtype)
|
||||
return
|
||||
# Predictors that don't own a torch model (e.g. VLM-backed predictors that
|
||||
# rely on an external server) treat .to() as a no-op.
|
||||
if hasattr(self, "manager") and self.manager is not None:
|
||||
return
|
||||
raise ValueError("Model not loaded")
|
||||
|
||||
def get_batch_size(self):
|
||||
batch_size = self.batch_size
|
||||
if batch_size is None:
|
||||
batch_size = self.default_batch_sizes["cpu"]
|
||||
if settings.TORCH_DEVICE_MODEL in self.default_batch_sizes:
|
||||
batch_size = self.default_batch_sizes[settings.TORCH_DEVICE_MODEL]
|
||||
return batch_size
|
||||
|
||||
def __call__(self, *args, **kwargs):
|
||||
raise NotImplementedError()
|
||||
@@ -0,0 +1,22 @@
|
||||
from typing import Optional
|
||||
|
||||
from transformers import PreTrainedModel
|
||||
from transformers.utils import is_flash_attn_2_available
|
||||
|
||||
|
||||
class SuryaPreTrainedModel(PreTrainedModel):
|
||||
# No-op if we pass attention, so we can set attention however we want in the config
|
||||
def _check_and_adjust_attn_implementation(
|
||||
self, attn_implementation: Optional[str], **kwargs
|
||||
):
|
||||
if attn_implementation is None:
|
||||
try:
|
||||
self._sdpa_can_dispatch(True)
|
||||
attn_implementation = "sdpa"
|
||||
except (ValueError, ImportError):
|
||||
attn_implementation = "eager"
|
||||
|
||||
if self._supports_flash_attn and is_flash_attn_2_available():
|
||||
attn_implementation = "flash_attention_2"
|
||||
|
||||
return attn_implementation
|
||||
@@ -0,0 +1,182 @@
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import tempfile
|
||||
import time
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from pathlib import Path
|
||||
|
||||
import requests
|
||||
from tqdm import tqdm
|
||||
|
||||
from surya.logging import get_logger
|
||||
from surya.settings import settings
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
# Lock file expiration time in seconds (10 minutes)
|
||||
LOCK_EXPIRATION = 600
|
||||
|
||||
|
||||
def join_urls(url1: str, url2: str):
|
||||
url1 = url1.rstrip("/")
|
||||
url2 = url2.lstrip("/")
|
||||
return f"{url1}/{url2}"
|
||||
|
||||
|
||||
def get_model_name(pretrained_model_name_or_path: str):
|
||||
return pretrained_model_name_or_path.split("/")[0]
|
||||
|
||||
|
||||
def download_file(remote_path: str, local_path: str, chunk_size: int = 1024 * 1024):
|
||||
local_path = Path(local_path)
|
||||
try:
|
||||
response = requests.get(remote_path, stream=True, allow_redirects=True)
|
||||
response.raise_for_status() # Raise an exception for bad status codes
|
||||
|
||||
# Get file size from headers for progress bar
|
||||
total_size = int(response.headers.get('content-length', 0))
|
||||
|
||||
# Create progress bar with file name and size info
|
||||
filename = local_path.name
|
||||
pbar = tqdm(
|
||||
total=total_size,
|
||||
unit='B',
|
||||
unit_scale=True,
|
||||
unit_divisor=1024,
|
||||
desc=f"Downloading {filename}",
|
||||
miniters=1
|
||||
)
|
||||
|
||||
with open(local_path, "wb") as f:
|
||||
downloaded = 0
|
||||
for chunk in response.iter_content(chunk_size=chunk_size):
|
||||
if chunk:
|
||||
f.write(chunk)
|
||||
downloaded += len(chunk)
|
||||
pbar.update(len(chunk))
|
||||
|
||||
pbar.close()
|
||||
return local_path
|
||||
except Exception as e:
|
||||
if local_path.exists():
|
||||
local_path.unlink()
|
||||
logger.error(f"Download error for file {remote_path}: {str(e)}")
|
||||
raise
|
||||
|
||||
|
||||
def check_manifest(local_dir: str):
|
||||
local_dir = Path(local_dir)
|
||||
manifest_path = local_dir / "manifest.json"
|
||||
if not os.path.exists(manifest_path):
|
||||
return False
|
||||
|
||||
try:
|
||||
with open(manifest_path, "r") as f:
|
||||
manifest = json.load(f)
|
||||
for file in manifest["files"]:
|
||||
if not os.path.exists(local_dir / file):
|
||||
return False
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
def download_directory(remote_path: str, local_dir: str):
|
||||
model_name = get_model_name(remote_path)
|
||||
s3_url = join_urls(settings.S3_BASE_URL, remote_path)
|
||||
# Check to see if it's already downloaded
|
||||
model_exists = check_manifest(local_dir)
|
||||
if model_exists:
|
||||
return
|
||||
|
||||
# Use tempfile.TemporaryDirectory to automatically clean up
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
# Download the manifest file
|
||||
manifest_file = join_urls(s3_url, "manifest.json")
|
||||
manifest_path = os.path.join(temp_dir, "manifest.json")
|
||||
download_file(manifest_file, manifest_path)
|
||||
|
||||
# List and download all files
|
||||
with open(manifest_path, "r") as f:
|
||||
manifest = json.load(f)
|
||||
|
||||
pbar = tqdm(
|
||||
desc=f"Downloading {model_name} model to {local_dir}",
|
||||
total=len(manifest["files"]),
|
||||
)
|
||||
|
||||
with ThreadPoolExecutor(
|
||||
max_workers=settings.PARALLEL_DOWNLOAD_WORKERS
|
||||
) as executor:
|
||||
futures = []
|
||||
for file in manifest["files"]:
|
||||
remote_file = join_urls(s3_url, file)
|
||||
local_file = os.path.join(temp_dir, file)
|
||||
futures.append(executor.submit(download_file, remote_file, local_file))
|
||||
|
||||
for future in futures:
|
||||
future.result()
|
||||
pbar.update(1)
|
||||
|
||||
pbar.close()
|
||||
|
||||
# Move all files to new directory
|
||||
for file in os.listdir(temp_dir):
|
||||
shutil.move(os.path.join(temp_dir, file), local_dir)
|
||||
|
||||
|
||||
class S3DownloaderMixin:
|
||||
s3_prefix = "s3://"
|
||||
|
||||
@classmethod
|
||||
def get_local_path(cls, pretrained_model_name_or_path) -> str:
|
||||
if pretrained_model_name_or_path.startswith(cls.s3_prefix):
|
||||
pretrained_model_name_or_path = pretrained_model_name_or_path.replace(
|
||||
cls.s3_prefix, ""
|
||||
)
|
||||
cache_dir = settings.MODEL_CACHE_DIR
|
||||
local_path = os.path.join(cache_dir, pretrained_model_name_or_path)
|
||||
os.makedirs(local_path, exist_ok=True)
|
||||
else:
|
||||
local_path = ""
|
||||
return local_path
|
||||
|
||||
@classmethod
|
||||
def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
|
||||
# Allow loading models directly from the hub, or using s3
|
||||
if not pretrained_model_name_or_path.startswith(cls.s3_prefix):
|
||||
return super().from_pretrained(
|
||||
pretrained_model_name_or_path, *args, **kwargs
|
||||
)
|
||||
|
||||
local_path = cls.get_local_path(pretrained_model_name_or_path)
|
||||
pretrained_model_name_or_path = pretrained_model_name_or_path.replace(
|
||||
cls.s3_prefix, ""
|
||||
)
|
||||
|
||||
# Retry logic for downloading the model folder
|
||||
retries = 3
|
||||
delay = 5
|
||||
attempt = 0
|
||||
success = False
|
||||
while not success and attempt < retries:
|
||||
try:
|
||||
download_directory(pretrained_model_name_or_path, local_path)
|
||||
success = True # If download succeeded
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Error downloading model from {pretrained_model_name_or_path}. Attempt {attempt + 1} of {retries}. Error: {e}"
|
||||
)
|
||||
attempt += 1
|
||||
if attempt < retries:
|
||||
logger.info(f"Retrying in {delay} seconds...")
|
||||
time.sleep(delay) # Wait before retrying
|
||||
else:
|
||||
logger.error(
|
||||
f"Failed to download {pretrained_model_name_or_path} after {retries} attempts."
|
||||
)
|
||||
raise e # Reraise exception after max retries
|
||||
|
||||
return super().from_pretrained(local_path, *args, **kwargs)
|
||||
@@ -0,0 +1,44 @@
|
||||
from typing import List
|
||||
|
||||
from surya.common.polygon import PolygonBox
|
||||
|
||||
|
||||
def clean_boxes(boxes: List[PolygonBox]) -> List[PolygonBox]:
|
||||
new_boxes = []
|
||||
for box_obj in boxes:
|
||||
xs = [point[0] for point in box_obj.polygon]
|
||||
ys = [point[1] for point in box_obj.polygon]
|
||||
if max(xs) == min(xs) or max(ys) == min(ys):
|
||||
continue
|
||||
|
||||
box = box_obj.bbox
|
||||
contained = False
|
||||
for other_box_obj in boxes:
|
||||
if other_box_obj.polygon == box_obj.polygon:
|
||||
continue
|
||||
|
||||
other_box = other_box_obj.bbox
|
||||
if box == other_box:
|
||||
continue
|
||||
if (
|
||||
box[0] >= other_box[0]
|
||||
and box[1] >= other_box[1]
|
||||
and box[2] <= other_box[2]
|
||||
and box[3] <= other_box[3]
|
||||
):
|
||||
contained = True
|
||||
break
|
||||
if not contained:
|
||||
new_boxes.append(box_obj)
|
||||
return new_boxes
|
||||
|
||||
|
||||
def expand_bbox(bbox, expansion_factor=0.01):
|
||||
expansion_low = 1 - expansion_factor
|
||||
expansion_high = 1 + expansion_factor
|
||||
return [
|
||||
bbox[0] * expansion_low,
|
||||
bbox[1] * expansion_low,
|
||||
bbox[2] * expansion_high,
|
||||
bbox[3] * expansion_high,
|
||||
]
|
||||
Reference in New Issue
Block a user