Fu DaiandClaude Opus 4.8
1a585693be
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>
78 lines
2.7 KiB
Python
78 lines
2.7 KiB
Python
"""Capture per-page OCR text, bboxes, and latency from the vLLM OCR endpoint."""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import base64
|
|
import json
|
|
import sys
|
|
from pathlib import Path
|
|
from typing import Any, Dict
|
|
|
|
import requests
|
|
|
|
|
|
def extract_capture(body: Dict[str, Any]) -> Dict[str, Any]:
|
|
data = body.get("data") or {}
|
|
if not isinstance(data, dict):
|
|
data = {}
|
|
blocks = ((data.get("ocr_text_json") or {}).get("blocks")) or []
|
|
boxes = [list(b["bbox"]) for b in blocks if isinstance(b, dict) and b.get("bbox")]
|
|
return {
|
|
"text": data.get("text_lines", "") or "",
|
|
"boxes": boxes,
|
|
"elapsed_seconds": data.get("elapsed_seconds"),
|
|
}
|
|
|
|
|
|
def write_capture(cap: Dict[str, Any], out_dir: Path, stem: str) -> None:
|
|
out_dir.mkdir(parents=True, exist_ok=True)
|
|
(out_dir / f"{stem}.txt").write_text(cap["text"], encoding="utf-8")
|
|
(out_dir / f"{stem}.json").write_text(
|
|
json.dumps({"boxes": cap["boxes"], "elapsed_seconds": cap["elapsed_seconds"]}),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
|
|
def _payload(image_path: Path) -> Dict[str, Any]:
|
|
suffix = image_path.suffix.lstrip(".").lower() or "png"
|
|
return {
|
|
"file": base64.b64encode(image_path.read_bytes()).decode("utf-8"),
|
|
"type": "jpg" if suffix == "jpeg" else suffix,
|
|
"skip_text_detection": False,
|
|
"skip_table_detection": False,
|
|
"recognize_math": False,
|
|
"ocr_with_boxes": True,
|
|
}
|
|
|
|
|
|
def capture_page(url: str, image_path: Path, out_dir: Path, timeout: float) -> Dict[str, Any]:
|
|
resp = requests.post(url, json=_payload(image_path), timeout=timeout)
|
|
cap = extract_capture(resp.json())
|
|
if cap["elapsed_seconds"] is None:
|
|
# Latency is a headline metric; a successful response with no timing means
|
|
# the response contract changed. Surface it loudly instead of silently
|
|
# dropping the page from the latency stats.
|
|
print(
|
|
f"WARNING: no elapsed_seconds in OCR response for {image_path.name}; "
|
|
"latency for this page will be missing",
|
|
file=sys.stderr,
|
|
)
|
|
write_capture(cap, out_dir, image_path.stem)
|
|
return cap
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description="Capture OCR text+boxes+latency per image.")
|
|
parser.add_argument("images", nargs="+", type=Path)
|
|
parser.add_argument("--url", required=True)
|
|
parser.add_argument("--out-dir", type=Path, required=True)
|
|
parser.add_argument("--timeout", type=float, default=900)
|
|
args = parser.parse_args()
|
|
for image in args.images:
|
|
cap = capture_page(args.url, image, args.out_dir, args.timeout)
|
|
print(f"captured {image.name}: {len(cap['boxes'])} boxes, {cap['elapsed_seconds']}s")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|