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,242 @@
|
||||
import io
|
||||
import tempfile
|
||||
from typing import List
|
||||
|
||||
import pypdfium2
|
||||
|
||||
from surya.debug.draw import draw_polys_on_image, draw_bboxes_on_image
|
||||
|
||||
from PIL import Image
|
||||
from surya.settings import settings
|
||||
from vllm_tools import predictors_vllm, ocr_vllm
|
||||
|
||||
|
||||
predictors = predictors_vllm
|
||||
|
||||
|
||||
def rescale_bbox(bbox, source_size, target_size):
|
||||
width_ratio = target_size[0] / source_size[0]
|
||||
height_ratio = target_size[1] / source_size[1]
|
||||
return [
|
||||
bbox[0] * width_ratio,
|
||||
bbox[1] * height_ratio,
|
||||
bbox[2] * width_ratio,
|
||||
bbox[3] * height_ratio,
|
||||
]
|
||||
|
||||
|
||||
def expand_bbox(bbox, margin=5):
|
||||
return [
|
||||
max(0, int(bbox[0]) - margin),
|
||||
max(0, int(bbox[1]) - margin),
|
||||
int(bbox[2]) + margin,
|
||||
int(bbox[3]) + margin,
|
||||
]
|
||||
|
||||
|
||||
|
||||
def page_counter(pdf_file):
|
||||
doc = open_pdf(pdf_file)
|
||||
doc_len = len(doc)
|
||||
doc.close()
|
||||
return doc_len
|
||||
|
||||
def ocr_errors(pdf_file, page_count, sample_len=512, max_samples=10, max_pages=15):
|
||||
from pdftext.extraction import plain_text_output
|
||||
|
||||
with tempfile.NamedTemporaryFile(suffix=".pdf") as f:
|
||||
f.write(pdf_file.getvalue())
|
||||
f.seek(0)
|
||||
|
||||
# Sample the text from the middle of the PDF
|
||||
page_middle = page_count // 2
|
||||
page_range = range(
|
||||
max(page_middle - max_pages, 0), min(page_middle + max_pages, page_count)
|
||||
)
|
||||
text = plain_text_output(f.name, page_range=page_range)
|
||||
|
||||
sample_gap = len(text) // max_samples
|
||||
if len(text) == 0 or sample_gap == 0:
|
||||
return "This PDF has no text or very little text", ["no text"]
|
||||
|
||||
if sample_gap < sample_len:
|
||||
sample_gap = sample_len
|
||||
|
||||
# Split the text into samples for the model
|
||||
samples = []
|
||||
for i in range(0, len(text), sample_gap):
|
||||
samples.append(text[i : i + sample_len])
|
||||
|
||||
results = predictors["ocr_error"](samples)
|
||||
label = "This PDF has good text."
|
||||
if results.labels.count("bad") / len(results.labels) > 0.2:
|
||||
label = "This PDF may have garbled or bad OCR text."
|
||||
return label, results.labels
|
||||
|
||||
|
||||
def text_detection(img):
|
||||
text_pred = predictors["detection"]([img])[0]
|
||||
text_polygons = [p.polygon for p in text_pred.bboxes]
|
||||
det_img = draw_polys_on_image(text_polygons, img.copy())
|
||||
return det_img, text_pred
|
||||
|
||||
|
||||
def layout_detection(img):
|
||||
pred = predictors["layout"]([img])[0]
|
||||
polygons = [p.polygon for p in pred.bboxes]
|
||||
labels = [
|
||||
f"{p.label}-{p.position}-{round(p.top_k[p.label], 2)}" for p in pred.bboxes
|
||||
]
|
||||
layout_img = draw_polys_on_image(
|
||||
polygons, img.copy(), labels=labels, label_font_size=18
|
||||
)
|
||||
return layout_img, pred
|
||||
|
||||
|
||||
def table_recognition(
|
||||
img, highres_img, skip_table_detection: bool
|
||||
):
|
||||
if skip_table_detection:
|
||||
layout_tables = [(0, 0, highres_img.size[0], highres_img.size[1])]
|
||||
table_imgs = [highres_img]
|
||||
else:
|
||||
_, layout_pred = layout_detection(img)
|
||||
layout_tables_lowres = [
|
||||
line.bbox
|
||||
for line in layout_pred.bboxes
|
||||
if line.label in ["Table", "TableOfContents"]
|
||||
]
|
||||
table_imgs = []
|
||||
layout_tables = []
|
||||
for tb in layout_tables_lowres:
|
||||
highres_bbox = rescale_bbox(tb, img.size, highres_img.size)
|
||||
# Slightly expand the box
|
||||
highres_bbox = expand_bbox(highres_bbox)
|
||||
table_imgs.append(highres_img.crop(highres_bbox))
|
||||
layout_tables.append(highres_bbox)
|
||||
|
||||
table_preds = predictors["table_rec"](table_imgs)
|
||||
table_img = img.copy()
|
||||
|
||||
for results, table_bbox in zip(table_preds, layout_tables):
|
||||
adjusted_bboxes = []
|
||||
labels = []
|
||||
colors = []
|
||||
|
||||
for item in results.cells:
|
||||
adjusted_bboxes.append(
|
||||
[
|
||||
(item.bbox[0] + table_bbox[0]),
|
||||
(item.bbox[1] + table_bbox[1]),
|
||||
(item.bbox[2] + table_bbox[0]),
|
||||
(item.bbox[3] + table_bbox[1]),
|
||||
]
|
||||
)
|
||||
labels.append(item.label)
|
||||
if "Row" in item.label:
|
||||
colors.append("blue")
|
||||
else:
|
||||
colors.append("red")
|
||||
table_img = draw_bboxes_on_image(
|
||||
adjusted_bboxes,
|
||||
highres_img,
|
||||
labels=labels,
|
||||
label_font_size=18,
|
||||
color=colors,
|
||||
)
|
||||
return table_img, table_preds
|
||||
|
||||
|
||||
# Function for OCR
|
||||
def ocr(
|
||||
img: Image.Image,
|
||||
highres_img: Image.Image,
|
||||
skip_text_detection: bool = False,
|
||||
recognize_math: bool = True,
|
||||
with_bboxes: bool = True,
|
||||
):
|
||||
return ocr_vllm(
|
||||
img,
|
||||
highres_img,
|
||||
skip_text_detection=skip_text_detection,
|
||||
recognize_math=recognize_math,
|
||||
with_bboxes=with_bboxes,
|
||||
)
|
||||
|
||||
|
||||
def open_pdf(pdf_file):
|
||||
stream = io.BytesIO(pdf_file.getvalue())
|
||||
return pypdfium2.PdfDocument(stream)
|
||||
|
||||
|
||||
def get_page_image(pdf_file, page_num, dpi=settings.IMAGE_DPI):
|
||||
doc = open_pdf(pdf_file)
|
||||
renderer = doc.render(
|
||||
pypdfium2.PdfBitmap.to_pil,
|
||||
page_indices=[page_num - 1],
|
||||
scale=dpi / 72,
|
||||
)
|
||||
png = list(renderer)[0]
|
||||
png_image = png.convert("RGB")
|
||||
doc.close()
|
||||
return png_image
|
||||
|
||||
|
||||
def page_counter(pdf_file):
|
||||
doc = open_pdf(pdf_file)
|
||||
doc_len = len(doc)
|
||||
doc.close()
|
||||
return doc_len
|
||||
|
||||
|
||||
import pandas as pd
|
||||
|
||||
def bbox_intersection(box1, box2):
|
||||
x1 = max(box1[0], box2[0])
|
||||
y1 = max(box1[1], box2[1])
|
||||
x2 = min(box1[2], box2[2])
|
||||
y2 = min(box1[3], box2[3])
|
||||
|
||||
if x1 < x2 and y1 < y2:
|
||||
return (x1, y1, x2, y2)
|
||||
else:
|
||||
return None
|
||||
|
||||
def area_of_bbox(box):
|
||||
return (box[2] - box[0]) * (box[3] - box[1])
|
||||
|
||||
def is_bbox_inside(box, parent_box):
|
||||
interaction_box = bbox_intersection(box, parent_box)
|
||||
if interaction_box is None:
|
||||
return False
|
||||
return area_of_bbox(interaction_box) / area_of_bbox(box) > 0.5
|
||||
|
||||
def center_of_bbox(box):
|
||||
return ((box[0] + box[2]) / 2, (box[1] + box[3]) / 2)
|
||||
|
||||
|
||||
def extract_text_from_image(image):
|
||||
layout_predictor = predictors["layout"]
|
||||
recognition_predictor = predictors["recognition"]
|
||||
|
||||
layout_prediction = layout_predictor([image])[0]
|
||||
prediction = recognition_predictor([image], [layout_prediction], full_page=False)[0]
|
||||
|
||||
items = [
|
||||
{
|
||||
"text": block.html,
|
||||
"position": block.reading_order,
|
||||
"order_value": block.bbox[1],
|
||||
}
|
||||
for block in prediction.blocks
|
||||
if not block.skipped and not block.error and block.html
|
||||
]
|
||||
if not items:
|
||||
return ""
|
||||
|
||||
df = pd.DataFrame(items)
|
||||
df = df.sort_values(by=['position', 'order_value'])
|
||||
ds = df.groupby('position').apply(lambda x: " ".join(x['text'].tolist()))
|
||||
full_text = "\n\n".join(ds.to_list())
|
||||
|
||||
return full_text
|
||||
Reference in New Issue
Block a user