utils.py 2.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263
  1. import fitz
  2. import numpy as np
  3. from loguru import logger
  4. def fitz_doc_to_image(doc, dpi=200) -> dict:
  5. """Convert fitz.Document to image, Then convert the image to numpy array.
  6. Args:
  7. doc (_type_): pymudoc page
  8. dpi (int, optional): reset the dpi of dpi. Defaults to 200.
  9. Returns:
  10. dict: {'img': numpy array, 'width': width, 'height': height }
  11. """
  12. mat = fitz.Matrix(dpi / 72, dpi / 72)
  13. pm = doc.get_pixmap(matrix=mat, alpha=False)
  14. # If the width or height exceeds 4500 after scaling, do not scale further.
  15. if pm.width > 4500 or pm.height > 4500:
  16. pm = doc.get_pixmap(matrix=fitz.Matrix(1, 1), alpha=False)
  17. # Convert pixmap samples directly to numpy array
  18. img = np.frombuffer(pm.samples, dtype=np.uint8).reshape(pm.height, pm.width, 3)
  19. img_dict = {'img': img, 'width': pm.width, 'height': pm.height}
  20. return img_dict
  21. def load_images_from_pdf(pdf_bytes: bytes, dpi=200, start_page_id=0, end_page_id=None) -> list:
  22. images = []
  23. with fitz.open('pdf', pdf_bytes) as doc:
  24. pdf_page_num = doc.page_count
  25. end_page_id = (
  26. end_page_id
  27. if end_page_id is not None and end_page_id >= 0
  28. else pdf_page_num - 1
  29. )
  30. if end_page_id > pdf_page_num - 1:
  31. logger.warning('end_page_id is out of range, use images length')
  32. end_page_id = pdf_page_num - 1
  33. for index in range(0, doc.page_count):
  34. if start_page_id <= index <= end_page_id:
  35. page = doc[index]
  36. mat = fitz.Matrix(dpi / 72, dpi / 72)
  37. pm = page.get_pixmap(matrix=mat, alpha=False)
  38. # If the width or height exceeds 4500 after scaling, do not scale further.
  39. if pm.width > 4500 or pm.height > 4500:
  40. pm = page.get_pixmap(matrix=fitz.Matrix(1, 1), alpha=False)
  41. # Convert pixmap samples directly to numpy array
  42. img = np.frombuffer(pm.samples, dtype=np.uint8).reshape(pm.height, pm.width, 3)
  43. img_dict = {'img': img, 'width': pm.width, 'height': pm.height}
  44. else:
  45. img_dict = {'img': [], 'width': 0, 'height': 0}
  46. images.append(img_dict)
  47. return images