ocr_detect_all_bboxes.py 1.6 KB

1234567891011121314151617181920212223242526272829303132333435
  1. from magic_pdf.libs.ocr_content_type import BlockType
  2. def ocr_prepare_bboxes_for_layout_split(img_blocks, table_blocks, discarded_blocks, text_blocks,
  3. title_blocks, interline_equation_blocks, page_w, page_h):
  4. all_bboxes = []
  5. for image in img_blocks:
  6. x0, y0, x1, y1 = image['bbox']
  7. all_bboxes.append([x0, y0, x1, y1, None, None, None, BlockType.Image, None, None, None, None])
  8. for table in table_blocks:
  9. x0, y0, x1, y1 = table['bbox']
  10. all_bboxes.append([x0, y0, x1, y1, None, None, None, BlockType.Table, None, None, None, None])
  11. for text in text_blocks:
  12. x0, y0, x1, y1 = text['bbox']
  13. all_bboxes.append([x0, y0, x1, y1, None, None, None, BlockType.Text, None, None, None, None])
  14. for title in title_blocks:
  15. x0, y0, x1, y1 = title['bbox']
  16. all_bboxes.append([x0, y0, x1, y1, None, None, None, BlockType.Title, None, None, None, None])
  17. for interline_equation in interline_equation_blocks:
  18. x0, y0, x1, y1 = interline_equation['bbox']
  19. all_bboxes.append([x0, y0, x1, y1, None, None, None, BlockType.InterlineEquation, None, None, None, None])
  20. '''discarded_blocks中只保留宽度超过1/3页面宽度的,高度超过10的,处于页面下半50%区域的(限定footnote)'''
  21. for discarded in discarded_blocks:
  22. x0, y0, x1, y1 = discarded['bbox']
  23. if (x1 - x0) > (page_w / 3) and (y1 - y0) > 10 and y0 > (page_h / 2):
  24. all_bboxes.append([x0, y0, x1, y1, None, None, None, BlockType.Footnote, None, None, None, None])
  25. return all_bboxes