| from typing import List | |
| from pypdfium2 import PdfDocument | |
| from surya.detection import batch_text_detection | |
| from marker.pdf.images import render_image | |
| from marker.schema.page import Page | |
| from marker.settings import settings | |
| def get_batch_size(): | |
| if settings.DETECTOR_BATCH_SIZE is not None: | |
| return settings.DETECTOR_BATCH_SIZE | |
| elif settings.TORCH_DEVICE_MODEL == "cuda": | |
| return 4 | |
| return 4 | |
| def surya_detection(doc: PdfDocument, pages: List[Page], det_model, batch_multiplier=1): | |
| processor = det_model.processor | |
| max_len = min(len(pages), len(doc)) | |
| images = [render_image(doc[pnum], dpi=settings.SURYA_DETECTOR_DPI) for pnum in range(max_len)] | |
| predictions = batch_text_detection(images, det_model, processor, batch_size=get_batch_size() * batch_multiplier) | |
| for (page, pred) in zip(pages, predictions): | |
| page.text_lines = pred | |