File size: 1,928 Bytes
c8a32e7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 |
from typing import List
from surya.layout import batch_layout_detection
from marker.pdf.images import render_image
from marker.schema.bbox import rescale_bbox
from marker.schema.page import Page
from marker.settings import settings
def get_batch_size():
if settings.LAYOUT_BATCH_SIZE is not None:
return settings.LAYOUT_BATCH_SIZE
elif settings.TORCH_DEVICE_MODEL == "cuda":
return 6
return 6
def surya_layout(doc, pages: List[Page], layout_model, batch_multiplier=1):
images = [render_image(doc[pnum], dpi=settings.SURYA_LAYOUT_DPI) for pnum in range(len(pages))]
text_detection_results = [p.text_lines for p in pages]
processor = layout_model.processor
layout_results = batch_layout_detection(images, layout_model, processor, detection_results=text_detection_results, batch_size=get_batch_size() * batch_multiplier)
for page, layout_result in zip(pages, layout_results):
page.layout = layout_result
def annotate_block_types(pages: List[Page]):
for page in pages:
max_intersections = {}
for i, block in enumerate(page.blocks):
for j, layout_block in enumerate(page.layout.bboxes):
layout_bbox = layout_block.bbox
layout_bbox = rescale_bbox(page.layout.image_bbox, page.bbox, layout_bbox)
intersection_pct = block.intersection_pct(layout_bbox)
if i not in max_intersections:
max_intersections[i] = (intersection_pct, j)
elif intersection_pct > max_intersections[i][0]:
max_intersections[i] = (intersection_pct, j)
for i, block in enumerate(page.blocks):
block = page.blocks[i]
block_type = "Text"
if i in max_intersections:
j = max_intersections[i][1]
block_type = page.layout.bboxes[j].label
block.block_type = block_type |