surya-ocr/surya/endpoint/service.py

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}