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>
243 lines
6.6 KiB
Python
243 lines
6.6 KiB
Python
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
|