layout_match_processor.py 1.5 KB

1234567891011121314151617181920212223242526272829303132333435363738394041
  1. import sys
  2. import math
  3. from para.commons import *
  4. if sys.version_info[0] >= 3:
  5. sys.stdout.reconfigure(encoding="utf-8") # type: ignore
  6. class LayoutFilterProcessor:
  7. def __init__(self) -> None:
  8. pass
  9. def batch_process_blocks(self, pdf_dict):
  10. for page_id, blocks in pdf_dict.items():
  11. if page_id.startswith("page_"):
  12. if "layout_bboxes" in blocks.keys() and "para_blocks" in blocks.keys():
  13. layout_bbox_objs = blocks["layout_bboxes"]
  14. if layout_bbox_objs is None:
  15. continue
  16. layout_bboxes = [bbox_obj["layout_bbox"] for bbox_obj in layout_bbox_objs]
  17. # Use math.ceil function to enlarge each value of x0, y0, x1, y1 of each layout_bbox
  18. layout_bboxes = [
  19. [math.ceil(x0), math.ceil(y0), math.ceil(x1), math.ceil(y1)] for x0, y0, x1, y1 in layout_bboxes
  20. ]
  21. para_blocks = blocks["para_blocks"]
  22. if para_blocks is None:
  23. continue
  24. for lb_bbox in layout_bboxes:
  25. for i, para_block in enumerate(para_blocks):
  26. para_bbox = para_block["bbox"]
  27. para_blocks[i]["in_layout"] = 0
  28. if is_in_bbox(para_bbox, lb_bbox):
  29. para_blocks[i]["in_layout"] = 1
  30. blocks["para_blocks"] = para_blocks
  31. return pdf_dict