cut_image.py 2.2 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152
  1. from magic_pdf.libs.commons import join_path
  2. from magic_pdf.libs.ocr_content_type import ContentType
  3. from magic_pdf.libs.pdf_image_tools import cut_image
  4. def ocr_cut_image_and_table(spans, page, page_id, pdf_bytes_md5, imageWriter):
  5. def return_path(type):
  6. return join_path(pdf_bytes_md5, type)
  7. for span in spans:
  8. span_type = span['type']
  9. if span_type == ContentType.Image:
  10. span['image_path'] = cut_image(span['bbox'], page_id, page, return_path=return_path('images'),
  11. imageWriter=imageWriter)
  12. elif span_type == ContentType.Table:
  13. span['image_path'] = cut_image(span['bbox'], page_id, page, return_path=return_path('tables'),
  14. imageWriter=imageWriter)
  15. return spans
  16. def txt_save_images_by_bboxes(page_num: int, page, pdf_bytes_md5: str,
  17. image_bboxes: list, images_overlap_backup: list, table_bboxes: list,
  18. equation_inline_bboxes: list,
  19. equation_interline_bboxes: list, imageWriter) -> dict:
  20. """
  21. 返回一个dict, key为bbox, 值是图片地址
  22. """
  23. image_info = []
  24. image_backup_info = []
  25. table_info = []
  26. inline_eq_info = []
  27. interline_eq_info = []
  28. # 图片的保存路径组成是这样的: {s3_or_local_path}/{book_name}/{images|tables|equations}/{page_num}_{bbox[0]}_{bbox[1]}_{bbox[2]}_{bbox[3]}.jpg
  29. def return_path(type):
  30. return join_path(pdf_bytes_md5, type)
  31. for bbox in image_bboxes:
  32. image_path = cut_image(bbox, page_num, page, return_path("images"), imageWriter)
  33. image_info.append({"bbox": bbox, "image_path": image_path})
  34. for bbox in images_overlap_backup:
  35. image_path = cut_image(bbox, page_num, page, return_path("images"), imageWriter)
  36. image_backup_info.append({"bbox": bbox, "image_path": image_path})
  37. for bbox in table_bboxes:
  38. image_path = cut_image(bbox, page_num, page, return_path("tables"), imageWriter)
  39. table_info.append({"bbox": bbox, "image_path": image_path})
  40. return image_info, image_backup_info, table_info, inline_eq_info, interline_eq_info