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