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,97 @@
|
||||
import json
|
||||
import logging
|
||||
import logging.config
|
||||
import time
|
||||
import uuid
|
||||
|
||||
import yaml
|
||||
from fastapi import FastAPI, Request
|
||||
from fastapi.exceptions import RequestValidationError
|
||||
from fastapi.responses import JSONResponse
|
||||
|
||||
from surya.endpoint.legacy import router as legacy_router
|
||||
from surya.endpoint.openai import router as openai_router
|
||||
from surya.endpoint.service import _error_response
|
||||
|
||||
with open("logger.yaml", "r", encoding="utf-8") as f:
|
||||
config = yaml.safe_load(f.read())
|
||||
logging.config.dictConfig(config)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
app = FastAPI(
|
||||
title="OCR SERVICE API SERVICE",
|
||||
version="1.0",
|
||||
docs_url="/v1/api/ai/swagger",
|
||||
openapi_url="/v1/api/ai/openapi.json",
|
||||
)
|
||||
|
||||
|
||||
@app.exception_handler(RequestValidationError)
|
||||
async def validation_exception_handler(request: Request, exc: RequestValidationError):
|
||||
request_id = getattr(request.state, "request_id", "-")
|
||||
logger.warning(
|
||||
"request_validation_failed request_id=%s method=%s path=%s errors=%s",
|
||||
request_id,
|
||||
request.method,
|
||||
request.url.path,
|
||||
exc.errors(),
|
||||
)
|
||||
return JSONResponse(
|
||||
status_code=200,
|
||||
content={
|
||||
"data": [],
|
||||
"message": json.dumps({"error info": exc.errors()}, default=str),
|
||||
"code": 422,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@app.middleware("http")
|
||||
async def allow_openapi_unauthorized(request: Request, call_next):
|
||||
if request.url.path == "/v1/api/ai/openapi.json":
|
||||
response = await call_next(request)
|
||||
return response
|
||||
return await call_next(request)
|
||||
|
||||
|
||||
@app.middleware("http")
|
||||
async def log_request_response(request, call_next):
|
||||
request_id = uuid.uuid4().hex
|
||||
request.state.request_id = request_id
|
||||
start = time.perf_counter()
|
||||
client = request.client.host if request.client else "-"
|
||||
logger.info(
|
||||
"request_start request_id=%s method=%s path=%s client=%s",
|
||||
request_id,
|
||||
request.method,
|
||||
request.url.path,
|
||||
client,
|
||||
)
|
||||
try:
|
||||
response = await call_next(request)
|
||||
except Exception as e:
|
||||
duration_ms = (time.perf_counter() - start) * 1000
|
||||
logger.exception(
|
||||
"request_unhandled_exception request_id=%s method=%s path=%s duration_ms=%.2f",
|
||||
request_id,
|
||||
request.method,
|
||||
request.url.path,
|
||||
duration_ms,
|
||||
exc_info=True,
|
||||
)
|
||||
return JSONResponse(status_code=500, content=_error_response(e))
|
||||
duration_ms = (time.perf_counter() - start) * 1000
|
||||
logger.info(
|
||||
"request_end request_id=%s method=%s path=%s status_code=%s duration_ms=%.2f",
|
||||
request_id,
|
||||
request.method,
|
||||
request.url.path,
|
||||
response.status_code,
|
||||
duration_ms,
|
||||
)
|
||||
return response
|
||||
|
||||
|
||||
app.include_router(legacy_router)
|
||||
app.include_router(openai_router)
|
||||
@@ -0,0 +1,157 @@
|
||||
import io
|
||||
import time
|
||||
|
||||
from fastapi import APIRouter, File, Request, UploadFile
|
||||
from PIL import Image
|
||||
|
||||
from tools import (
|
||||
ocr,
|
||||
text_detection,
|
||||
layout_detection,
|
||||
table_recognition,
|
||||
extract_text_from_image,
|
||||
)
|
||||
from vllm_batcher import vllm_ocr_batcher
|
||||
from vllm_tools import vllm_backend_info
|
||||
from surya.endpoint.schemas import ApiResponse, Info
|
||||
from surya.endpoint.service import (
|
||||
_dump_model_or_list,
|
||||
_error_response,
|
||||
_load_base64_image,
|
||||
_log_exception,
|
||||
_ocr_response_data,
|
||||
_success_response,
|
||||
logger,
|
||||
)
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
@router.post("/home")
|
||||
def home():
|
||||
return "<h1>Welcome to SURYA OCR API!</h1>"
|
||||
|
||||
|
||||
@router.post("/v1/api/ai/suya_ocr", include_in_schema=False, response_model=ApiResponse)
|
||||
@router.post("/v1/api/ai/suya_ocr/", response_model=ApiResponse)
|
||||
def run_ocr(p: Info, request: Request):
|
||||
try:
|
||||
pil_image, pil_image_highres = _load_base64_image(p)
|
||||
rec_img, pred, box_img = ocr(
|
||||
pil_image,
|
||||
pil_image_highres,
|
||||
p.skip_text_detection,
|
||||
p.recognize_math,
|
||||
with_bboxes=p.ocr_with_boxes,
|
||||
)
|
||||
return _success_response(_ocr_response_data(pred))
|
||||
except Exception as e:
|
||||
results = _error_response(e)
|
||||
_log_exception("suya_ocr", request, e, results)
|
||||
return results
|
||||
|
||||
|
||||
@router.get("/v1/api/ai/suya_ocr_vllm/health", response_model=ApiResponse)
|
||||
async def suya_ocr_vllm_health():
|
||||
return _success_response(vllm_backend_info())
|
||||
|
||||
|
||||
@router.post("/v1/api/ai/suya_ocr_vllm", include_in_schema=False, response_model=ApiResponse)
|
||||
@router.post("/v1/api/ai/suya_ocr_vllm/", response_model=ApiResponse)
|
||||
def run_ocr_vllm(p: Info, request: Request):
|
||||
try:
|
||||
pil_image, pil_image_highres = _load_base64_image(p)
|
||||
start = time.perf_counter()
|
||||
_, pred, _ = vllm_ocr_batcher.submit(
|
||||
pil_image,
|
||||
pil_image_highres,
|
||||
skip_text_detection=p.skip_text_detection,
|
||||
recognize_math=p.recognize_math,
|
||||
with_bboxes=False,
|
||||
request_id=getattr(request.state, "request_id", "-"),
|
||||
)
|
||||
logger.info(
|
||||
"api_vllm_submit_wait_complete request_id=%s duration_ms=%.2f",
|
||||
getattr(request.state, "request_id", "-"),
|
||||
(time.perf_counter() - start) * 1000,
|
||||
)
|
||||
return _success_response(_ocr_response_data(pred))
|
||||
except Exception as e:
|
||||
results = _error_response(e)
|
||||
_log_exception("suya_ocr_vllm", request, e, results)
|
||||
return results
|
||||
|
||||
|
||||
@router.post("/v1/api/ai/suya_text_det", include_in_schema=False, response_model=ApiResponse)
|
||||
@router.post("/v1/api/ai/suya_text_det/", response_model=ApiResponse)
|
||||
def run_text_det(p: Info, request: Request):
|
||||
try:
|
||||
pil_image, _ = _load_base64_image(p)
|
||||
det_img, text_pred = text_detection(pil_image)
|
||||
text_lines = text_pred.model_dump(exclude=["heatmap", "affinity_map"])
|
||||
return _success_response({"text_lines": text_lines})
|
||||
except Exception as e:
|
||||
results = _error_response(e)
|
||||
_log_exception("suya_text_det", request, e, results)
|
||||
return results
|
||||
|
||||
|
||||
@router.post("/v1/api/ai/suya_layout_det", include_in_schema=False, response_model=ApiResponse)
|
||||
@router.post("/v1/api/ai/suya_layout_det/", response_model=ApiResponse)
|
||||
def run_layout_det(p: Info, request: Request):
|
||||
try:
|
||||
pil_image, _ = _load_base64_image(p)
|
||||
layout_img, pred = layout_detection(pil_image)
|
||||
text_lines = pred.model_dump(exclude=["segmentation_map"])
|
||||
return _success_response({"text_lines": text_lines})
|
||||
except Exception as e:
|
||||
results = _error_response(e)
|
||||
_log_exception("suya_layout_det", request, e, results)
|
||||
return results
|
||||
|
||||
|
||||
@router.post("/v1/api/ai/suya_table_rec", include_in_schema=False, response_model=ApiResponse)
|
||||
@router.post("/v1/api/ai/suya_table_rec/", response_model=ApiResponse)
|
||||
def run_table_rec(p: Info, request: Request):
|
||||
try:
|
||||
pil_image, pil_image_highres = _load_base64_image(p)
|
||||
|
||||
table_img, pred = table_recognition(
|
||||
pil_image, pil_image_highres, p.skip_table_detection
|
||||
)
|
||||
|
||||
text_json = _dump_model_or_list(pred)
|
||||
text_lines = ""
|
||||
if isinstance(pred, list):
|
||||
text_lines = "\n".join(
|
||||
[item.html for item in pred if getattr(item, "html", None)]
|
||||
)
|
||||
elif hasattr(pred, "text_lines"):
|
||||
text_lines = "\n".join([p.text for p in pred.text_lines])
|
||||
|
||||
return _success_response({"ocr_text_json": text_json, "text_lines": text_lines})
|
||||
except Exception as e:
|
||||
results = _error_response(e)
|
||||
_log_exception("suya_table_rec", request, e, results)
|
||||
return results
|
||||
|
||||
|
||||
@router.post("/image2text", response_model=ApiResponse)
|
||||
async def image_to_text(request: Request, file: UploadFile = File(...)):
|
||||
try:
|
||||
contents = await file.read()
|
||||
image = Image.open(io.BytesIO(contents))
|
||||
logger.info(
|
||||
"image2text_upload request_id=%s filename=%s content_type=%s size_bytes=%s image_size=%s",
|
||||
getattr(request.state, "request_id", "-"),
|
||||
file.filename,
|
||||
file.content_type,
|
||||
len(contents),
|
||||
image.size,
|
||||
)
|
||||
full_text = extract_text_from_image(image)
|
||||
return _success_response(full_text)
|
||||
except Exception as e:
|
||||
results = _error_response(e)
|
||||
_log_exception("image2text", request, e, results)
|
||||
return results
|
||||
@@ -0,0 +1,108 @@
|
||||
import time
|
||||
import uuid
|
||||
from typing import Any, Dict, List, Literal, Optional, Union
|
||||
|
||||
from fastapi import APIRouter, Request
|
||||
from fastapi.responses import JSONResponse
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
|
||||
from surya.endpoint import service
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
class _ImageUrl(BaseModel):
|
||||
model_config = ConfigDict(extra="allow")
|
||||
url: str
|
||||
|
||||
|
||||
class _ContentPart(BaseModel):
|
||||
model_config = ConfigDict(extra="allow")
|
||||
type: str
|
||||
image_url: Optional[_ImageUrl] = None
|
||||
text: Optional[str] = None
|
||||
|
||||
|
||||
class _ChatMessage(BaseModel):
|
||||
model_config = ConfigDict(extra="allow")
|
||||
role: str
|
||||
content: Union[str, List[_ContentPart]]
|
||||
|
||||
|
||||
class ChatCompletionRequest(BaseModel):
|
||||
# The OpenAI SDK merges `extra_body={...}` into the TOP LEVEL of the request
|
||||
# JSON, so the OCR control flags live here rather than nested under a key.
|
||||
model_config = ConfigDict(extra="allow")
|
||||
model: str = "surya-ocr"
|
||||
messages: List[_ChatMessage]
|
||||
stream: Optional[bool] = False
|
||||
skip_text_detection: bool = False
|
||||
recognize_math: bool = False
|
||||
skip_table_detection: bool = False
|
||||
ocr_with_boxes: bool = True
|
||||
mode: Literal["block", "full_page", "table"] = "block"
|
||||
|
||||
|
||||
def _openai_error(message: str, err_type: str, status_code: int) -> JSONResponse:
|
||||
return JSONResponse(
|
||||
status_code=status_code,
|
||||
content={"error": {"message": message, "type": err_type, "code": None}},
|
||||
)
|
||||
|
||||
|
||||
def _build_chat_completion(
|
||||
content: str, model: str, ocr_json: Any, elapsed: Optional[float]
|
||||
) -> Dict[str, Any]:
|
||||
# Token accounting happens inside vLLM and is not surfaced to this layer, so
|
||||
# usage counts are reported as 0 (prompt_tokens intentionally 0).
|
||||
return {
|
||||
"id": f"chatcmpl-{uuid.uuid4().hex}",
|
||||
"object": "chat.completion",
|
||||
"created": int(time.time()),
|
||||
"model": model,
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"finish_reason": "stop",
|
||||
"message": {"role": "assistant", "content": content},
|
||||
}
|
||||
],
|
||||
"usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
|
||||
"surya": {"ocr_text_json": ocr_json, "elapsed_seconds": elapsed},
|
||||
}
|
||||
|
||||
|
||||
@router.post("/v1/chat/completions")
|
||||
def chat_completions(req: ChatCompletionRequest, request: Request):
|
||||
request_id = getattr(request.state, "request_id", "-")
|
||||
if req.stream:
|
||||
return _openai_error(
|
||||
"streaming is not supported by this OCR endpoint",
|
||||
"invalid_request_error",
|
||||
400,
|
||||
)
|
||||
try:
|
||||
image = service.extract_last_image(req.messages)
|
||||
except service.OcrRequestError as exc:
|
||||
return _openai_error(str(exc), "invalid_request_error", 400)
|
||||
|
||||
try:
|
||||
if req.mode == "table":
|
||||
data = service.table_image(image, req.skip_table_detection)
|
||||
else:
|
||||
skip = req.skip_text_detection or req.mode == "full_page"
|
||||
data = service.ocr_via_batcher(
|
||||
image,
|
||||
skip_text_detection=skip,
|
||||
recognize_math=req.recognize_math,
|
||||
request_id=request_id,
|
||||
)
|
||||
except service.OcrRequestError as exc:
|
||||
return _openai_error(str(exc), "invalid_request_error", 400)
|
||||
except Exception as exc:
|
||||
service._log_exception("chat_completions", request, exc, {})
|
||||
return _openai_error(str(exc), "internal_error", 500)
|
||||
|
||||
content = data.get("text_lines", "") or ""
|
||||
ocr_json = data.get("ocr_text_json") if req.ocr_with_boxes else None
|
||||
return _build_chat_completion(content, req.model, ocr_json, data.get("elapsed_seconds"))
|
||||
@@ -0,0 +1,46 @@
|
||||
import base64
|
||||
import binascii
|
||||
from typing import Any, Literal, Optional
|
||||
|
||||
from pydantic import (
|
||||
BaseModel,
|
||||
Field,
|
||||
PrivateAttr,
|
||||
StrictBool,
|
||||
field_validator,
|
||||
model_validator,
|
||||
)
|
||||
|
||||
|
||||
class Info(BaseModel):
|
||||
_decoded_file: bytes = PrivateAttr(default=b"")
|
||||
|
||||
file: str # base64 string
|
||||
type: Literal["png", "jpg", "jpeg", "gif"] # image type like 'png', 'jpg'
|
||||
skip_text_detection: Optional[StrictBool] = Field(False)
|
||||
skip_table_detection: Optional[StrictBool] = Field(False)
|
||||
recognize_math: Optional[StrictBool] = Field(False)
|
||||
ocr_with_boxes: Optional[StrictBool] = Field(True)
|
||||
|
||||
@field_validator("type", mode="before")
|
||||
@classmethod
|
||||
def normalize_type(cls, value: str) -> str:
|
||||
if not isinstance(value, str):
|
||||
raise ValueError("type must be a string")
|
||||
return value.lower()
|
||||
|
||||
@model_validator(mode="after")
|
||||
def decode_base64_file(self):
|
||||
if not self.file:
|
||||
raise ValueError("file must be a non-empty base64 string")
|
||||
try:
|
||||
self._decoded_file = base64.b64decode(self.file, validate=True)
|
||||
except binascii.Error as exc:
|
||||
raise ValueError("file must be valid base64") from exc
|
||||
return self
|
||||
|
||||
|
||||
class ApiResponse(BaseModel):
|
||||
data: Any
|
||||
message: str
|
||||
code: int
|
||||
@@ -0,0 +1,166 @@
|
||||
import base64
|
||||
import binascii
|
||||
import io
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
import traceback
|
||||
from typing import Any, Dict
|
||||
|
||||
import requests
|
||||
from fastapi import Request
|
||||
from PIL import Image
|
||||
|
||||
import tools
|
||||
from vllm_batcher import vllm_ocr_batcher
|
||||
from vllm_tools import page_ocr_to_response
|
||||
|
||||
logger = logging.getLogger("surya.endpoint")
|
||||
|
||||
|
||||
class OcrRequestError(Exception):
|
||||
"""Client-side error in an OCR request. Maps to HTTP 400."""
|
||||
|
||||
|
||||
# ---- helpers moved verbatim from api.py ------------------------------------
|
||||
|
||||
def _load_base64_image(p, copy_highres: bool = False):
|
||||
start = time.perf_counter()
|
||||
filetype = p.type.lower()
|
||||
allowed_types = ["png", "jpg", "jpeg", "gif"]
|
||||
if filetype not in allowed_types:
|
||||
raise ValueError("Unsupported file type")
|
||||
pil_image = Image.open(io.BytesIO(p._decoded_file)).convert("RGB")
|
||||
logger.info(
|
||||
"api_decode_image duration_ms=%.2f size_bytes=%s image_size=%s copy_highres=%s",
|
||||
(time.perf_counter() - start) * 1000,
|
||||
len(p._decoded_file),
|
||||
pil_image.size,
|
||||
copy_highres,
|
||||
)
|
||||
return pil_image, pil_image.copy() if copy_highres else pil_image
|
||||
|
||||
|
||||
def _error_response(e: Exception):
|
||||
tb = traceback.extract_tb(e.__traceback__)
|
||||
frame = tb[-1] if tb else None
|
||||
return {
|
||||
"data": [],
|
||||
"message": json.dumps(
|
||||
{
|
||||
"error info": str(e),
|
||||
"error at": frame.filename if frame else "unknown",
|
||||
"errort line at": frame.lineno if frame else 0,
|
||||
}
|
||||
),
|
||||
"code": 500,
|
||||
}
|
||||
|
||||
|
||||
def _success_response(data: Any) -> Dict[str, Any]:
|
||||
return {"data": data, "message": "success", "code": 200}
|
||||
|
||||
|
||||
def _log_exception(endpoint: str, request: Request, e: Exception, results: Dict[str, Any]):
|
||||
request_id = getattr(request.state, "request_id", "-")
|
||||
logger.exception(
|
||||
"endpoint_failed request_id=%s endpoint=%s code=%s message=%s",
|
||||
request_id,
|
||||
endpoint,
|
||||
results.get("code"),
|
||||
results.get("message"),
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
|
||||
def _ocr_response_data(pred):
|
||||
if hasattr(pred, "text_lines"):
|
||||
text_json = pred.model_dump()
|
||||
text_lines = "\n".join([p.text for p in pred.text_lines])
|
||||
return {"ocr_text_json": text_json, "text_lines": text_lines}
|
||||
return page_ocr_to_response(pred)
|
||||
|
||||
|
||||
def _dump_model_or_list(value):
|
||||
if hasattr(value, "model_dump"):
|
||||
return value.model_dump()
|
||||
if isinstance(value, list):
|
||||
return [
|
||||
item.model_dump() if hasattr(item, "model_dump") else item
|
||||
for item in value
|
||||
]
|
||||
return value
|
||||
|
||||
|
||||
# ---- new helpers for the OpenAI endpoint -----------------------------------
|
||||
|
||||
def load_image_from_url(url: str) -> Image.Image:
|
||||
"""Decode an OpenAI `image_url` (data: base64 URL or http(s): URL) to RGB."""
|
||||
if url.startswith("data:"):
|
||||
header, _, b64 = url.partition(",")
|
||||
if "base64" not in header or not b64:
|
||||
raise OcrRequestError("image_url data URL must be base64-encoded")
|
||||
try:
|
||||
raw = base64.b64decode(b64, validate=True)
|
||||
except binascii.Error as exc:
|
||||
raise OcrRequestError("image_url contains invalid base64") from exc
|
||||
elif url.startswith("http://") or url.startswith("https://"):
|
||||
try:
|
||||
resp = requests.get(url, timeout=10)
|
||||
resp.raise_for_status()
|
||||
except Exception as exc:
|
||||
raise OcrRequestError(f"could not fetch image_url: {exc}") from exc
|
||||
raw = resp.content
|
||||
else:
|
||||
raise OcrRequestError("image_url must be a data: or http(s): URL")
|
||||
try:
|
||||
return Image.open(io.BytesIO(raw)).convert("RGB")
|
||||
except Exception as exc:
|
||||
raise OcrRequestError(f"could not decode image: {exc}") from exc
|
||||
|
||||
|
||||
def extract_last_image(messages) -> Image.Image:
|
||||
"""Return the last image found across the chat messages' content parts."""
|
||||
url = None
|
||||
for msg in messages:
|
||||
content = msg.content
|
||||
if isinstance(content, list):
|
||||
for part in content:
|
||||
if getattr(part, "type", None) == "image_url" and part.image_url:
|
||||
url = part.image_url.url
|
||||
if url is None:
|
||||
raise OcrRequestError("no image_url content part found in messages")
|
||||
return load_image_from_url(url)
|
||||
|
||||
|
||||
def ocr_via_batcher(
|
||||
image: Image.Image,
|
||||
*,
|
||||
skip_text_detection: bool,
|
||||
recognize_math: bool,
|
||||
request_id: str,
|
||||
) -> Dict[str, Any]:
|
||||
"""Run page OCR through the shared request batcher. The annotated image is
|
||||
discarded, so we request with_bboxes=False to maximise batch coalescing with
|
||||
the primary /suya_ocr_vllm endpoint."""
|
||||
_, pred, _ = vllm_ocr_batcher.submit(
|
||||
image,
|
||||
image,
|
||||
skip_text_detection=skip_text_detection,
|
||||
recognize_math=recognize_math,
|
||||
with_bboxes=False,
|
||||
request_id=request_id,
|
||||
)
|
||||
return _ocr_response_data(pred)
|
||||
|
||||
|
||||
def table_image(image: Image.Image, skip_table_detection: bool) -> Dict[str, Any]:
|
||||
"""Run table-structure recognition; mirrors the /suya_table_rec response shape."""
|
||||
_, pred = tools.table_recognition(image, image.copy(), skip_table_detection)
|
||||
text_json = _dump_model_or_list(pred)
|
||||
text_lines = ""
|
||||
if isinstance(pred, list):
|
||||
text_lines = "\n".join([item.html for item in pred if getattr(item, "html", None)])
|
||||
elif hasattr(pred, "text_lines"):
|
||||
text_lines = "\n".join([line.text for line in pred.text_lines])
|
||||
return {"ocr_text_json": text_json, "text_lines": text_lines, "elapsed_seconds": None}
|
||||
Reference in New Issue
Block a user